package jsplugin import ( "bytes" "context" "encoding/base64" "encoding/json" "io" "math" "mime/multipart" "net/http" "net/http/httptest" "net/textproto" "strings" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/pkg/billingexpr" pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" "github.com/QuantumNous/new-api/plugins" "github.com/QuantumNous/new-api/relay/channel" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/service" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) const mockPlugin = ` export const meta = { apiVersion: 1, key: "mock-task", name: "Mock Task", version: "1.0.0", author: {name: "Test"}, channelTypes: [1001], models: ["mock-v1"], fetchMode: "per_task", protocols: ["openai_video"], usageSchema: {seconds: {type: "number", unit: "second"}, mode: {enum: ["std", "pro"]}}, }; export function buildSubmitRequest(ctx) { if (!ctx.requestBody.prompt) throw new Error("prompt required"); return { url: ctx.baseUrl + "/submit", method: "POST", headers: {"X-Plugin": "submit"}, body: {prompt: ctx.requestBody.prompt}, action: "text_to_video", model: "mock-v1", rewriteModel: "mock-upstream" }; } export function parseSubmitResponse(ctx, resp) { return { taskId: resp.body.id, taskData: {accepted: true, status: resp.statusCode}, }; } export function extractUsage(ctx) { return {seconds: 5, mode: "pro"}; } export function extractUsageOnSubmit(ctx, data) { return {seconds: data.seconds || 7}; } export function extractUsageOnComplete(task, result) { return {upstreamUnits: 23}; } export function buildQueryRequest(ctx) { return {url: ctx.baseUrl + "/tasks/" + ctx.taskId, method: "GET", headers: {"X-Plugin": "query"}}; } export function parseTaskResult(ctx, body) { return {taskId: body.id, status: "SUCCESS", progress: "100%", url: body.url}; } export function listArtifacts() { return []; } export function buildContentRequest() { throw new Error("artifact_not_found"); } export const protocols = {openai_video: { decodeRequest: function(ctx) { return {kind: "submit", model: ctx.model, requestBody: ctx.body.value}; }, render: function(ctx, task) { return {id: task.task_id, status: "completed"}; } }}; ` func TestTaskAdaptorRejectsDeprecatedClientResponse(t *testing.T) { source := strings.Replace(mockPlugin, `taskData: {accepted: true, status: resp.statusCode},`, `taskData: {}, clientResponse: {id: ctx.publicTaskId},`, 1) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "task_public"}} adaptor.Init(info) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil) response := &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"id":"upstream"}`))} parsed, taskErr := adaptor.ParseResponse(c, response, info) assert.Nil(t, parsed) require.NotNil(t, taskErr) require.Error(t, taskErr.Error) assert.Contains(t, taskErr.Error.Error(), "must not return clientResponse") } func TestTaskAdaptorBuildsMultipartFromOpaqueFileReference(t *testing.T) { source := ` export const meta = {apiVersion:1,key:"multipart",name:"Multipart",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task"}; export function buildSubmitRequest(ctx) { return {url:ctx.baseUrl+"/submit",bodyType:"multipart",parts:[{name:"model",value:"m"},{name:"input_reference",fileRef:ctx.files[0].ref}]}; } export function parseSubmitResponse(ctx,r){return {taskId:"1"}} export function buildQueryRequest(){return {url:"https://example.com"}} export function parseTaskResult(){return {status:"SUCCESS"}} ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} adaptor.Init(info) var input bytes.Buffer writer := multipart.NewWriter(&input) file, err := writer.CreateFormFile("input_reference", "ref.png") require.NoError(t, err) _, err = file.Write([]byte("image-bytes")) require.NoError(t, err) require.NoError(t, writer.Close()) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", bytes.NewReader(input.Bytes())) c.Request.Header.Set("Content-Type", writer.FormDataContentType()) c.Set("task_request", relaycommon.TaskSubmitReq{Prompt: "p"}) body, err := adaptor.BuildRequestBody(c, info) require.NoError(t, err) requestBytes, err := io.ReadAll(body) require.NoError(t, err) reader := multipart.NewReader(bytes.NewReader(requestBytes), strings.TrimPrefix(c.GetHeader("Content-Type"), "multipart/form-data; boundary=")) form, err := reader.ReadForm(1024) require.NoError(t, err) assert.Equal(t, []string{"m"}, form.Value["model"]) require.Len(t, form.File["input_reference"], 1) opened, err := form.File["input_reference"][0].Open() require.NoError(t, err) content, err := io.ReadAll(opened) require.NoError(t, err) assert.Equal(t, "image-bytes", string(content)) } func TestTaskAdaptorInlinesJSONFilePlaceholders(t *testing.T) { const fileBytes = "image-bytes" encoded := base64.StdEncoding.EncodeToString([]byte(fileBytes)) source := ` export const meta = {apiVersion:1,key:"json-inline",name:"JSON Inline",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task"}; export function buildSubmitRequest(ctx) { return {url:ctx.baseUrl+"/submit",body:{ prompt:"p", image:{__fileRef:ctx.files[0].ref,encoding:"base64"}, nested:{items:[{__fileRef:ctx.files[0].ref,encoding:"dataUrl",mimeType:"image/png"}]}, dataUrl:{__fileRef:ctx.files[0].ref,encoding:"dataUrl"} }}; } export function parseSubmitResponse(){return {taskId:"1"}} export function buildQueryRequest(){return {url:"https://example.com"}} export function parseTaskResult(){return {status:"SUCCESS"}} ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} adaptor.Init(info) c := newMultipartFileContext(t, "input_reference", "ref.png", "image/jpeg", []byte(fileBytes)) c.Set("task_request", map[string]any{"prompt": "p"}) body, err := adaptor.BuildRequestBody(c, info) require.NoError(t, err) requestBytes, err := io.ReadAll(body) require.NoError(t, err) var decoded map[string]any require.NoError(t, common.Unmarshal(requestBytes, &decoded)) assert.Equal(t, "p", decoded["prompt"]) assert.Equal(t, encoded, decoded["image"]) nested := decoded["nested"].(map[string]any) items := nested["items"].([]any) require.Len(t, items, 1) assert.Equal(t, "data:image/png;base64,"+encoded, items[0]) assert.Equal(t, "data:image/jpeg;base64,"+encoded, decoded["dataUrl"]) } func TestTaskAdaptorJSONFilePlaceholderErrors(t *testing.T) { tests := []struct { name string part string fileSize int globalMB int wantContain string }{ {name: "unknown ref", part: `{__fileRef:"request_file:missing",encoding:"base64"}`, fileSize: 4, wantContain: `unknown file reference "request_file:missing"`}, {name: "extra key", part: `{__fileRef:"request_file:input_reference",encoding:"base64",extra:true}`, fileSize: 4, wantContain: "invalid file placeholder"}, {name: "missing encoding", part: `{__fileRef:"request_file:input_reference"}`, fileSize: 4, wantContain: "encoding"}, {name: "oversize maxBytes", part: `{__fileRef:"request_file:input_reference",encoding:"base64",maxBytes:3}`, fileSize: 4, wantContain: "3 byte limit"}, {name: "oversize global", part: `{__fileRef:"request_file:input_reference",encoding:"base64"}`, fileSize: 2 << 20, globalMB: 1, wantContain: "1048576 byte limit"}, {name: "multiple references cap", part: `{a:{__fileRef:"request_file:input_reference",encoding:"base64"},b:{__fileRef:"request_file:input_reference",encoding:"base64"}}`, fileSize: 700 << 10, globalMB: 1, wantContain: "1048576 byte limit"}, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { if testCase.globalMB > 0 { previous := constant.MaxFileDownloadMB constant.MaxFileDownloadMB = testCase.globalMB t.Cleanup(func() { constant.MaxFileDownloadMB = previous }) } source := strings.Replace(` export const meta = {apiVersion:1,key:"json-inline-err",name:"JSON Inline Err",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task"}; export function buildSubmitRequest() { return {url:"https://provider.example/submit",body:PLACEHOLDER}; } export function parseSubmitResponse(){return {taskId:"1"}} export function buildQueryRequest(){return {url:"https://example.com"}} export function parseTaskResult(){return {status:"SUCCESS"}} `, "PLACEHOLDER", testCase.part, 1) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} adaptor.Init(info) c := newMultipartFileContext(t, "input_reference", "ref.bin", "application/octet-stream", bytes.Repeat([]byte("x"), testCase.fileSize)) c.Set("task_request", map[string]any{"prompt": "p"}) _, err = adaptor.BuildRequestBody(c, info) require.Error(t, err) assert.Contains(t, err.Error(), testCase.wantContain) }) } } func newMultipartFileContext(t *testing.T, field, filename, contentType string, content []byte) *gin.Context { t.Helper() var input bytes.Buffer writer := multipart.NewWriter(&input) part, err := writer.CreatePart(textproto.MIMEHeader{ "Content-Disposition": {`form-data; name="` + field + `"; filename="` + filename + `"`}, "Content-Type": {contentType}, }) require.NoError(t, err) _, err = part.Write(content) require.NoError(t, err) require.NoError(t, writer.Close()) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", bytes.NewReader(input.Bytes())) c.Request.Header.Set("Content-Type", writer.FormDataContentType()) return c } func TestTaskAdaptorDoesNotEmitInjectedMultipartDispositionHeaders(t *testing.T) { tests := []struct { name string part string }{ {name: "part name", part: `{name:"prompt\r\nX-Injected: yes",value:"hello"}`}, {name: "filename", part: `{name:"input_reference",fileRef:"request_file:input_reference",filename:"safe.png\r\nX-Injected: yes"}`}, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { source := strings.Replace(` export const meta = {apiVersion:1,key:"multipart-safe",name:"Multipart Safe",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task"}; export function buildSubmitRequest(ctx) { return {url:ctx.baseUrl+"/submit",bodyType:"multipart",parts:[PART]}; } export function parseSubmitResponse(){return {taskId:"1"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}} `, "PART", testCase.part, 1) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} adaptor.Init(info) var input bytes.Buffer writer := multipart.NewWriter(&input) file, err := writer.CreateFormFile("input_reference", "input.png") require.NoError(t, err) _, err = file.Write([]byte("image")) require.NoError(t, err) require.NoError(t, writer.Close()) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", bytes.NewReader(input.Bytes())) c.Request.Header.Set("Content-Type", writer.FormDataContentType()) c.Set("task_request", map[string]any{"model": "m"}) body, buildErr := adaptor.BuildRequestBody(c, info) if buildErr != nil { assert.Contains(t, buildErr.Error(), "multipart") return } encoded, err := io.ReadAll(body) require.NoError(t, err) assert.NotContains(t, string(encoded), "\r\nX-Injected: yes") }) } } func TestTaskAdaptorRejectsPostDistributionEndpointModelDrift(t *testing.T) { source := ` export const meta = {apiVersion:1,key:"endpoint-drift",name:"Endpoint Drift",version:"1.0.0",author:{name:"Test"},models:["claimed-model"],fetchMode:"per_task"}; export function buildSubmitRequest(ctx) { return {url:ctx.baseUrl+"/submit",method:"POST",model:"outside-model",rewriteModel:"allowed-upstream-rewrite"}; } export function parseSubmitResponse(){return {taskId:"1"}} export function buildQueryRequest(){return {url:"https://example.com"}} export function parseTaskResult(){return {status:"SUCCESS"}} ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, OriginModelName: "claimed-model", } adaptor.Init(info) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) c.Set("task_request", map[string]any{"model": "claimed-model"}) c.Set("resolved_task_model", "claimed-model") c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{ Plugin: plugin, }) taskErr := adaptor.ValidateRequestAndSetAction(c, info) require.NotNil(t, taskErr) assert.Equal(t, "claimed-model", info.OriginModelName) assert.Contains(t, taskErr.Message, "does not match") } func TestTaskAdaptorReDecodesFinalCandidateAndRejectsModelDrift(t *testing.T) { source := ` export const meta = {apiVersion:1,key:"redecode",name:"Redecode",version:"1.0.0",author:{name:"Test"},models:["claimed-model"],fetchMode:"per_task",protocols:[{name:"openai_responses",supports:["sync","background"]}]}; let calls = 0; export const protocols = {openai_responses:{decodeRequest:function(ctx){calls++;return {kind:"submit",model:calls === 1 ? ctx.model : "drifted-model",requestBody:ctx.body.value};},renderFinal:function(){return {};}}}; export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit"}} export function parseSubmitResponse(){return {taskId:"one"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}} ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) protocolContext := pluginruntime.ProtocolRequestContext{ RouteRequestContext: pluginruntime.RouteRequestContext{Body: map[string]any{"kind": "json", "value": map[string]any{"model": "claimed-model"}}, RequestBody: map[string]any{"model": "claimed-model"}}, Protocol: "openai_responses", Model: "claimed-model", } _, err = plugin.Engine.CallPath(context.Background(), "protocols", []string{"openai_responses", "decodeRequest"}, protocolContext.JSValue()) require.NoError(t, err) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{Plugin: plugin, Protocol: "openai_responses", Model: "claimed-model"}) c.Set(pluginruntime.ContextKeyProtocolRequest, protocolContext) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, OriginModelName: "claimed-model"} adaptor := New(plugin) adaptor.Init(info) taskErr := adaptor.ValidateRequestAndSetAction(c, info) require.NotNil(t, taskErr) assert.Equal(t, http.StatusBadRequest, taskErr.StatusCode) assert.Contains(t, taskErr.Message, "pinned model") } func TestTaskAdaptorRejectsRendererFromFinalProtocolDecoder(t *testing.T) { source := ` export const meta = {apiVersion:1,key:"renderer-reject",name:"Renderer Reject",version:"1.0.0",author:{name:"Test"},models:["claimed-model"],fetchMode:"per_task",protocols:[{name:"openai_responses",supports:["sync","background"]}]}; export const protocols = {openai_responses:{decodeRequest:function(ctx){return {kind:"submit",model:ctx.model,requestBody:ctx.body.value,renderer:"legacy"};},renderFinal:function(){return {};}}}; export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit"}} export function parseSubmitResponse(){return {taskId:"one"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}} ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) protocolContext := pluginruntime.ProtocolRequestContext{ RouteRequestContext: pluginruntime.RouteRequestContext{Body: map[string]any{"kind": "json", "value": map[string]any{"model": "claimed-model"}}}, Protocol: "openai_responses", Model: "claimed-model", } c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{Plugin: plugin, Protocol: "openai_responses", Model: "claimed-model"}) c.Set(pluginruntime.ContextKeyProtocolRequest, protocolContext) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, OriginModelName: "claimed-model"} adaptor := New(plugin) adaptor.Init(info) taskErr := adaptor.ValidateRequestAndSetAction(c, info) require.NotNil(t, taskErr) assert.Equal(t, http.StatusBadRequest, taskErr.StatusCode) assert.Contains(t, taskErr.Message, "must not return renderer") } func TestTaskAdaptorBuildContentRequestHookAndMissingFallback(t *testing.T) { source := strings.Replace(mockPlugin, `export function listArtifacts() { return []; } export function buildContentRequest() { throw new Error("artifact_not_found"); }`, `export function listArtifacts(task) { return [{key: "video", type: "video", mimeType: "video/mp4"}]; } export function buildContentRequest(ctx) { if (ctx.data.id !== "raw-upstream" || ctx.upstreamTaskId !== "upstream-task" || ctx.producerVersion !== "0.9.0") throw new Error("bad task context"); return {url: ctx.baseUrl + "/content/" + ctx.artifactKey, method: ctx.clientRequest.method, headers: {"X-Content": "plugin"}}; }`, 1) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) adaptor.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example", ApiKey: "key"}}) taskData, err := common.Marshal(map[string]any{"id": "raw-upstream"}) require.NoError(t, err) task := &model.Task{ TaskID: "task-public", Status: model.TaskStatusSuccess, Data: taskData, PrivateData: model.TaskPrivateData{ UpstreamTaskID: "upstream-task", Execution: &model.TaskExecutionSnapshot{TaskPlugin: &model.TaskPluginSnapshot{ Version: "0.9.0", }}, }, } artifacts, err := adaptor.ListArtifacts(task) require.NoError(t, err) require.Equal(t, []channel.TaskArtifact{{Key: "video", Type: "video", MimeType: "video/mp4"}}, artifacts) descriptor, err := adaptor.BuildContentRequest(task, "video", channel.TaskArtifactClientRequest{Method: http.MethodHead}) require.NoError(t, err) require.NotNil(t, descriptor) assert.Equal(t, "https://provider.example/content/video", descriptor.URL) assert.Equal(t, http.MethodHead, descriptor.Method) assert.Equal(t, "plugin", descriptor.Headers["X-Content"]) withoutHook, err := pluginruntime.NewRegistry().Register(` export const meta = {apiVersion:1,key:"no-artifacts",name:"No Artifacts",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task"}; export function buildSubmitRequest(){return {url:"https://provider.example"};} export function parseSubmitResponse(){return {taskId:"1"};} export function buildQueryRequest(){return {url:"https://provider.example"};} export function parseTaskResult(){return {status:"SUCCESS"};} `, pluginruntime.Options{}) require.NoError(t, err) fallback := New(withoutHook) fallback.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}}) artifacts, err = fallback.ListArtifacts(&model.Task{}) require.NoError(t, err) assert.Nil(t, artifacts) descriptor, err = fallback.BuildContentRequest(&model.Task{}, "video", channel.TaskArtifactClientRequest{Method: http.MethodGet}) require.NoError(t, err) assert.Nil(t, descriptor) } func TestTaskAdaptorRejectsInvalidArtifactProjection(t *testing.T) { testCases := []struct { name string projection string }{ {name: "duplicate key", projection: `[{key:"video",type:"video"},{key:"video",type:"video"}]`}, {name: "array index identity", projection: `[{key:"video",type:"video",index:0}]`}, {name: "upstream url", projection: `[{key:"video",type:"video",url:"https://cdn.example/video.mp4"}]`}, {name: "invalid key", projection: `[{key:"video/0",type:"video"}]`}, {name: "unsupported type", projection: `[{key:"video",type:"text"}]`}, } for _, testCase := range testCases { t.Run(testCase.name, func(t *testing.T) { source := strings.Replace(mockPlugin, `export function listArtifacts() { return []; } export function buildContentRequest() { throw new Error("artifact_not_found"); }`, `export function listArtifacts() { return `+testCase.projection+`; } export function buildContentRequest(ctx) { return {url:ctx.baseUrl+"/content",method:"GET"}; } `, 1) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) _, err = adaptor.ListArtifacts(&model.Task{TaskID: "task", Status: model.TaskStatusSuccess, Data: []byte(`{}`)}) require.Error(t, err) }) } } func TestTaskAdaptorAllowsExplicitCredentiallessCDNRequest(t *testing.T) { source := strings.Replace(mockPlugin, `export function listArtifacts() { return []; } export function buildContentRequest() { throw new Error("artifact_not_found"); }`, `export function listArtifacts() { return [{key:"video",type:"video"}]; } export function buildContentRequest(ctx) { return {url:"https://cdn.example/video.mp4",method:ctx.clientRequest.method,credentialless:true}; } `, 1) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) adaptor.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}}) descriptor, err := adaptor.BuildContentRequest( &model.Task{TaskID: "task", Data: []byte(`{}`)}, "video", channel.TaskArtifactClientRequest{Method: http.MethodGet}, ) require.NoError(t, err) require.NotNil(t, descriptor) assert.True(t, descriptor.Credentialless) assert.Equal(t, "https://cdn.example/video.mp4", descriptor.URL) } func TestTaskAdaptorMapsJSContract(t *testing.T) { gin.SetMode(gin.TestMode) service.InitHttpClient() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/submit": assert.Equal(t, "submit", r.Header.Get("X-Plugin")) body, err := io.ReadAll(r.Body) require.NoError(t, err) assert.JSONEq(t, `{"prompt":"hello"}`, string(body)) _, _ = w.Write([]byte(`{"id":"upstream-1"}`)) case "/tasks/upstream-1": assert.Equal(t, "query", r.Header.Get("X-Plugin")) _, _ = w.Write([]byte(`{"id":"upstream-1","url":"https://cdn.example/video.mp4"}`)) default: http.NotFound(w, r) } })) defer server.Close() registry := pluginruntime.NewRegistry() plugin, err := registry.Register(mockPlugin, pluginruntime.Options{Key: "mock-task", Version: "1.0.0"}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: server.URL, ApiKey: "secret"}, OriginModelName: "client-model", TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "task_public"}} adaptor.Init(info) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil) c.Set("task_request", relaycommon.TaskSubmitReq{Prompt: "hello"}) require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info)) assert.Equal(t, "text_to_video", info.Action) assert.Equal(t, "mock-v1", info.OriginModelName) assert.Equal(t, "mock-upstream", info.UpstreamModelName) assert.Equal(t, []string{"mock-v1"}, adaptor.GetModelList()) assert.Equal(t, "Mock Task", adaptor.GetChannelName()) assert.Equal(t, map[string]float64{"seconds": 5}, adaptor.EstimateBilling(c, info)) requestBody, err := adaptor.BuildRequestBody(c, info) require.NoError(t, err) url, err := adaptor.BuildRequestURL(info) require.NoError(t, err) assert.Equal(t, server.URL+"/submit", url) req := httptest.NewRequest(http.MethodPost, url, nil) require.NoError(t, adaptor.BuildRequestHeader(c, req, info)) assert.Equal(t, "submit", req.Header.Get("X-Plugin")) resp, err := adaptor.DoRequest(c, info, requestBody) require.NoError(t, err) parsed, taskErr := adaptor.ParseResponse(c, resp, info) require.Nil(t, taskErr) require.NotNil(t, parsed) assert.Equal(t, "upstream-1", parsed.UpstreamTaskID) assert.JSONEq(t, `{"accepted":true,"status":200}`, string(parsed.TaskData)) assert.Nil(t, parsed.ClientResponse) assert.Empty(t, recorder.Body.String(), "response parsing must not write before the durable task barrier") assert.Equal(t, map[string]float64{"seconds": 7}, adaptor.AdjustBillingOnSubmit(info, []byte(`{"seconds":7}`))) queryResp, err := adaptor.FetchTask(server.URL, "secret", &model.Task{ Action: info.Action, PrivateData: model.TaskPrivateData{UpstreamTaskID: parsed.UpstreamTaskID}, }, "") require.NoError(t, err) queryBody, err := io.ReadAll(queryResp.Body) require.NoError(t, err) require.NoError(t, queryResp.Body.Close()) result, err := adaptor.ParseTaskResult(&model.Task{}, queryResp, queryBody) require.NoError(t, err) assert.Equal(t, "SUCCESS", result.Status) assert.Equal(t, "https://cdn.example/video.mp4", result.Url) assert.Zero(t, adaptor.AdjustBillingOnComplete(&model.Task{}, result)) assert.Equal(t, 23, result.TotalTokens) rendered, err := adaptor.ConvertToOpenAIVideo(&model.Task{TaskID: "task_public", Status: model.TaskStatusSuccess}) require.NoError(t, err) assert.JSONEq(t, `{ "id":"task_public", "object":"video", "model":"", "status":"completed", "progress":0, "created_at":0 }`, string(rendered)) _, err = plugin.Engine.Export(context.Background(), "meta") require.NoError(t, err) } func TestTaskAdaptorPreservesSoraVideoResponseFields(t *testing.T) { source, err := plugins.Source("sora") require.NoError(t, err) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) for _, tc := range []struct { status model.TaskStatus want string }{ {model.TaskStatusInProgress, "in_progress"}, {model.TaskStatusSuccess, "completed"}, {model.TaskStatusFailure, "failed"}, } { t.Run(string(tc.status), func(t *testing.T) { task := &model.Task{ TaskID: "task_public", Status: tc.status, Progress: "42%", CreatedAt: 100, FinishTime: 200, Properties: model.Properties{OriginModelName: "origin-model"}, PrivateData: model.TaskPrivateData{UpstreamTaskID: "upstream-task-id"}, Data: []byte(`{ "id":"upstream-task-id","task_id":"upstream-task-id", "object":"provider-video","model":"provider-model", "status":"completed","progress":100,"created_at":10,"completed_at":20, "url":"https://cdn.example/video.mp4", "metadata":{"url":"https://cdn.example/video.mp4","URL":"https://cdn.example/uppercase.mp4"}, "provider_payload":{"task_id":"upstream-task-id","items":[{"enabled":false,"count":0,"value":null}]}, "seconds":8,"resolution":"720p","aspect_ratio":"16:9", "reference_images":["https://cdn.example/reference.png"], "error":{"code":"provider_error","message":"provider rejected request","detail":{"retryable":false}} }`), } rendered, err := adaptor.ConvertToOpenAIVideo(task) require.NoError(t, err) var fields map[string]any require.NoError(t, common.Unmarshal(rendered, &fields)) assert.Equal(t, "task_public", fields["id"]) assert.NotContains(t, fields, "task_id") assert.Equal(t, "video", fields["object"]) assert.Equal(t, "origin-model", fields["model"]) assert.Equal(t, tc.want, fields["status"]) assert.Equal(t, float64(42), fields["progress"]) assert.Equal(t, float64(100), fields["created_at"]) if tc.status == model.TaskStatusSuccess { assert.Equal(t, float64(200), fields["completed_at"]) } else { assert.NotContains(t, fields, "completed_at") } assert.Equal(t, "https://cdn.example/video.mp4", fields["url"]) assert.Equal(t, map[string]any{ "url": "https://cdn.example/video.mp4", "URL": "https://cdn.example/uppercase.mp4", }, fields["metadata"]) assert.Equal(t, map[string]any{ "task_id": "task_public", "items": []any{map[string]any{"enabled": false, "count": float64(0), "value": nil}}, }, fields["provider_payload"]) assert.Equal(t, float64(8), fields["seconds"]) assert.Equal(t, "720p", fields["resolution"]) assert.Equal(t, "16:9", fields["aspect_ratio"]) assert.Equal(t, []any{"https://cdn.example/reference.png"}, fields["reference_images"]) assert.Equal(t, map[string]any{ "code": "provider_error", "message": "provider rejected request", "detail": map[string]any{"retryable": false}, }, fields["error"]) }) } } func TestTaskAdaptorRejectsNonObjectOpenAIVideoRendererOutput(t *testing.T) { for _, value := range []string{"null", "[]", `"video"`, "42", "false"} { t.Run(value, func(t *testing.T) { source := strings.Replace(mockPlugin, `return {id: task.task_id, status: "completed"};`, "return "+value+";", 1) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) _, err = New(plugin).ConvertToOpenAIVideo(&model.Task{TaskID: "task_public"}) require.ErrorContains(t, err, "invalid OpenAI video object") }) } } func TestTaskAdaptorPreservesOpenAIVideoFailureSlotsAndOwnsLifecycle(t *testing.T) { source := strings.Replace(mockPlugin, `render: function(ctx, task) { return {id: task.task_id, status: "completed"}; }`, `render: function() { return {id:"provider", object:"provider", model:"provider-model", status:"completed", progress:100, created_at:99, completed_at:20, error:{code:"provider_error",message:"provider rejected request"}}; }`, 1) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) task := &model.Task{ TaskID: "task_public", Status: model.TaskStatusFailure, FailReason: "provider secret", CreatedAt: 10, UpdatedAt: 20, Properties: model.Properties{OriginModelName: "origin-model"}, } rendered, err := adaptor.ConvertToOpenAIVideo(task) require.NoError(t, err) assert.JSONEq(t, `{"id":"task_public","object":"video","model":"origin-model","status":"failed","progress":0,"created_at":10,"error":{"message":"provider rejected request","code":"provider_error"}}`, string(rendered)) } func TestTaskAdaptorBoundsNativeUsageBeforeQuotaCalculation(t *testing.T) { source := ` export const meta = { apiVersion: 1, key: "bounded-usage", name: "Bounded Usage", version: "1.0.0", author: {name: "Test"}, models: ["model"], fetchMode: "per_task", usageSchema: { duration: {type: "number", unit: "second"}, count: {type: "number", unit: "count"}, tokens: {type: "number", unit: "token"}, mode: {enum: ["std", "pro"]}, }, usageExamples: [{label: "std · 1s", facts: {duration: 1, count: 1, tokens: 1, mode: "std"}}], }; export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl + "/submit", method: "POST", body: {}}; } export function parseSubmitResponse() { return {taskId: "1"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } export function extractUsage(ctx) { const entries = (ctx.requestBody || {}).hookUsageEntries || []; const facts = {}; entries.forEach(function(entry) { facts[entry.name] = entry.value; }); return facts; } export function extractUsageOnSubmit(ctx, data) { return (data || {}).usage || {}; } export function extractUsageOnComplete(task, result, body) { return (body || {}).completionUsage || {}; } ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) newRequest := func(t *testing.T, requestBody map[string]any) (*TaskAdaptor, *gin.Context, *relaycommon.RelayInfo) { t.Helper() adaptor := New(plugin) info := &relaycommon.RelayInfo{ ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, } adaptor.Init(info) context, _ := gin.CreateTestContext(httptest.NewRecorder()) context.Request = httptest.NewRequest(http.MethodPost, "/native/submit", nil) context.Set("task_request", requestBody) return adaptor, context, info } requestTests := []struct { name string body map[string]any }{ { name: "duration in resolved metadata", body: map[string]any{"metadata": map[string]any{"duration": relaycommon.MaxTaskDurationSeconds + 1}}, }, { name: "count in resolved metadata", body: map[string]any{"metadata": map[string]any{"count": dto.MaxImageN + 1}}, }, { name: "declared enum in resolved metadata", body: map[string]any{"metadata": map[string]any{"mode": "turbo"}}, }, { name: "implicit duration key without declaration", body: map[string]any{"durationSeconds": relaycommon.MaxTaskDurationSeconds + 1}, }, { name: "implicit count key without declaration", body: map[string]any{"image_count": dto.MaxImageN + 1}, }, { name: "negative resolved duration", body: map[string]any{"duration": -1}, }, { name: "non-finite resolved duration", body: map[string]any{"duration": math.Inf(1)}, }, { name: "metadata cannot hide behind valid top-level duration", body: map[string]any{ "duration": relaycommon.MaxTaskDurationSeconds, "metadata": map[string]any{"duration": relaycommon.MaxTaskDurationSeconds + 1}, }, }, { name: "nested passthrough duration", body: map[string]any{ "metadata": map[string]any{ "parameters": map[string]any{"duration": relaycommon.MaxTaskDurationSeconds + 1}, }, }, }, } for _, testCase := range requestTests { t.Run(testCase.name, func(t *testing.T) { adaptor, context, info := newRequest(t, testCase.body) taskErr := adaptor.ValidateRequestAndSetAction(context, info) require.NotNil(t, taskErr) assert.Equal(t, "plugin_usage_invalid", taskErr.Code) }) } hookTests := []struct { name string usage map[string]any }{ { name: "duration returned only by extractUsage", usage: map[string]any{"duration": float64(relaycommon.MaxTaskDurationSeconds + 1)}, }, { name: "count returned only by extractUsage", usage: map[string]any{"count": float64(dto.MaxImageN + 1)}, }, { name: "enum returned only by extractUsage", usage: map[string]any{"mode": "turbo"}, }, { name: "undeclared numeric ratio uses conservative host ceiling", usage: map[string]any{"custom_ratio": float64(relaycommon.MaxTaskDurationSeconds + 1)}, }, { name: "negative hook ratio", usage: map[string]any{"custom_ratio": -1.0}, }, { name: "non-finite hook ratio", usage: map[string]any{"custom_ratio": math.NaN()}, }, } for _, testCase := range hookTests { t.Run(testCase.name, func(t *testing.T) { entries := make([]any, 0, len(testCase.usage)) for key, value := range testCase.usage { entries = append(entries, map[string]any{"name": key, "value": value}) } adaptor, context, info := newRequest(t, map[string]any{"hookUsageEntries": entries}) require.Nil(t, adaptor.ValidateRequestAndSetAction(context, info)) ratios, err := adaptor.EstimateBillingValidated(context, info) require.Error(t, err) assert.Nil(t, ratios) }) } t.Run("numeric strings remain valid in vendor request fields", func(t *testing.T) { adaptor, context, info := newRequest(t, map[string]any{ "metadata": map[string]any{ "duration": "5", "count": "2", "mode": "std", }, }) assert.Nil(t, adaptor.ValidateRequestAndSetAction(context, info)) }) t.Run("numeric strings from usage hooks are rejected", func(t *testing.T) { adaptor, context, info := newRequest(t, map[string]any{ "hookUsageEntries": []any{map[string]any{"name": "duration", "value": "5"}}, }) require.Nil(t, adaptor.ValidateRequestAndSetAction(context, info)) ratios, err := adaptor.EstimateBillingValidated(context, info) require.Error(t, err) assert.Nil(t, ratios) }) t.Run("declared token facts use int32 saturation instead of duration cap", func(t *testing.T) { adaptor, context, info := newRequest(t, map[string]any{ "hookUsageEntries": []any{ map[string]any{"name": "tokens", "value": float64(500000)}, }, }) require.Nil(t, adaptor.ValidateRequestAndSetAction(context, info)) facts, err := adaptor.ExtractUsageFactsValidated(context, info) require.NoError(t, err) assert.EqualValues(t, 500000, facts["tokens"]) ratios, err := adaptor.EstimateBillingValidated(context, info) require.NoError(t, err) assert.Equal(t, 500000.0, ratios["tokens"]) }) t.Run("declared token facts saturate at the int32 quota bound", func(t *testing.T) { adaptor, context, info := newRequest(t, map[string]any{ "hookUsageEntries": []any{ map[string]any{"name": "tokens", "value": float64(common.MaxQuota) + 1}, }, }) require.Nil(t, adaptor.ValidateRequestAndSetAction(context, info)) facts, err := adaptor.ExtractUsageFactsValidated(context, info) require.NoError(t, err) assert.EqualValues(t, common.MaxQuota, facts["tokens"]) }) t.Run("canonical maxima and enum are accepted", func(t *testing.T) { adaptor, context, info := newRequest(t, map[string]any{ "duration": relaycommon.MaxTaskDurationSeconds, "count": dto.MaxImageN, "mode": "std", "hookUsageEntries": []any{ map[string]any{"name": "duration", "value": float64(relaycommon.MaxTaskDurationSeconds)}, map[string]any{"name": "count", "value": float64(dto.MaxImageN)}, map[string]any{"name": "mode", "value": "pro"}, }, }) require.Nil(t, adaptor.ValidateRequestAndSetAction(context, info)) ratios, err := adaptor.EstimateBillingValidated(context, info) require.NoError(t, err) assert.Equal(t, map[string]float64{ "duration": relaycommon.MaxTaskDurationSeconds, "count": dto.MaxImageN, }, ratios) }) t.Run("runtime error does not expose plugin-controlled usage key", func(t *testing.T) { adaptor, context, info := newRequest(t, map[string]any{ "hookUsageEntries": []any{ map[string]any{ "name": "https://private.invalid/?token=secret", "value": float64(relaycommon.MaxTaskDurationSeconds + 1), }, }, }) require.Nil(t, adaptor.ValidateRequestAndSetAction(context, info)) _, err := adaptor.EstimateBillingValidated(context, info) require.Error(t, err) assert.NotContains(t, err.Error(), "private.invalid") assert.NotContains(t, err.Error(), "secret") }) for _, testCase := range []struct { name string usage map[string]any }{ { name: "oversized completion duration is discarded", usage: map[string]any{"duration": relaycommon.MaxTaskDurationSeconds + 1}, }, { name: "oversized completion count is discarded", usage: map[string]any{"count": dto.MaxImageN + 1}, }, { name: "completion hook numeric string is discarded", usage: map[string]any{"duration": "5"}, }, } { t.Run(testCase.name, func(t *testing.T) { adaptor, _, _ := newRequest(t, map[string]any{}) body, marshalErr := common.Marshal(map[string]any{"completionUsage": testCase.usage}) require.NoError(t, marshalErr) result, parseErr := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, body) require.NoError(t, parseErr) assert.Nil(t, result.UsageFacts) assert.Zero(t, result.TotalTokens) }) } t.Run("declared completion token unit is saturated instead of discarded", func(t *testing.T) { adaptor, _, _ := newRequest(t, map[string]any{}) body, err := common.Marshal(map[string]any{"completionUsage": map[string]any{"tokens": 500000}}) require.NoError(t, err) result, err := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, body) require.NoError(t, err) assert.EqualValues(t, 500000, result.UsageFacts["tokens"]) }) t.Run("declared credit facts keep sub-integer precision", func(t *testing.T) { source := ` export const meta = { apiVersion: 1, key: "credit-decimals", name: "Credit Decimals", version: "1.0.0", author: {name: "Test"}, models: ["model"], fetchMode: "per_task", usageSchema: {units: {type: "number", unit: "credit"}}, usageExamples: [{label: "3.5 credits", facts: {units: 3.5}}], }; export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl + "/submit"}; } export function parseSubmitResponse() { return {taskId: "task"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } export function extractUsage() { return {units: 3.5}; } export function extractUsageOnComplete() { return {units: 3.5}; } ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, } adaptor.Init(info) context, _ := gin.CreateTestContext(httptest.NewRecorder()) context.Request = httptest.NewRequest(http.MethodPost, "/native/submit", nil) context.Set("task_request", map[string]any{"model": "model"}) require.Nil(t, adaptor.ValidateRequestAndSetAction(context, info)) facts, err := adaptor.ExtractUsageFactsValidated(context, info) require.NoError(t, err) assert.Equal(t, 3.5, facts["units"]) body, err := common.Marshal(map[string]any{}) require.NoError(t, err) result, err := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, body) require.NoError(t, err) assert.Equal(t, 3.5, result.UsageFacts["units"]) }) t.Run("completion token facts are saturated instead of duration-capped", func(t *testing.T) { adaptor, _, _ := newRequest(t, map[string]any{}) body, err := common.Marshal(map[string]any{"completionUsage": map[string]any{"upstreamUnits": 5000}}) require.NoError(t, err) result, err := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, body) require.NoError(t, err) assert.Equal(t, 5000, result.TotalTokens) assert.EqualValues(t, 5000, result.UsageFacts["upstreamUnits"]) }) t.Run("invalid post-submit adjustment is discarded before recalculation", func(t *testing.T) { adaptor, _, info := newRequest(t, map[string]any{}) ratios := adaptor.AdjustBillingOnSubmit(info, []byte(`{"usage":{"duration":1000000000000000}}`)) assert.Nil(t, ratios) }) } func TestTaskAdaptorSeparatesExpressionFactsFromLegacyBillingRatios(t *testing.T) { source := ` export const meta = { apiVersion: 1, key: "usage-purpose", name: "Usage Purpose", version: "1.0.0", author: {name: "Test"}, models: ["usage-model"], fetchMode: "per_task", usageSchema: {seconds: {type: "number", unit: "second"}}, }; export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl + "/submit"}; } export function parseSubmitResponse() { return {taskId: "task"}; } export function buildQueryRequest(ctx) { return {url: ctx.baseUrl + "/query"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } export function extractUsage(ctx) { return ctx.usagePurpose === "billing_ratios" ? {legacy_multiplier: 2} : {seconds: 5}; } ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, } adaptor.Init(info) context, _ := gin.CreateTestContext(httptest.NewRecorder()) context.Request = httptest.NewRequest(http.MethodPost, "/submit", nil) context.Set("task_request", map[string]any{"model": "usage-model"}) require.Nil(t, adaptor.ValidateRequestAndSetAction(context, info)) facts, err := adaptor.ExtractUsageFactsValidated(context, info) require.NoError(t, err) assert.EqualValues(t, 5, facts["seconds"]) assert.Len(t, facts, 1) ratios, err := adaptor.EstimateBillingValidated(context, info) require.NoError(t, err) assert.Equal(t, map[string]float64{"legacy_multiplier": 2}, ratios) } func TestTaskAdaptorAcceptsNormalizedLegacyTokenCounters(t *testing.T) { source := ` export const meta = {apiVersion:1,key:"normalized-tokens",name:"Normalized Tokens",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task"}; export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl + "/submit"}; } export function parseSubmitResponse() { return {taskId: "task"}; } export function buildQueryRequest(ctx) { return {url: ctx.baseUrl + "/query"}; } export function parseTaskResult(ctx, body) { return {status: "SUCCESS", completionTokens: body.completion, totalTokens: body.total}; } ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) result, err := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, []byte(`{"completion":13,"total":17}`)) require.NoError(t, err) assert.Equal(t, 13, result.CompletionTokens) assert.Equal(t, 17, result.TotalTokens) assert.Nil(t, result.UsageFacts) } func TestSubmitContextExposesOriginTasks(t *testing.T) { plugin, err := pluginruntime.NewRegistry().Register(mockPlugin, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example", ApiKey: "secret"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{ OriginTasks: []relaycommon.OriginTaskRef{{ TaskID: "task_pub_1", UpstreamTaskID: "cgt-upstream-1", Action: "text_to_video", Status: "SUCCESS", Data: []byte(`{"id":"cgt-upstream-1"}`), }}, }, } adaptor.Init(info) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil) ctx := adaptor.submitContext(c, info) originTasks, ok := ctx["originTasks"].([]map[string]any) require.True(t, ok) require.Len(t, originTasks, 1) assert.Equal(t, "task_pub_1", originTasks[0]["taskId"]) assert.Equal(t, "cgt-upstream-1", originTasks[0]["upstreamTaskId"]) assert.Equal(t, "text_to_video", originTasks[0]["action"]) assert.Equal(t, "SUCCESS", originTasks[0]["status"]) assert.Equal(t, map[string]any{"id": "cgt-upstream-1"}, originTasks[0]["data"]) } func TestSubmitContextOmitsOriginTasksWhenEmpty(t *testing.T) { plugin, err := pluginruntime.NewRegistry().Register(mockPlugin, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example", ApiKey: "secret"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, } adaptor.Init(info) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil) ctx := adaptor.submitContext(c, info) _, ok := ctx["originTasks"] assert.False(t, ok) } func TestSubmitContextOriginTasksNilDataOnInvalidJSON(t *testing.T) { plugin, err := pluginruntime.NewRegistry().Register(mockPlugin, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example", ApiKey: "secret"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{ OriginTasks: []relaycommon.OriginTaskRef{{ TaskID: "task_pub_1", UpstreamTaskID: "cgt-upstream-1", Action: "text_to_video", Status: "SUCCESS", Data: []byte("not-json"), }}, }, } adaptor.Init(info) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil) ctx := adaptor.submitContext(c, info) originTasks, ok := ctx["originTasks"].([]map[string]any) require.True(t, ok) require.Len(t, originTasks, 1) assert.Nil(t, originTasks[0]["data"]) } func TestTaskAdaptorRejectsRequestHostOverride(t *testing.T) { source := strings.Replace(mockPlugin, `ctx.baseUrl + "/submit"`, `"https://attacker.example/steal"`, 1) plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{Key: "mock-task", Version: "1.0.0"}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil) c.Set("task_request", relaycommon.TaskSubmitReq{Prompt: "hello"}) taskErr := adaptor.ValidateRequestAndSetAction(c, info) require.NotNil(t, taskErr) assert.Contains(t, taskErr.Message, "not allowed") } const batchMockPlugin = ` export const meta = { apiVersion: 1, key: "mock-batch", name: "Mock Batch", version: "1.0.0", author: {name: "Test"}, channelTypes: [1002], models: ["batch-v1"], fetchMode: "batch" }; export function buildSubmitRequest(ctx) { return { url: ctx.baseUrl + "/submit", method: "POST", body: {} }; } export function parseSubmitResponse(ctx, resp) { return { taskId: resp.body.id }; } export function buildQueryRequest(ctx) { return { url: ctx.baseUrl + "/tasks/" + ctx.taskId }; } export function parseTaskResult(ctx, body) { return { taskId: body.id, status: "SUCCESS" }; } export function buildBatchQueryRequest(ctx, tasks) { return { url: ctx.baseUrl + "/batch", method: "POST", headers: { "X-Plugin": "batch" }, body: { ids: (tasks || []).map(function (task) { return task.taskId; }) } }; } export function parseBatchResult(ctx, body) { return body.items.map(function (item) { return { taskId: item.id, action: item.action, status: item.status, progress: item.progress, url: (item.urls || [])[0] || "", finishTime: item.finish || 0, data: item }; }); } export function extractUsageOnComplete(task, result, body) { return {upstreamUnits: body.usage || 0}; } ` // Covers the bridge half of the batch contract: FetchBatchTasks must build the // upstream request from the plugin descriptor, and ParseBatchResult must key // results by taskId, preserve the explicit result URL, and skip entries without // a task id. func TestTaskAdaptorBatchBridge(t *testing.T) { service.InitHttpClient() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(t, "/batch", r.URL.Path) require.Equal(t, http.MethodPost, r.Method) assert.Equal(t, "batch", r.Header.Get("X-Plugin")) body, err := io.ReadAll(r.Body) require.NoError(t, err) assert.JSONEq(t, `{"ids":["task-a","task-b"]}`, string(body)) _, _ = w.Write([]byte(`{"items":[ {"id":"task-a","action":"music","status":"SUCCESS","progress":"100%","urls":["https://cdn.example/a1.mp3","https://cdn.example/a2.mp3"],"finish":1700000000,"usage":23}, {"id":"task-b","status":"IN_PROGRESS","progress":"40%"}, {"id":"","status":"SUCCESS"} ]}`)) })) defer server.Close() plugin, err := pluginruntime.NewRegistry().Register(batchMockPlugin, pluginruntime.Options{Key: "mock-batch", Version: "1.0.0"}) require.NoError(t, err) adaptor := New(plugin) require.Equal(t, "batch", adaptor.FetchMode()) tasks := []*model.Task{ {PrivateData: model.TaskPrivateData{UpstreamTaskID: "task-a"}}, {PrivateData: model.TaskPrivateData{UpstreamTaskID: "task-b"}}, } resp, err := adaptor.FetchBatchTasks(server.URL, "secret", tasks, "") require.NoError(t, err) defer resp.Body.Close() payload, err := io.ReadAll(resp.Body) require.NoError(t, err) results, err := adaptor.ParseBatchResult(tasks, resp, payload) require.NoError(t, err) require.Len(t, results, 2, "entry without taskId must be skipped") done := results["task-a"] require.NotNil(t, done) assert.Equal(t, "music", done.Action) assert.Equal(t, "SUCCESS", done.TaskInfo.Status) assert.Equal(t, "100%", done.TaskInfo.Progress) assert.Equal(t, "https://cdn.example/a1.mp3", done.TaskInfo.Url) assert.Equal(t, int64(1700000000), done.FinishTime) assert.EqualValues(t, 23, done.TaskInfo.UsageFacts["upstreamUnits"]) assert.Equal(t, 23, done.TaskInfo.TotalTokens) require.NotNil(t, done.Data) pending := results["task-b"] require.NotNil(t, pending) assert.Equal(t, "IN_PROGRESS", pending.TaskInfo.Status) assert.Equal(t, "40%", pending.TaskInfo.Progress) assert.Empty(t, pending.TaskInfo.Url) } const mappingOrderAdaptorPlugin = ` export const meta = {apiVersion:1,key:"map-order-adaptor",name:"Map Order Adaptor",version:"1.0.0",author:{name:"Test"},models:["declared-model"],fetchMode:"per_task"}; export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl+"/submit", method:"POST", body:{upstreamModel: ctx.upstreamModel, model: ctx.model}}; } export function parseSubmitResponse(){return {taskId:"1"};} export function buildQueryRequest(){return {url:"https://provider.example"};} export function parseTaskResult(){return {status:"SUCCESS"};} ` func mappingOrderSubmitBody(t *testing.T, origin, mapping string) []byte { t.Helper() plugin, err := pluginruntime.NewRegistry().Register(mappingOrderAdaptorPlugin, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, OriginModelName: origin, } adaptor.Init(info) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil) if mapping != "" { c.Set("model_mapping", mapping) } c.Set("task_request", map[string]any{"prompt": "p"}) info.UpstreamModelName = info.OriginModelName require.NoError(t, helper.ModelMappedHelper(c, info, nil)) require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info)) body, err := adaptor.BuildRequestBody(c, info) require.NoError(t, err) raw, err := io.ReadAll(body) require.NoError(t, err) return raw } func TestTaskAdaptorBuildSubmitReceivesMappedUpstreamModel(t *testing.T) { gin.SetMode(gin.TestMode) mapped := mappingOrderSubmitBody(t, "alias-model", `{"alias-model":"mid-model","mid-model":"declared-model"}`) var decoded map[string]any require.NoError(t, common.Unmarshal(mapped, &decoded)) assert.Equal(t, "declared-model", decoded["upstreamModel"]) assert.Equal(t, "alias-model", decoded["model"]) withoutMapping := mappingOrderSubmitBody(t, "declared-model", "") emptyMapping := mappingOrderSubmitBody(t, "declared-model", "{}") assert.Equal(t, withoutMapping, emptyMapping) require.NoError(t, common.Unmarshal(withoutMapping, &decoded)) assert.Equal(t, "declared-model", decoded["upstreamModel"]) assert.Equal(t, "declared-model", decoded["model"]) } // Polling has no relay info, so query hooks can only branch on the model when // the host forwards the persisted task identities from the fetch body. func TestTaskAdaptorFetchTaskExposesModelIdentities(t *testing.T) { service.InitHttpClient() var requested string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requested = r.URL.RequestURI() _, _ = w.Write([]byte(`{}`)) })) defer server.Close() source := ` export const meta = {apiVersion:1,key:"query-model",name:"Query Model",version:"1.0.0",author:{name:"Test"},models:["alias"],fetchMode:"per_task"}; export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit"}} export function parseSubmitResponse(){return {taskId:"1"}} export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/tasks/"+ctx.model+"/"+ctx.upstreamModel+"/"+ctx.taskId,method:"GET"}} export function parseTaskResult(){return {status:"SUCCESS"}} ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) testCases := []struct { name string task *model.Task want string }{ { name: "mapped model", task: &model.Task{ Properties: model.Properties{OriginModelName: "alias", UpstreamModelName: "declared-model"}, PrivateData: model.TaskPrivateData{UpstreamTaskID: "t1"}, }, want: "/tasks/alias/declared-model/t1", }, { name: "unmapped model falls back to the origin name", task: &model.Task{ Properties: model.Properties{OriginModelName: "alias"}, PrivateData: model.TaskPrivateData{UpstreamTaskID: "t1"}, }, want: "/tasks/alias/alias/t1", }, } for _, testCase := range testCases { t.Run(testCase.name, func(t *testing.T) { resp, fetchErr := adaptor.FetchTask(server.URL, "secret", testCase.task, "") require.NoError(t, fetchErr) require.NoError(t, resp.Body.Close()) assert.Equal(t, testCase.want, requested) }) } } func TestTaskAdaptorQueryContextOmitsRequestBody(t *testing.T) { service.InitHttpClient() var captured map[string]any server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.NoError(t, common.DecodeJson(r.Body, &captured)) _, _ = w.Write([]byte(`{"ok":true}`)) })) defer server.Close() source := ` export const meta = {apiVersion:1,key:"query-ctx",name:"Query Ctx",version:"1.0.0",author:{name:"Test"},models:["alias"],fetchMode:"per_task"}; export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit"}} export function parseSubmitResponse(){return {taskId:"1",state:{req_key:"from-submit"}}} export function buildQueryRequest(ctx){ return {url:ctx.baseUrl+"/query",method:"POST",body:{ keys: Object.keys(ctx).sort(), taskId: ctx.taskId, publicTaskId: ctx.publicTaskId, action: ctx.action, model: ctx.model, upstreamModel: ctx.upstreamModel, data: ctx.data, state: ctx.state, hasRequestBody: Object.prototype.hasOwnProperty.call(ctx, "requestBody") }}; } export function parseTaskResult(ctx, body, response){ return {status:"IN_PROGRESS",reason:String(response && response.status),url:ctx.taskId,state:{round:"poll"}}; } ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) adaptor.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: server.URL, ApiKey: "secret"}}) task := &model.Task{ TaskID: "task_public", Action: constant.TaskActionImageToVideo, Properties: model.Properties{ OriginModelName: "alias", UpstreamModelName: "declared", }, Data: []byte(`{"snapshot":true}`), PrivateData: model.TaskPrivateData{ UpstreamTaskID: "upstream-1", PluginState: []byte(`{"req_key":"kept"}`), }, } resp, err := adaptor.FetchTask(server.URL, "secret", task, "") require.NoError(t, err) require.NoError(t, resp.Body.Close()) assert.Equal(t, "upstream-1", captured["taskId"]) assert.Equal(t, "task_public", captured["publicTaskId"]) assert.Equal(t, constant.NormalizeTaskAction(constant.TaskActionImageToVideo), captured["action"]) assert.Equal(t, "alias", captured["model"]) assert.Equal(t, "declared", captured["upstreamModel"]) assert.Equal(t, map[string]any{"snapshot": true}, captured["data"]) assert.Equal(t, map[string]any{"req_key": "kept"}, captured["state"]) assert.Equal(t, false, captured["hasRequestBody"]) keys, ok := captured["keys"].([]any) require.True(t, ok) assert.NotContains(t, keys, "requestBody") result, err := adaptor.ParseTaskResult(task, &http.Response{StatusCode: http.StatusTeapot, Header: make(http.Header)}, []byte(`{"ok":true}`)) require.NoError(t, err) assert.Equal(t, "IN_PROGRESS", result.Status) assert.Equal(t, "418", result.Reason) assert.Equal(t, "upstream-1", result.Url) assert.JSONEq(t, `{"round":"poll"}`, string(result.PluginState)) } func TestTaskAdaptorParseSubmitResponsePersistsState(t *testing.T) { service.InitHttpClient() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte(`{"id":"upstream-1"}`)) })) defer server.Close() source := ` export const meta = {apiVersion:1,key:"submit-state",name:"Submit State",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task"}; export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit",method:"POST",body:{}}} export function parseSubmitResponse(){return {taskId:"upstream-1",state:{req_key:"from-submit"}}} export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"}} export function parseTaskResult(){return {status:"SUCCESS"}} ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: server.URL, ApiKey: "secret"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} adaptor.Init(info) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil) c.Set("task_request", relaycommon.TaskSubmitReq{Prompt: "hello"}) require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info)) body, err := adaptor.BuildRequestBody(c, info) require.NoError(t, err) resp, err := adaptor.DoRequest(c, info, body) require.NoError(t, err) parsed, taskErr := adaptor.ParseResponse(c, resp, info) require.Nil(t, taskErr) require.NotNil(t, parsed) assert.JSONEq(t, `{"req_key":"from-submit"}`, string(parsed.PluginState)) } func TestTaskAdaptorBatchQueryReceivesTaskObjects(t *testing.T) { service.InitHttpClient() var captured map[string]any server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.NoError(t, common.DecodeJson(r.Body, &captured)) _, _ = w.Write([]byte(`{"items":[]}`)) })) defer server.Close() source := ` export const meta = {apiVersion:1,key:"batch-ctx",name:"Batch Ctx",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"batch"}; export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit"}} export function parseSubmitResponse(){return {taskId:"1"}} export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/q"}} export function parseTaskResult(){return {status:"SUCCESS"}} export function buildBatchQueryRequest(ctx, tasks){ return {url:ctx.baseUrl+"/batch",method:"POST",body:{ ids: (tasks||[]).map(function(task){return task.taskId;}), models: (tasks||[]).map(function(task){return task.model;}), hasRequestBody: (tasks||[]).some(function(task){return Object.prototype.hasOwnProperty.call(task,"requestBody");}) }}; } export function parseBatchResult(){return [];} ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) adaptor := New(plugin) tasks := []*model.Task{ {Properties: model.Properties{OriginModelName: "model-a"}, PrivateData: model.TaskPrivateData{UpstreamTaskID: "task-a"}}, {Properties: model.Properties{OriginModelName: "model-b"}, PrivateData: model.TaskPrivateData{UpstreamTaskID: "task-b"}}, } resp, err := adaptor.FetchBatchTasks(server.URL, "secret", tasks, "") require.NoError(t, err) require.NoError(t, resp.Body.Close()) assert.Equal(t, []any{"task-a", "task-b"}, captured["ids"]) assert.Equal(t, []any{"model-a", "model-b"}, captured["models"]) assert.Equal(t, false, captured["hasRequestBody"]) } func TestTaskAdaptorUsageProfilesFollowExecutionModel(t *testing.T) { const source = ` export const meta = { apiVersion:1, key:"profile-usage", name:"Profile Usage", version:"1.0.0", author:{name:"Test"}, models:["image", "video"], fetchMode:"batch", usageSchema:{units:{type:"number",unit:"second"}}, usageProfiles:[ {models:["image"],schema:{units:{type:"number",unit:"count"},mode:{enum:["image"]}}}, {models:["video"],schema:{units:{type:"number",unit:"token"},mode:{enum:["video"]}}, examples:[{label:"video",facts:{units:1,mode:"video"}}]} ] }; export function buildSubmitRequest(ctx) { return {url:ctx.baseUrl+"/submit",rewriteModel:ctx.requestBody.rewriteTo || ""}; } export function parseSubmitResponse(){return {taskId:"task"};} export function buildQueryRequest(){return {url:"https://provider.example/task"};} export function buildBatchQueryRequest(){return {url:"https://provider.example/tasks"};} export function parseTaskResult(){return {status:"SUCCESS"};} export function parseBatchResult(ctx,body){return body.items;} export function extractUsage(ctx){return {units:ctx.requestBody.hookUnits,mode:ctx.requestBody.hookMode,legacyRatio:2};} export function extractUsageOnSubmit(ctx,body){return body.usage;} export function extractUsageOnComplete(ctx,result,body){return body.usage;} ` plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{}) require.NoError(t, err) newRequest := func(t *testing.T, upstream string, body map[string]any) (*TaskAdaptor, *gin.Context, *relaycommon.RelayInfo) { t.Helper() info := &relaycommon.RelayInfo{ OriginModelName: "public-alias", ChannelMeta: &relaycommon.ChannelMeta{ UpstreamModelName: upstream, ChannelBaseUrl: "https://provider.example", }, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, } adaptor := New(plugin) adaptor.Init(info) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/native/submit", nil) c.Set("task_request", body) return adaptor, c, info } for _, tc := range []struct { name, upstream, rewrite, mode string units float64 wantError bool }{ {"mapped image", "image", "", "image", 2, false}, {"mapped video token units", "video", "", "video", 500000, false}, {"rewritten model", "image", "video", "video", 500000, false}, {"image count ceiling", "image", "", "image", float64(dto.MaxImageN + 1), true}, {"profile enum", "image", "", "video", 2, true}, } { t.Run(tc.name, func(t *testing.T) { adaptor, c, info := newRequest(t, tc.upstream, map[string]any{ "hookUnits": tc.units, "hookMode": tc.mode, "rewriteTo": tc.rewrite, }) require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info)) facts, err := adaptor.ExtractUsageFactsValidated(c, info) if tc.wantError { require.Error(t, err) return } require.NoError(t, err) assert.Equal(t, map[string]any{"units": tc.units, "mode": tc.mode, "legacyRatio": 2.0}, facts) ratios, err := adaptor.EstimateBillingValidated(c, info) require.NoError(t, err) assert.Equal(t, map[string]float64{"units": tc.units, "legacyRatio": 2}, ratios) body, err := common.Marshal(map[string]any{"usage": facts}) require.NoError(t, err) assert.Equal(t, ratios, adaptor.AdjustBillingOnSubmit(info, body)) }) } t.Run("request revalidated after model rewrite", func(t *testing.T) { adaptor, c, info := newRequest(t, "video", map[string]any{ "rewriteTo": "image", "metadata": map[string]any{"units": dto.MaxImageN + 1}, }) taskErr := adaptor.ValidateRequestAndSetAction(c, info) require.NotNil(t, taskErr) assert.Equal(t, "plugin_usage_invalid", taskErr.Code) }) t.Run("request enum belongs to the rewritten model", func(t *testing.T) { adaptor, c, info := newRequest(t, "video", map[string]any{ "rewriteTo": "image", "mode": "image", "units": 2, }) require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info)) assert.Equal(t, "image", info.UpstreamModelName) }) t.Run("completion selects each persisted model including legacy origin fallback", func(t *testing.T) { adaptor, _, _ := newRequest(t, "image", map[string]any{}) for _, tc := range []struct { upstream, origin, mode string units float64 accepted bool }{ {"video", "public-alias", "video", 500000, true}, {"image", "public-alias", "image", 2, true}, {"", "video", "video", 500000, true}, {"image", "public-alias", "image", float64(dto.MaxImageN + 1), false}, {"video", "public-alias", "image", 1, false}, } { task := &model.Task{Properties: model.Properties{OriginModelName: tc.origin, UpstreamModelName: tc.upstream}} body, err := common.Marshal(map[string]any{"usage": map[string]any{"units": tc.units, "mode": tc.mode}}) require.NoError(t, err) result, err := adaptor.ParseTaskResult(task, &http.Response{StatusCode: http.StatusOK}, body) require.NoError(t, err) if tc.accepted { assert.Equal(t, map[string]any{"units": tc.units, "mode": tc.mode}, result.UsageFacts) } else { assert.Nil(t, result.UsageFacts) } } }) t.Run("legacy completion facts remain usable by stored expressions", func(t *testing.T) { updated, err := pluginruntime.NewRegistry().Register(strings.Replace(source, `export function extractUsageOnComplete(ctx,result,body){return body.usage;}`, `export function extractUsageOnComplete(){return {legacyRatio:3};}`, 1), pluginruntime.Options{}) require.NoError(t, err) _, _, info := newRequest(t, "image", map[string]any{}) adaptor := New(updated) adaptor.Init(info) result, err := adaptor.ParseTaskResult(&model.Task{Properties: model.Properties{UpstreamModelName: "image"}}, &http.Response{StatusCode: http.StatusOK}, []byte(`{}`)) require.NoError(t, err) cost, _, err := billingexpr.RunExprWithRequest(`tier("legacy", u("legacyRatio") * 2)`, billingexpr.TokenParams{}, billingexpr.RequestInput{Usage: result.UsageFacts}) require.NoError(t, err) assert.Equal(t, 6.0, cost) }) t.Run("mixed batch selects schema by task rather than adaptor", func(t *testing.T) { adaptor, _, _ := newRequest(t, "image", map[string]any{}) tasks := []*model.Task{ {TaskID: "video-task", Properties: model.Properties{UpstreamModelName: "video"}}, {TaskID: "image-task", Properties: model.Properties{UpstreamModelName: "image"}}, {TaskID: "invalid-image", Properties: model.Properties{UpstreamModelName: "image"}}, } body := []byte(`{"items":[ {"taskId":"image-task","status":"SUCCESS","data":{"usage":{"units":2,"mode":"image"}}}, {"taskId":"video-task","status":"SUCCESS","data":{"usage":{"units":500000,"mode":"video"}}}, {"taskId":"invalid-image","status":"SUCCESS","data":{"usage":{"units":129,"mode":"image"}}} ]}`) results, err := adaptor.ParseBatchResult(tasks, &http.Response{StatusCode: http.StatusOK}, body) require.NoError(t, err) require.Len(t, results, 3) assert.Equal(t, map[string]any{"units": 500000.0, "mode": "video"}, results["video-task"].TaskInfo.UsageFacts) assert.Equal(t, map[string]any{"units": 2.0, "mode": "image"}, results["image-task"].TaskInfo.UsageFacts) assert.Nil(t, results["invalid-image"].TaskInfo.UsageFacts) }) } const documentStreamPlugin = ` export const meta = {apiVersion:1,key:"document-stream",name:"Document stream",version:"1.0.0",author:{name:"Test"},models:["document"],fetchMode:"per_task",submitResponseTypes:["json","sse"],usageSchema:{units:{type:"number",unit:"count"}}}; export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/compile",responseType:ctx.requestBody.responseType || "sse"};} export function parseSubmitEvent(ctx,event,previous) { const chunk = JSON.parse(event.data); if (chunk.error) throw new Error("provider stream failure"); if (chunk.badState) return {state:null}; if (chunk.largeState) return {state:{document:"x".repeat(1048577)},done:chunk.complete === true}; if (chunk.escapedState) return {state:{document:"<".repeat(200000)},done:true}; if (chunk.resetState) return {state:{document:"",units:0},done:true}; const state = Object.assign({}, previous || {document:"",units:0}); state.document += chunk.part || ""; if (chunk.units !== undefined) state.units = chunk.units; state.event = event.event; state.id = event.id; return {state:state,done:chunk.complete === true}; } export function parseSubmitResponse(ctx,response){return {taskId:"vendor-document",taskData:response.body,immediate:{status:"SUCCESS"},state:{revision:1}};} export function parseTaskResult(){return {status:"SUCCESS"};} export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"};} export function extractUsageOnComplete(ctx,result,body){ if(ctx.model!=="alias" || ctx.upstreamModel!=="document" || ctx.taskId!=="vendor-document" || ctx.publicTaskId!=="public-document" || ctx.state.revision!==1 || ctx.data.document!==body.document) throw new Error("invalid completion context"); return {units:body.units}; } ` func TestTaskSubmitStreamContract(t *testing.T) { plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{}) require.NoError(t, err) for _, tc := range []struct { name, body, contentType string valid bool }{ {"multiline and CRLF", ": heartbeat\r\nid: doc-1\r\nevent: update\r\ndata: {\"part\":\r\ndata: \"hello\",\"units\":2}\r\n\r\ndata: {\"part\":\"world\",\"units\":3,\"complete\":true}\n\n", "text/event-stream; charset=utf-8", true}, {"zero actual units", "data: {\"part\":\"free\",\"units\":0,\"complete\":true}\n\n", "text/event-stream", true}, {"premature EOF", "data: {\"part\":\"partial\"}\n\n", "text/event-stream", false}, {"unterminated event", "data: {\"complete\":true}", "text/event-stream", false}, {"provider error", "data: {\"error\":true}\n\n", "text/event-stream", false}, {"invalid hook result", "data: {\"badState\":true}\n\n", "text/event-stream", false}, {"state limit", "data: {\"largeState\":true}\n\n", "text/event-stream", false}, {"intermediate state limit", "data: {\"largeState\":true}\n\ndata: {\"resetState\":true}\n\n", "text/event-stream", false}, {"encoded state limit", "data: {\"escapedState\":true}\n\n", "text/event-stream", false}, {"event limit", "data: " + strings.Repeat("x", maxTaskPluginPersistedJSONBytes) + "\n\n", "text/event-stream", false}, {"wrong response type", "{}", "application/json", false}, {"undeclared stream", "data: {}\n\n", "text/event-stream", false}, } { t.Run(tc.name, func(t *testing.T) { info := &relaycommon.RelayInfo{OriginModelName: "alias", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "document", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "public-document"}} adaptor := New(plugin) adaptor.Init(info) c, recorder := gin.CreateTestContext(httptest.NewRecorder()) _ = recorder c.Request = httptest.NewRequest(http.MethodPost, "/compile", nil) requestBody := map[string]any{} if tc.name == "undeclared stream" { requestBody["responseType"] = "json" } c.Set("task_request", requestBody) require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info)) response := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": {tc.contentType}}, Body: io.NopCloser(strings.NewReader(tc.body))} parsed, taskErr := adaptor.ParseResponse(c, response, info) assert.False(t, c.Writer.Written()) if !tc.valid { require.NotNil(t, taskErr) assert.True(t, taskErr.NoRetry) return } require.Nil(t, taskErr) require.NotNil(t, parsed.Immediate) if tc.name == "zero actual units" { assert.Equal(t, map[string]any{"units": float64(0)}, parsed.Immediate.UsageFacts) return } assert.Equal(t, map[string]any{"units": float64(3)}, parsed.Immediate.UsageFacts) assert.JSONEq(t, `{"document":"helloworld","units":3,"event":"message","id":"doc-1"}`, string(parsed.TaskData)) }) } } func TestTaskSubmitStreamCancellationClosesReader(t *testing.T) { plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{}) require.NoError(t, err) reader, writer := io.Pipe() defer writer.Close() ctx, cancel := context.WithCancel(t.Context()) cancel() _, err = New(plugin).readSubmitEvents(ctx, &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: reader}, nil) require.ErrorIs(t, err, context.Canceled) _, err = writer.Write([]byte("data: {}\n\n")) require.Error(t, err) } func TestTaskSubmitStreamIdleTimeout(t *testing.T) { previous := constant.StreamingTimeout constant.StreamingTimeout = 1 t.Cleanup(func() { constant.StreamingTimeout = previous }) plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{}) require.NoError(t, err) reader, writer := io.Pipe() defer writer.Close() _, err = New(plugin).readSubmitEvents(t.Context(), &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: reader}, nil) require.ErrorContains(t, err, "idle timeout") } func TestPluginJSONValuesPreserveCodecNormalizationAndIsolation(t *testing.T) { for _, value := range []any{ map[string]any{"units": int64(3), "nested": []any{true, " / 图像", math.Copysign(0, -1), nil}}, map[string]any{"empty": []any{}, "null": []any(nil), "object": map[string]any(nil)}, map[string]any{"large": int64(math.MaxInt64), "invalid UTF-8": string([]byte{0xff, 0xfe})}, map[string]any{string([]byte{0xff}): "invalid key"}, map[string]any{"bytes": []byte{0, 1, 255}, "number": json.Number("9007199254740993")}, json.RawMessage(`{"units":2,"enabled":false}`), struct { Units int `json:"units"` }{Units: 0}, } { encoded, err := common.Marshal(value) require.NoError(t, err) var expected any require.NoError(t, common.Unmarshal(encoded, &expected)) assert.Equal(t, expected, jsonValue(value)) } source := map[string]any{"items": []any{map[string]any{"label": "original"}}} copy := jsonValue(source).(map[string]any) copy["items"].([]any)[0].(map[string]any)["label"] = "changed" assert.Equal(t, "original", source["items"].([]any)[0].(map[string]any)["label"]) } func TestTaskSubmitHooksReceiveIndependentRequestSnapshots(t *testing.T) { source := ` export const meta={apiVersion:1,key:"request-copy",name:"Request copy",version:"1.0.0",author:{name:"Test"},models:["copy"],fetchMode:"per_task",usageSchema:{units:{type:"number",unit:"count"}}}; export function buildSubmitRequest(ctx){ctx.requestBody.units=3;ctx.requestBody.nested.label="built";return {url:ctx.baseUrl+"/submit",body:ctx.requestBody};} export function extractUsage(ctx){if(ctx.requestBody.nested.label!=="original")throw new Error("request changed");return {units:ctx.requestBody.units};} export function parseSubmitResponse(){return {taskId:"copy"};} export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"};} export function parseTaskResult(){return {status:"SUCCESS"};} ` plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{}) require.NoError(t, err) request := map[string]any{"units": int64(2), "nested": map[string]any{"label": "original"}} c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/submit", nil) c.Set("task_request", request) info := &relaycommon.RelayInfo{OriginModelName: "copy", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "copy", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}} adaptor := New(plugin) adaptor.Init(info) require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info)) facts, err := adaptor.ExtractUsageFactsValidated(c, info) require.NoError(t, err) assert.Equal(t, map[string]any{"units": float64(2)}, facts) body, err := adaptor.BuildRequestBody(c, info) require.NoError(t, err) encoded, err := io.ReadAll(body) require.NoError(t, err) assert.JSONEq(t, `{"units":3,"nested":{"label":"built"}}`, string(encoded)) assert.Equal(t, int64(2), request["units"]) assert.Equal(t, "original", request["nested"].(map[string]any)["label"]) } func TestTaskSubmitDeltaStreamContract(t *testing.T) { source := strings.Replace(documentStreamPlugin, `submitResponseTypes:["json","sse"],`, `submitResponseTypes:["json","sse"],requiredCapabilities:["submit-sse-delta@1"],`, 1) + ` export function parseSubmitEventDelta(ctx,event,previous) { if(previous && previous.document !== undefined) throw new Error("full result leaked into control state"); const chunk=JSON.parse(event.data); let changes=chunk.changes; if(chunk.largeResult) changes=[{op:"set",path:[],value:{document:"<".repeat(200000)}}]; if(chunk.resetAfterLarge) changes.push({op:"set",path:[],value:{document:"",units:0}}); const result={changes:changes,state:{events:(previous ? previous.events : 0)+1},done:chunk.complete===true}; if(chunk.largeControl) result.state={text:"x".repeat(65537)}; if(chunk.extra) result.extra=true; return result; }` plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{}) require.NoError(t, err) const first = `{"changes":[{"op":"set","path":[],"value":{"document":"hello","units":2}}]}` const last = `{"changes":[{"op":"appendText","path":["document"],"value":"world"},{"op":"set","path":["units"],"value":0}],"complete":true}` for _, tc := range []struct { name string frames []string valid bool }{ {"control state and zero usage", []string{first, last}, true}, {"oversized intermediate result", []string{`{"largeResult":true}`, last}, false}, {"oversized operation before reset", []string{`{"largeResult":true,"resetAfterLarge":true,"complete":true}`}, false}, {"control state limit", []string{first, `{"changes":[],"largeControl":true,"complete":true}`}, false}, {"missing changes", []string{`{"complete":true}`}, false}, {"extra result fields", []string{first, `{"changes":[],"extra":true,"complete":true}`}, false}, {"unfinished delta stream", []string{first}, false}, } { t.Run(tc.name, func(t *testing.T) { var stream strings.Builder for _, frame := range tc.frames { stream.WriteString("data: " + frame + "\n\n") } info := &relaycommon.RelayInfo{OriginModelName: "alias", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "document", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "public-document"}} adaptor := New(plugin) adaptor.Init(info) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/compile", nil) c.Set("task_request", map[string]any{}) require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info)) response := &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(stream.String()))} parsed, taskErr := adaptor.ParseResponse(c, response, info) assert.False(t, c.Writer.Written()) if !tc.valid { require.NotNil(t, taskErr) assert.True(t, taskErr.NoRetry) return } require.Nil(t, taskErr) require.NotNil(t, parsed.Immediate) assert.JSONEq(t, `{"document":"helloworld","units":0}`, string(parsed.TaskData)) assert.Equal(t, map[string]any{"units": float64(0)}, parsed.Immediate.UsageFacts) }) } } func TestAlibabaSubmitDeltaDoesNotMutateControlState(t *testing.T) { source, err := plugins.Source("alibaba") require.NoError(t, err) plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{}) require.NoError(t, err) const before = `{"request_id":"old","usage":{"image_count":1},"output":{"finished":false,"choices":[{"message":{"role":"assistant","content":[{"text":"hello"}]}}]}}` const control = `{"choices":[{"count":1,"lastText":true,"finishReason":""}],"hasUsage":true}` var previous, initial any require.NoError(t, common.UnmarshalJsonStr(control, &previous)) require.NoError(t, common.UnmarshalJsonStr(before, &initial)) accumulated := pluginruntime.NewJSONState(maxTaskPluginPersistedJSONBytes) require.NoError(t, accumulated.Apply(t.Context(), []any{map[string]any{"op": "set", "path": []any{}, "value": initial}})) value, err := plugin.Engine.Call(t.Context(), "parseSubmitEventDelta", map[string]any{}, map[string]any{ "event": "message", "data": `{"request_id":"new","usage":{"image_count":2},"output":{"choices":[{"finish_reason":"stop","message":{"content":[{"text":" world"},{"image":"https://cdn.example/image.png"}]}}]}}`, }, previous) require.NoError(t, err) unchanged, err := common.Marshal(previous) require.NoError(t, err) assert.JSONEq(t, control, string(unchanged)) result := value.(map[string]any) assert.Equal(t, true, result["done"]) require.NoError(t, accumulated.Apply(t.Context(), result["changes"])) final, err := accumulated.Value() require.NoError(t, err) encoded, err := common.Marshal(final) require.NoError(t, err) assert.JSONEq(t, `{"request_id":"new","usage":{"image_count":2},"output":{"finished":true,"choices":[{"finish_reason":"stop","message":{"role":"assistant","content":[{"text":"hello world"},{"image":"https://cdn.example/image.png"}]}}]}}`, string(encoded)) encoded, err = common.Marshal(result["state"]) require.NoError(t, err) assert.JSONEq(t, `{"choices":[{"count":2,"lastText":false,"finishReason":"stop"}],"hasUsage":true}`, string(encoded)) }