cneicy/new-api-custom
1
0 Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

feat: publish sanitized custom.97 source

Browse Source
This commit is contained in:
cneicy committed 2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
1022 files changed
+69903 -19207
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

relay/channel/task/jsplugin/adaptor.go
+374 -142
View File
@@ -2,9 +2,10 @@ package jsplugin
import (
"bytes"
"cmp"
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"maps"
@@ -14,6 +15,7 @@ import (
"net/http"
"net/textproto"
"net/url"
"reflect"
"regexp"
"slices"
"strconv"
@@ -46,6 +48,26 @@ type requestDescriptor struct {
RewriteModel string `json:"rewriteModel"`
BodyType string `json:"bodyType"`
Parts []requestPart `json:"parts"`
// BodyText is the body as JSON.stringify wrote it, for a plugin that
// preserves JSON order: the upstream gets the members in that order.
BodyText json.RawMessage `json:"-"`
}
// orderedRequestDescriptor is requestDescriptor as a plugin that preserves
// JSON order returns it: decoded from the JSON.stringify text of the result,
// with the body kept as that text.
type orderedRequestDescriptor struct {
ResponseType string `json:"responseType"`
URL string `json:"url"`
Method string `json:"method"`
Headers map[string]string `json:"headers"`
Body json.RawMessage `json:"body"`
Credentialless bool `json:"credentialless"`
Action string `json:"action"`
Model string `json:"model"`
RewriteModel string `json:"rewriteModel"`
BodyType string `json:"bodyType"`
Parts []requestPart `json:"parts"`
}
type requestPart struct {
@@ -88,6 +110,12 @@ type TaskAdaptor struct {
routeRequest *pluginruntime.RouteRequestContext
requestHeaders map[string]string
files []map[string]any
// hookRequestBody is hookRequestSource as hooks receive it.
hookRequestSource map[string]any
hookRequestBody any
// requestBodyText is the decoded requestBody as JSON text, which hooks
// of a plugin that preserves JSON order receive instead.
requestBodyText json.RawMessage
}
func New(plugin *pluginruntime.LoadedPlugin) *TaskAdaptor { return &TaskAdaptor{plugin: plugin} }
@@ -98,7 +126,10 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
if pinned, ok := pinnedValue.(pluginruntime.PinnedEndpoint); ok && pinned.Plugin == a.plugin {
if protocolValue, present := c.Get(pluginruntime.ContextKeyProtocolRequest); present {
if protocolContext, valid := protocolValue.(pluginruntime.ProtocolRequestContext); valid {
resolvedValue, callErr := a.plugin.Engine.CallPath(context.WithoutCancel(c.Request.Context()), "protocols", []string{pinned.Protocol, "decodeRequest"}, protocolContext.JSValue())
resolvedValue, requestBodyText, callErr := a.plugin.Engine.CallPathWithMemberJSON(
context.WithoutCancel(c.Request.Context()), 0, a.plugin.Meta.JSONTextMember("requestBody"),
"protocols", []string{pinned.Protocol, "decodeRequest"}, protocolContext.JSValueFor(a.plugin.Meta),
)
resolved, resolvedOK := resolvedValue.(map[string]any)
resolvedModel, modelOK := resolved["model"].(string)
if callErr != nil || !resolvedOK || !modelOK || resolvedModel != pinned.Model {
@@ -110,6 +141,9 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
if body, present := resolved["requestBody"]; present {
c.Set("task_request", body)
}
if requestBodyText != nil {
c.Set(pluginruntime.ContextKeyRequestBodyText, requestBodyText)
}
if action, valid := resolved["action"].(string); valid && strings.TrimSpace(action) != "" {
c.Set("task_action", action)
info.Action = action
@@ -126,18 +160,23 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
request, hasRequest := c.Get("task_request")
hasUsageProfiles := len(a.plugin.Meta.UsageProfiles) > 0
if hasRequest && !hasUsageProfiles {
if err := a.validateResolvedUsageRequest(request, ""); err != nil {
if err := a.validateResolvedUsageRequest(request); err != nil {
return service.TaskErrorWrapperLocal(err, "plugin_usage_invalid", http.StatusBadRequest)
}
}
if _, err := a.buildSubmit(c, info); err != nil {
return service.TaskErrorWrapperLocal(err, "plugin_request_invalid", http.StatusBadRequest)
// utils.fetch failures in the build hook name the host fetched: the
// channel base URL or one of the plugin's allowed hosts, which an
// installed plugin may point at private hosts.
taskErr := service.TaskErrorWrapper(err, "plugin_request_invalid", http.StatusBadRequest, append(common.UpstreamAddresses(c), a.plugin.Meta.AllowedHosts...)...)
taskErr.LocalError = true
return taskErr
}
// The descriptor may rewrite the model. Validate profiled requests against
// that final model before any quota calculation or upstream submission.
// that final model before any quota calculation or upstream submission; an
// upstream endpoint ID without a profile keeps the client model's profile.
if hasRequest && hasUsageProfiles {
usageModel := cmp.Or(info.UpstreamModelName, info.OriginModelName)
if err := a.validateResolvedUsageRequest(request, usageModel); err != nil {
if err := a.validateResolvedUsageRequest(request, info.UpstreamModelName, info.OriginModelName); err != nil {
return service.TaskErrorWrapperLocal(err, "plugin_usage_invalid", http.StatusBadRequest)
}
}
@@ -156,7 +195,7 @@ func (a *TaskAdaptor) EstimateBilling(c *gin.Context, info *relaycommon.RelayInf
func (a *TaskAdaptor) EstimateBillingValidated(c *gin.Context, info *relaycommon.RelayInfo) (map[string]float64, error) {
usageContext := a.submitContext(c, info)
usageContext["usagePurpose"] = "billing_ratios"
return a.usageRatios(c.Request.Context(), cmp.Or(info.UpstreamModelName, info.OriginModelName), "extractUsage", usageContext)
return a.usageRatios(c.Request.Context(), []string{info.UpstreamModelName, info.OriginModelName}, "extractUsage", usageContext)
}
func (a *TaskAdaptor) ExtractUsageFacts(c *gin.Context, info *relaycommon.RelayInfo) map[string]any {
@@ -185,10 +224,11 @@ func (a *TaskAdaptor) ExtractUsageFactsValidated(c *gin.Context, info *relaycomm
if !ok {
return nil, fmt.Errorf("plugin usage hook must return an object")
}
if _, err = a.validatedUsageRatios(facts, cmp.Or(info.UpstreamModelName, info.OriginModelName)); err != nil {
validated, _, err := a.validatedUsageRatios(facts, info.UpstreamModelName, info.OriginModelName)
if err != nil {
return nil, err
}
return facts, nil
return validated, nil
}
func (a *TaskAdaptor) AdjustBillingOnSubmit(info *relaycommon.RelayInfo, taskData []byte) map[string]float64 {
@@ -196,7 +236,7 @@ func (a *TaskAdaptor) AdjustBillingOnSubmit(info *relaycommon.RelayInfo, taskDat
if err := common.Unmarshal(taskData, &data); err != nil {
data = string(taskData)
}
ratios, err := a.usageRatios(context.Background(), cmp.Or(info.UpstreamModelName, info.OriginModelName), "extractUsageOnSubmit", a.submitContext(nil, info), data)
ratios, err := a.usageRatios(context.Background(), []string{info.UpstreamModelName, info.OriginModelName}, "extractUsageOnSubmit", a.submitContext(nil, info), data)
if err != nil {
a.logRejectedUsage("extractUsageOnSubmit", err)
return nil
@@ -212,11 +252,11 @@ func (a *TaskAdaptor) AdjustBillingOnComplete(task *model.Task, result *relaycom
if err != nil {
return 0
}
usageModel := ""
if task != nil {
usageModel = cmp.Or(task.Properties.UpstreamModelName, task.Properties.OriginModelName)
if task == nil {
a.applyCompletionUsageFacts(result, value)
return 0
}
a.applyCompletionUsageFacts(result, value, usageModel)
a.applyCompletionUsageFacts(result, value, task.Properties.UpstreamModelName, task.Properties.OriginModelName)
return 0
}
@@ -227,13 +267,18 @@ func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, erro
return a.submit.URL, pluginruntime.ValidateRequestURL(a.submit.URL, info.ChannelBaseUrl, a.plugin.Meta.AllowedHosts)
}
func (a *TaskAdaptor) BuildRequestHeader(_ *gin.Context, req *http.Request, _ *relaycommon.RelayInfo) error {
func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, _ *relaycommon.RelayInfo) error {
if a.submit == nil {
return fmt.Errorf("plugin submit request was not built")
}
for name, value := range a.submit.Headers {
req.Header.Set(name, value)
}
// BuildRequestBody rebuilt the multipart body with a new boundary and put its
// Content-Type on the client request; a plugin's own Content-Type cannot match it.
if a.submit.BodyType == "multipart" {
req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
}
return nil
}
@@ -267,18 +312,18 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
}
continue
}
field := strings.TrimPrefix(part.FileRef, "request_file:")
field, index, validRef := pluginruntime.ParseFileReference(part.FileRef)
files := form.File[field]
if len(files) == 0 {
if !validRef || index >= len(files) {
return nil, fmt.Errorf("unknown file reference %q", part.FileRef)
}
file, openErr := files[0].Open()
file, openErr := files[index].Open()
if openErr != nil {
return nil, openErr
}
filename := part.Filename
if filename == "" {
filename = files[0].Filename
filename = files[index].Filename
}
header := make(textproto.MIMEHeader)
disposition := mime.FormatMediaType("form-data", map[string]string{"name": part.Name, "filename": filename})
@@ -287,7 +332,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
return nil, fmt.Errorf("invalid multipart name or filename")
}
header.Set("Content-Disposition", disposition)
header.Set("Content-Type", files[0].Header.Get("Content-Type"))
header.Set("Content-Type", files[index].Header.Get("Content-Type"))
destination, copyErr := writer.CreatePart(header)
if copyErr == nil {
_, copyErr = io.Copy(destination, file)
@@ -309,7 +354,18 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
if text, ok := descriptor.Body.(string); ok {
return strings.NewReader(text), nil
}
inlined, err := inlineJSONFilePlaceholders(c, descriptor.Body)
value := descriptor.Body
if descriptor.BodyText != nil {
if !bytes.Contains(descriptor.BodyText, []byte(`"__fileRef":`)) {
return bytes.NewReader(descriptor.BodyText), nil
}
// File placeholders resolve only in the decoded value, which writes
// the members sorted; uploads carry no client JSON order to keep.
if err := common.Unmarshal(descriptor.BodyText, &value); err != nil {
return nil, err
}
}
inlined, err := inlineJSONFilePlaceholders(c, value)
if err != nil {
return nil, err
}
@@ -391,12 +447,12 @@ func encodeFilePlaceholder(placeholder map[string]any, form *multipart.Form, lim
if form == nil {
return "", fmt.Errorf("unknown file reference %q", ref)
}
field := strings.TrimPrefix(ref, "request_file:")
field, index, validRef := pluginruntime.ParseFileReference(ref)
files := form.File[field]
if len(files) == 0 {
if !validRef || index >= len(files) {
return "", fmt.Errorf("unknown file reference %q", ref)
}
header := files[0]
header := files[index]
maxBytes := limit
if raw, exists := placeholder["maxBytes"]; exists {
n, ok := usageNumber(raw, false)
@@ -484,7 +540,7 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r
}
}
if err != nil {
failure := service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusBadGateway)
failure := service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusBadGateway, common.UpstreamAddresses(c)...)
failure.NoRetry = streaming
return nil, failure
}
@@ -499,7 +555,7 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r
resp.StatusCode,
time.Since(started).Milliseconds(),
)
return nil, service.TaskErrorWrapper(err, "plugin_submit_response_failed", http.StatusBadGateway)
return nil, service.TaskErrorWrapper(err, "plugin_submit_response_failed", http.StatusBadGateway, common.UpstreamAddresses(c)...)
}
if object, ok := value.(map[string]any); ok {
if _, forbidden := object["clientResponse"]; forbidden {
@@ -518,7 +574,7 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r
resp.StatusCode,
time.Since(started).Milliseconds(),
)
return nil, service.TaskErrorWrapper(err, "plugin_submit_response_invalid", http.StatusBadGateway)
return nil, service.TaskErrorWrapper(err, "plugin_submit_response_invalid", http.StatusBadGateway, common.UpstreamAddresses(c)...)
}
var taskData []byte
if parsed.TaskData != nil {
@@ -527,12 +583,12 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r
if err == nil {
err = fmt.Errorf("task data exceeds size limit")
}
return nil, service.TaskErrorWrapper(err, "plugin_submit_response_invalid", http.StatusBadGateway)
return nil, service.TaskErrorWrapper(err, "plugin_submit_response_invalid", http.StatusBadGateway, common.UpstreamAddresses(c)...)
}
}
pluginState, _ := encodeReturnedPluginState(value)
if len(pluginState) > maxTaskPluginPersistedJSONBytes {
logger.LogWarn(c, fmt.Sprintf("task plugin %s rejected oversized submit state (%d bytes)", a.plugin.Meta.Key, len(pluginState)))
logger.LogWarn(c, common.LogText("task plugin %s rejected oversized submit state (%d bytes)", a.plugin.Meta.Key, len(pluginState)))
pluginState = nil
}
response := &channel.TaskSubmitResponse{
@@ -562,16 +618,16 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r
task.PrivateData.PluginState = pluginState
ctx, err := a.queryContext(task, info.ApiKey, info.ChannelBaseUrl, info.ChannelSetting.Proxy)
if err != nil {
logger.LogWarn(c, fmt.Sprintf("task plugin %s completion context failed; retaining reserved quota: %v", a.plugin.Meta.Key, err))
logger.LogWarn(c, common.LogText("task plugin %s completion context failed; retaining reserved quota: %v", a.plugin.Meta.Key, err))
return response, nil
}
facts, err := a.plugin.Engine.Call(c.Request.Context(), "extractUsageOnComplete", ctx, jsonValue(immediate), parsed.TaskData)
if err != nil {
logger.LogWarn(c, fmt.Sprintf("task plugin %s completion usage failed; retaining reserved quota: %v", a.plugin.Meta.Key, err))
logger.LogWarn(c, common.LogText("task plugin %s completion usage failed; retaining reserved quota: %v", a.plugin.Meta.Key, err))
return response, nil
}
if err := a.applyCompletionUsageFacts(immediate, facts, cmp.Or(info.UpstreamModelName, info.OriginModelName)); err != nil {
logger.LogWarn(c, fmt.Sprintf("task plugin %s completion usage rejected; retaining reserved quota: %v", a.plugin.Meta.Key, err))
if err := a.applyCompletionUsageFacts(immediate, facts, info.UpstreamModelName, info.OriginModelName); err != nil {
logger.LogWarn(c, common.LogText("task plugin %s completion usage rejected; retaining reserved quota: %v", a.plugin.Meta.Key, err))
}
return response, nil
}
@@ -593,11 +649,11 @@ func (a *TaskAdaptor) FetchBatchTasks(baseURL, key string, tasks []*model.Task,
if err != nil {
return nil, err
}
value, err := a.plugin.Engine.Call(context.Background(), "buildBatchQueryRequest", ctx, taskContexts)
descriptor, err := a.requestDescriptor(pluginruntime.WithFetcher(context.Background(), a.hookFetcher(baseURL, proxy)), "buildBatchQueryRequest", ctx, taskContexts)
if err != nil {
return nil, err
}
return a.doFetchDescriptor(baseURL, proxy, value)
return a.doFetchDescriptor(baseURL, proxy, descriptor)
}
func (a *TaskAdaptor) FetchTask(baseURL, key string, task *model.Task, proxy string) (*http.Response, error) {
@@ -605,18 +661,60 @@ func (a *TaskAdaptor) FetchTask(baseURL, key string, task *model.Task, proxy str
if err != nil {
return nil, err
}
value, err := a.plugin.Engine.Call(context.Background(), "buildQueryRequest", ctx)
descriptor, err := a.requestDescriptor(pluginruntime.WithFetcher(context.Background(), a.hookFetcher(baseURL, proxy)), "buildQueryRequest", ctx)
if err != nil {
return nil, err
}
return a.doFetchDescriptor(baseURL, proxy, value)
return a.doFetchDescriptor(baseURL, proxy, descriptor)
}
func (a *TaskAdaptor) doFetchDescriptor(baseURL, proxy string, value any) (*http.Response, error) {
var descriptor requestDescriptor
if err := convert(value, &descriptor); err != nil {
return nil, err
// hookFetcher sends a driver hook's utils.fetch requests to the hosts its
// request descriptors may use, through the channel's proxy like poll requests.
func (a *TaskAdaptor) hookFetcher(baseURL, proxy string) pluginruntime.Fetcher {
return func(ctx context.Context, request pluginruntime.FetchRequest) (pluginruntime.FetchResponse, error) {
if err := pluginruntime.ValidateRequestURL(request.URL, baseURL, a.plugin.Meta.AllowedHosts); err != nil {
return pluginruntime.FetchResponse{}, err
}
client, err := service.GetHttpClientWithProxySettings(proxy, a.channelSetting())
if err != nil {
return pluginruntime.FetchResponse{}, err
}
// A redirect reaches the plugin as it is; fetching its location is a
// new request, checked and counted again.
noRedirects := *client
noRedirects.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
req, err := http.NewRequestWithContext(ctx, request.Method, request.URL, nil)
if err != nil {
return pluginruntime.FetchResponse{}, err
}
for name, value := range request.Headers {
req.Header.Set(name, value)
}
started := time.Now()
resp, err := noRedirects.Do(req)
if err != nil {
logger.LogDebug(ctx, "task_plugin subsystem=fetch event=failed plugin=%q hook=%q method=%q host=%q reason=transport_error elapsed_ms=%d",
a.plugin.Meta.Key, request.Hook, request.Method, req.URL.Host, time.Since(started).Milliseconds())
return pluginruntime.FetchResponse{}, err
}
defer resp.Body.Close()
body, err := io.ReadAll(io.LimitReader(resp.Body, pluginruntime.MaxFetchBodyBytes+1))
if err != nil {
return pluginruntime.FetchResponse{}, err
}
headers := make(map[string]string, len(resp.Header))
for name, values := range resp.Header {
if len(values) > 0 {
headers[name] = values[0]
}
}
logger.LogDebug(ctx, "task_plugin subsystem=fetch event=done plugin=%q hook=%q method=%q host=%q status=%d bytes=%d elapsed_ms=%d",
a.plugin.Meta.Key, request.Hook, request.Method, req.URL.Host, resp.StatusCode, len(body), time.Since(started).Milliseconds())
return pluginruntime.FetchResponse{Status: resp.StatusCode, Headers: headers, Body: body}, nil
}
}
func (a *TaskAdaptor) doFetchDescriptor(baseURL, proxy string, descriptor requestDescriptor) (*http.Response, error) {
if err := pluginruntime.ValidateRequestURL(descriptor.URL, baseURL, a.plugin.Meta.AllowedHosts); err != nil {
return nil, err
}
@@ -624,6 +722,8 @@ func (a *TaskAdaptor) doFetchDescriptor(baseURL, proxy string, value any) (*http
if descriptor.Body != nil {
if bodyText, ok := descriptor.Body.(string); ok {
requestBody = strings.NewReader(bodyText)
} else if descriptor.BodyText != nil {
requestBody = bytes.NewReader(descriptor.BodyText)
} else {
encoded, marshalErr := common.Marshal(descriptor.Body)
if marshalErr != nil {
@@ -643,7 +743,7 @@ func (a *TaskAdaptor) doFetchDescriptor(baseURL, proxy string, value any) (*http
for name, value := range descriptor.Headers {
req.Header.Set(name, value)
}
client, err := service.GetHttpClientWithProxy(proxy)
client, err := service.GetHttpClientWithProxySettings(proxy, a.channelSetting())
if err != nil {
return nil, err
}
@@ -690,11 +790,6 @@ func (a *TaskAdaptor) ParseBatchResult(tasks []*model.Task, resp *http.Response,
if err != nil {
return nil, err
}
value, err := a.plugin.Engine.Call(context.Background(), "parseBatchResult", ctx, input, hookHTTPResponse(resp))
if err != nil {
logger.LogDebug(context.Background(), "task_plugin subsystem=adaptor event=parse_batch_failed plugin=%q reason=hook_failed body_bytes=%d elapsed_ms=%d", a.plugin.Meta.Key, len(body), time.Since(started).Milliseconds())
return nil, err
}
var parsed []struct {
TaskID string `json:"taskId"`
Action string `json:"action"`
@@ -708,8 +803,13 @@ func (a *TaskAdaptor) ParseBatchResult(tasks []*model.Task, resp *http.Response,
Data any `json:"data"`
State any `json:"state"`
}
if err = convert(value, &parsed); err != nil {
logger.LogDebug(context.Background(), "task_plugin subsystem=adaptor event=parse_batch_failed plugin=%q reason=invalid_result body_bytes=%d elapsed_ms=%d", a.plugin.Meta.Key, len(body), time.Since(started).Milliseconds())
if err = a.plugin.Engine.CallInto(pluginruntime.WithFetcher(context.Background(), a.hookFetcher(baseURL, proxy)), &parsed, "parseBatchResult", ctx, input, hookHTTPResponse(resp)); err != nil {
reason := "hook_failed"
var invalid *pluginruntime.ResultError
if errors.As(err, &invalid) {
reason = "invalid_result"
}
logger.LogDebug(context.Background(), "task_plugin subsystem=adaptor event=parse_batch_failed plugin=%q reason=%s body_bytes=%d elapsed_ms=%d", a.plugin.Meta.Key, reason, len(body), time.Since(started).Milliseconds())
return nil, err
}
results := make(map[string]*service.BatchTaskResult, len(parsed))
@@ -722,7 +822,7 @@ func (a *TaskAdaptor) ParseBatchResult(tasks []*model.Task, resp *http.Response,
if item.State != nil {
pluginState, marshalErr := common.Marshal(item.State)
if marshalErr != nil || len(pluginState) > maxTaskPluginPersistedJSONBytes {
logger.LogWarn(context.Background(), fmt.Sprintf("task plugin %s rejected invalid or oversized poll state", a.plugin.Meta.Key))
logger.LogWarn(context.Background(), common.LogText("task plugin %s rejected invalid or oversized poll state", a.plugin.Meta.Key))
} else {
info.PluginState = pluginState
}
@@ -741,8 +841,9 @@ func (a *TaskAdaptor) ParseBatchResult(tasks []*model.Task, resp *http.Response,
}
facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", itemCtx, jsonValue(&info), usageBody)
if hookErr == nil {
usageModel, _ := itemCtx["upstreamModel"].(string)
a.applyCompletionUsageFacts(&info, facts, usageModel)
upstreamModel, _ := itemCtx["upstreamModel"].(string)
originModel, _ := itemCtx["model"].(string)
a.applyCompletionUsageFacts(&info, facts, upstreamModel, originModel)
}
}
results[item.TaskID] = &service.BatchTaskResult{TaskInfo: info, Action: item.Action, SubmitTime: item.SubmitTime, StartTime: item.StartTime, FinishTime: item.FinishTime, Data: item.Data}
@@ -774,7 +875,7 @@ func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, bod
if err != nil {
return nil, err
}
value, err := a.plugin.Engine.Call(context.Background(), "parseTaskResult", ctx, input, hookHTTPResponse(resp))
value, err := a.plugin.Engine.Call(pluginruntime.WithFetcher(context.Background(), a.hookFetcher(baseURL, proxy)), "parseTaskResult", ctx, input, hookHTTPResponse(resp))
if err != nil {
logger.LogDebug(context.Background(), "task_plugin subsystem=adaptor event=parse_task_failed plugin=%q reason=hook_failed body_bytes=%d elapsed_ms=%d", a.plugin.Meta.Key, len(body), time.Since(started).Milliseconds())
return nil, err
@@ -797,7 +898,7 @@ func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, bod
}
if pluginState, present := encodeReturnedPluginState(value); present {
if len(pluginState) > maxTaskPluginPersistedJSONBytes {
logger.LogWarn(context.Background(), fmt.Sprintf("task plugin %s rejected oversized poll state (%d bytes)", a.plugin.Meta.Key, len(pluginState)))
logger.LogWarn(context.Background(), common.LogText("task plugin %s rejected oversized poll state (%d bytes)", a.plugin.Meta.Key, len(pluginState)))
} else {
result.PluginState = pluginState
}
@@ -807,8 +908,9 @@ func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, bod
if a.hasHook(context.Background(), "extractUsageOnComplete") {
facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", ctx, jsonValue(result), input)
if hookErr == nil {
usageModel, _ := ctx["upstreamModel"].(string)
a.applyCompletionUsageFacts(result, facts, usageModel)
upstreamModel, _ := ctx["upstreamModel"].(string)
originModel, _ := ctx["model"].(string)
a.applyCompletionUsageFacts(result, facts, upstreamModel, originModel)
}
}
taskStatus := model.TaskStatus(result.Status)
@@ -823,8 +925,8 @@ func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, bod
return result, nil
}
func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, facts any, modelName string) error {
values, err := a.validatedCompletionUsageFacts(facts, modelName)
func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, facts any, models ...string) error {
values, err := a.validatedCompletionUsageFacts(facts, models...)
if err != nil {
a.logRejectedUsage("extractUsageOnComplete", err)
return err
@@ -861,16 +963,19 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
if err != nil {
return nil, err
}
value, err := a.plugin.Engine.CallPath(context.Background(), "protocols", []string{"openai_video", "render"}, map[string]any{"protocol": "openai_video", "operation": "retrieve"}, jsonValue(view))
encodedView, err := common.Marshal(view)
if err != nil {
return nil, err
}
encoded, err := common.Marshal(value)
// Decoding into any leaves only encoding failures to CallPathInto, so they
// keep the codec's message.
var value any
err = a.plugin.Engine.CallPathInto(context.Background(), &value, "protocols", []string{"openai_video", "render"}, map[string]any{"protocol": "openai_video", "operation": "retrieve"}, pluginruntime.RawJSON(encodedView))
if err != nil {
return nil, err
}
var rendered map[string]any
if err = common.Unmarshal(encoded, &rendered); err != nil || rendered == nil {
rendered, _ := value.(map[string]any)
if rendered == nil {
return nil, fmt.Errorf("plugin returned an invalid OpenAI video object")
}
// Keep provider extensions intact while the host owns the public task's
@@ -906,7 +1011,7 @@ func (a *TaskAdaptor) ListArtifacts(task *model.Task) ([]channel.TaskArtifact, e
return validateTaskArtifacts(value)
}
func (a *TaskAdaptor) BuildContentRequest(task *model.Task, artifactKey string, clientRequest channel.TaskArtifactClientRequest) (*channel.TaskContentRequest, error) {
func (a *TaskAdaptor) BuildContentRequest(ctx context.Context, task *model.Task, artifactKey string, clientRequest channel.TaskArtifactClientRequest) (*channel.TaskContentRequest, error) {
if !a.hasHook(context.Background(), "buildContentRequest") {
return nil, nil
}
@@ -916,32 +1021,22 @@ func (a *TaskAdaptor) BuildContentRequest(task *model.Task, artifactKey string,
if !taskArtifactKeyPattern.MatchString(artifactKey) {
return nil, fmt.Errorf("invalid artifact key")
}
ctx, err := taskArtifactContext(task)
hookContext, err := taskArtifactContext(task)
if err != nil {
return nil, err
}
ctx["upstreamTaskId"] = task.GetUpstreamTaskID()
ctx["artifactKey"] = artifactKey
ctx["baseUrl"] = a.info.ChannelBaseUrl
ctx["clientRequest"] = jsonValue(clientRequest)
proxy := a.info.ChannelSetting.Proxy
auth, err := resolveAuth(a.plugin.Meta.Auth, a.info.ApiKey, proxy)
if err != nil {
hookContext["upstreamTaskId"] = task.GetUpstreamTaskID()
hookContext["artifactKey"] = artifactKey
hookContext["baseUrl"] = a.info.ChannelBaseUrl
hookContext["clientRequest"] = jsonValue(clientRequest)
if err = a.applyUpstreamCredentials(hookContext, a.info.ChannelType, a.info.ApiKey, a.info.ChannelSetting.Proxy); err != nil {
return nil, err
}
ctx["auth"] = auth
ctx["authHeader"] = auth["authHeader"]
if a.plugin.Meta.Auth.Type == "" || a.plugin.Meta.Auth.Type == "none" || a.plugin.Meta.Auth.Type == "api_key" {
ctx["apiKey"] = a.info.ApiKey
}
value, err := a.plugin.Engine.Call(context.Background(), "buildContentRequest", ctx)
fetchContext := pluginruntime.WithFetcher(ctx, a.hookFetcher(a.info.ChannelBaseUrl, a.info.ChannelSetting.Proxy))
descriptor, err := a.requestDescriptor(fetchContext, "buildContentRequest", hookContext)
if err != nil {
return nil, err
}
var descriptor requestDescriptor
if err = convert(value, &descriptor); err != nil {
return nil, err
}
method := strings.ToUpper(strings.TrimSpace(descriptor.Method))
if method == "" {
method = strings.ToUpper(strings.TrimSpace(clientRequest.Method))
@@ -970,6 +1065,8 @@ func (a *TaskAdaptor) BuildContentRequest(task *model.Task, artifactKey string,
if descriptor.Body != nil {
if text, ok := descriptor.Body.(string); ok {
body = []byte(text)
} else if descriptor.BodyText != nil {
body = descriptor.BodyText
} else {
body, err = common.Marshal(descriptor.Body)
if err != nil {
@@ -1056,29 +1153,17 @@ func (a *TaskAdaptor) queryContext(task *model.Task, key, baseURL, proxy string)
key = task.PrivateData.Key
}
}
auth, err := resolveAuth(a.plugin.Meta.Auth, key, proxy)
if err != nil {
if err := a.applyUpstreamCredentials(ctx, a.channelType(), key, proxy); err != nil {
return nil, err
}
ctx["auth"] = auth
ctx["authHeader"] = auth["authHeader"]
if a.plugin.Meta.Auth.Type == "" || a.plugin.Meta.Auth.Type == "none" || a.plugin.Meta.Auth.Type == "api_key" {
ctx["apiKey"] = key
}
return ctx, nil
}
func (a *TaskAdaptor) batchQueryContext(key, baseURL, proxy string, tasks []map[string]any) (map[string]any, error) {
ctx := map[string]any{"baseUrl": baseURL, "tasks": tasks}
auth, err := resolveAuth(a.plugin.Meta.Auth, key, proxy)
if err != nil {
if err := a.applyUpstreamCredentials(ctx, a.channelType(), key, proxy); err != nil {
return nil, err
}
ctx["auth"] = auth
ctx["authHeader"] = auth["authHeader"]
if a.plugin.Meta.Auth.Type == "" || a.plugin.Meta.Auth.Type == "none" || a.plugin.Meta.Auth.Type == "api_key" {
ctx["apiKey"] = key
}
return ctx, nil
}
@@ -1089,6 +1174,50 @@ func (a *TaskAdaptor) queryCredentials() (key, baseURL, proxy string) {
return a.info.ApiKey, a.info.ChannelBaseUrl, a.info.ChannelSetting.Proxy
}
// channelType is the executing channel's type when polling or content
// fetches run with channel metadata; zero otherwise, which is a vendor upstream.
func (a *TaskAdaptor) channelType() int {
if a.info == nil || !a.info.HasChannelMeta() {
return 0
}
return a.info.ChannelType
}
// channelSetting is the executing channel's settings when the adaptor runs
// with channel metadata; zero otherwise.
func (a *TaskAdaptor) channelSetting() kitdto.ChannelSettings {
if a.info == nil || !a.info.HasChannelMeta() {
return kitdto.ChannelSettings{}
}
return a.info.ChannelSetting
}
// applyUpstreamCredentials sets ctx.upstream together with the credentials
// every driver hook reads. A New API channel (type 60) is another gateway:
// the plugin must address its own prefixed native routes there, and the
// channel key is that gateway's token, so it is sent as a Bearer header for
// every auth type instead of running the plugin's vendor auth scheme.
func (a *TaskAdaptor) applyUpstreamCredentials(ctx map[string]any, channelType int, key, proxy string) error {
if channelType == constant.ChannelTypeNewAPI {
ctx["upstream"] = map[string]any{"kind": pluginruntime.UpstreamKindNewAPI}
ctx["auth"] = map[string]any{"authHeader": "Bearer " + key}
ctx["authHeader"] = "Bearer " + key
ctx["apiKey"] = key
return nil
}
ctx["upstream"] = map[string]any{"kind": pluginruntime.UpstreamKindVendor}
auth, err := resolveAuth(a.plugin.Meta.Auth, key, proxy)
if err != nil {
return err
}
ctx["auth"] = auth
ctx["authHeader"] = auth["authHeader"]
if a.plugin.Meta.Auth.Type == "" || a.plugin.Meta.Auth.Type == "none" || a.plugin.Meta.Auth.Type == "api_key" {
ctx["apiKey"] = key
}
return nil
}
func hookHTTPResponse(resp *http.Response) map[string]any {
headers := map[string]string{}
status := 0
@@ -1176,27 +1305,62 @@ func validateTaskArtifacts(value any) ([]channel.TaskArtifact, error) {
return artifacts, nil
}
// requestDescriptor runs a request-building hook. A plugin that preserves JSON
// order has its descriptor decoded from the JSON.stringify text of the result,
// with the body kept as that text (BodyText) for the upstream.
func (a *TaskAdaptor) requestDescriptor(ctx context.Context, hook string, args ...any) (requestDescriptor, error) {
if !a.plugin.Meta.PreservesJSONOrder() {
var descriptor requestDescriptor
err := a.plugin.Engine.CallInto(ctx, &descriptor, hook, args...)
return descriptor, err
}
var ordered orderedRequestDescriptor
if err := a.plugin.Engine.CallJSONInto(ctx, &ordered, hook, args...); err != nil {
return requestDescriptor{}, err
}
descriptor := requestDescriptor{
ResponseType: ordered.ResponseType, URL: ordered.URL, Method: ordered.Method, Headers: ordered.Headers,
Credentialless: ordered.Credentialless, Action: ordered.Action, Model: ordered.Model,
RewriteModel: ordered.RewriteModel, BodyType: ordered.BodyType, Parts: ordered.Parts,
}
switch {
case len(ordered.Body) == 0 || string(ordered.Body) == "null":
case ordered.Body[0] == '"':
// A string body is the plugin's own text, sent as it is.
var text string
if err := common.Unmarshal(ordered.Body, &text); err != nil {
return requestDescriptor{}, err
}
descriptor.Body = text
default:
descriptor.Body, descriptor.BodyText = ordered.Body, ordered.Body
}
return descriptor, nil
}
func (a *TaskAdaptor) buildSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*requestDescriptor, error) {
if a.submit != nil {
return a.submit, nil
}
started := time.Now()
value, err := a.plugin.Engine.Call(c.Request.Context(), "buildSubmitRequest", a.submitContext(c, info))
if err != nil {
logger.LogDebug(
c,
"task_plugin subsystem=adaptor event=build_submit_failed plugin=%q stage=build_submit_request reason=hook_failed elapsed_ms=%d",
a.plugin.Meta.Key,
time.Since(started).Milliseconds(),
)
return nil, err
// buildSubmitRequest runs before quota is reserved, so a failed fetch
// costs nothing; requests that start vendor work stay the descriptor.
fetchContext := c.Request.Context()
if info.HasChannelMeta() {
fetchContext = pluginruntime.WithFetcher(fetchContext, a.hookFetcher(info.ChannelBaseUrl, info.ChannelSetting.Proxy))
}
var descriptor requestDescriptor
if err = convert(value, &descriptor); err != nil {
descriptor, err := a.requestDescriptor(fetchContext, "buildSubmitRequest", a.submitContext(c, info))
if err != nil {
reason := "hook_failed"
var invalid *pluginruntime.ResultError
if errors.As(err, &invalid) {
reason = "invalid_descriptor"
}
logger.LogDebug(
c,
"task_plugin subsystem=adaptor event=build_submit_failed plugin=%q stage=build_submit_request reason=invalid_descriptor elapsed_ms=%d",
"task_plugin subsystem=adaptor event=build_submit_failed plugin=%q stage=build_submit_request reason=%s elapsed_ms=%d",
a.plugin.Meta.Key,
reason,
time.Since(started).Milliseconds(),
)
return nil, err
@@ -1282,7 +1446,11 @@ func (a *TaskAdaptor) submitContext(c *gin.Context, info *relaycommon.RelayInfo)
}
}
if taskRequest, exists := c.Get("task_request"); exists {
routeRequest.RequestBody = jsonValue(taskRequest)
routeRequest.RequestBody = taskRequest
}
a.requestBodyText = nil
if text, exists := c.Get(pluginruntime.ContextKeyRequestBodyText); exists && a.plugin.Meta.PreservesJSONOrder() {
a.requestBodyText, _ = text.(json.RawMessage)
}
if c.Request != nil {
if routeRequest.Path == "" {
@@ -1309,8 +1477,8 @@ func (a *TaskAdaptor) submitContext(c *gin.Context, info *relaycommon.RelayInfo)
if form, err := common.ParseMultipartFormReusable(c); err == nil {
defer form.RemoveAll()
for field, headers := range form.File {
for _, header := range headers {
files = append(files, map[string]any{"ref": "request_file:" + field, "field": field, "filename": header.Filename, "mimeType": header.Header.Get("Content-Type"), "size": header.Size})
for index, header := range headers {
files = append(files, map[string]any{"ref": pluginruntime.FileReference(field, index), "field": field, "filename": header.Filename, "mimeType": header.Header.Get("Content-Type"), "size": header.Size})
}
}
}
@@ -1323,7 +1491,25 @@ func (a *TaskAdaptor) submitContext(c *gin.Context, info *relaycommon.RelayInfo)
a.files = append(a.files[:0], files...)
}
ctx := routeRequest.JSValue()
ctx["requestBody"] = jsonValue(routeRequest.RequestBody)
// Hooks never write into the request body, so the same object normalizes
// to the same value on every hook call of this request.
requestBody := routeRequest.RequestBody
source, isObject := requestBody.(map[string]any)
if a.requestBodyText != nil {
// Driver hooks read the request body, so the engine parses the text
// in place for each of them, members in order.
requestBody = pluginruntime.RawJSON(a.requestBodyText)
} else if isObject && a.hookRequestSource != nil && reflect.ValueOf(source).UnsafePointer() == reflect.ValueOf(a.hookRequestSource).UnsafePointer() {
requestBody = a.hookRequestBody
} else {
if !isPlainJSONValue(requestBody, 0) {
requestBody = jsonValue(requestBody)
}
if isObject {
a.hookRequestSource, a.hookRequestBody = source, requestBody
}
}
ctx["requestBody"] = requestBody
ctx["requestHeaders"] = requestHeaders
ctx["files"] = files
ctx["action"] = info.Action
@@ -1352,21 +1538,13 @@ func (a *TaskAdaptor) submitContext(c *gin.Context, info *relaycommon.RelayInfo)
ctx["upstreamModel"] = info.UpstreamModelName
ctx["baseUrl"] = info.ChannelBaseUrl
ctx["userSetting"] = info.UserSetting
proxy := ""
proxy = info.ChannelSetting.Proxy
if auth, err := resolveAuth(a.plugin.Meta.Auth, info.ApiKey, proxy); err == nil {
ctx["auth"] = auth
ctx["authHeader"] = auth["authHeader"]
if a.plugin.Meta.Auth.Type == "" || a.plugin.Meta.Auth.Type == "none" || a.plugin.Meta.Auth.Type == "api_key" {
ctx["apiKey"] = info.ApiKey
}
} else {
if err := a.applyUpstreamCredentials(ctx, info.ChannelType, info.ApiKey, info.ChannelSetting.Proxy); err != nil {
ctx["authError"] = err.Error()
}
return ctx
}
func (a *TaskAdaptor) usageRatios(ctx context.Context, modelName, hook string, args ...any) (map[string]float64, error) {
func (a *TaskAdaptor) usageRatios(ctx context.Context, models []string, hook string, args ...any) (map[string]float64, error) {
if !a.hasHook(ctx, hook) {
return nil, nil
}
@@ -1384,7 +1562,7 @@ func (a *TaskAdaptor) usageRatios(ctx context.Context, modelName, hook string, a
logger.LogDebug(ctx, "task_plugin subsystem=adaptor event=usage_hook_failed plugin=%q hook=%q reason=result_not_object elapsed_ms=%d", a.plugin.Meta.Key, hook, time.Since(started).Milliseconds())
return nil, fmt.Errorf("plugin usage hook must return an object")
}
ratios, err := a.validatedUsageRatios(facts, modelName)
_, ratios, err := a.validatedUsageRatios(facts, models...)
if err != nil {
logger.LogDebug(ctx, "task_plugin subsystem=adaptor event=usage_hook_failed plugin=%q hook=%q reason=invalid_usage elapsed_ms=%d", a.plugin.Meta.Key, hook, time.Since(started).Milliseconds())
return nil, err
@@ -1401,8 +1579,8 @@ func (a *TaskAdaptor) usageRatios(ctx context.Context, modelName, hook string, a
return ratios, nil
}
func (a *TaskAdaptor) validateResolvedUsageRequest(request any, modelName string) error {
schema, _ := a.plugin.Meta.UsageForModel(modelName)
func (a *TaskAdaptor) validateResolvedUsageRequest(request any, models ...string) error {
schema, _ := a.plugin.Meta.UsageForModels(models...)
return a.validateResolvedUsageValue(jsonValue(request), schema)
}
@@ -1410,11 +1588,21 @@ func (a *TaskAdaptor) validateResolvedUsageValue(value any, usageSchema map[stri
switch typed := value.(type) {
case map[string]any:
for key, item := range typed {
limit, canonical := canonicalUsageLimit(key)
// A duration of -1 asks the vendor to pick the length. Plugins that
// declare duration-auto@1 reserve the longest duration for it and
// settle on the usage the vendor reports; usage facts stay
// non-negative.
if canonical && limit == relaycommon.MaxTaskDurationSeconds && a.plugin.Meta.AcceptsAutoDuration() {
if number, _ := usageNumber(item, true); number == -1 {
continue
}
}
if schema, declared := usageSchema[key]; declared {
if _, err := validateUsageValue(item, schema, true); err != nil {
return err
}
} else if limit, canonical := canonicalUsageLimit(key); canonical {
} else if canonical {
if err := validateUsageLimit(item, limit, true); err != nil {
return err
}
@@ -1433,17 +1621,22 @@ func (a *TaskAdaptor) validateResolvedUsageValue(value any, usageSchema map[stri
return nil
}
func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any, modelName string) (map[string]float64, error) {
usageSchema, _ := a.plugin.Meta.UsageForModel(modelName)
// validatedUsageRatios returns the normalized facts in a new map because the
// hook result may be a host value, such as the request body, that must stay
// unchanged.
func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any, models ...string) (map[string]any, map[string]float64, error) {
usageSchema, _ := a.plugin.Meta.UsageForModels(models...)
validated := make(map[string]any, len(facts))
ratios := make(map[string]float64)
for key, value := range facts {
validated[key] = value
if schema, declared := usageSchema[key]; declared {
number, err := validateUsageValue(value, schema, false)
if err != nil {
return nil, err
return nil, nil, err
}
if schema.Type == "number" {
facts[key] = number
validated[key] = number
if number > 0 {
ratios[key] = number
}
@@ -1462,20 +1655,20 @@ func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any, modelName strin
limit = relaycommon.MaxTaskDurationSeconds
}
if err := validateUsageNumberLimit(number, limit); err != nil {
return nil, err
return nil, nil, err
}
// Keep numeric facts usable by existing expressions even when the
// selected profile no longer declares that legacy extension field.
facts[key] = number
validated[key] = number
if number > 0 {
ratios[key] = number
}
}
return ratios, nil
return validated, ratios, nil
}
func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any, modelName string) (map[string]any, error) {
usageSchema, _ := a.plugin.Meta.UsageForModel(modelName)
func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any, models ...string) (map[string]any, error) {
usageSchema, _ := a.plugin.Meta.UsageForModels(models...)
if facts == nil {
return nil, nil
}
@@ -1602,8 +1795,10 @@ func usageNumber(value any, allowNumericString bool) (float64, bool) {
}
}
var usageKeySeparators = strings.NewReplacer("_", "", "-", "")
func canonicalUsageLimit(key string) (int, bool) {
normalized := strings.NewReplacer("_", "", "-", "").Replace(strings.ToLower(key))
normalized := usageKeySeparators.Replace(strings.ToLower(key))
switch normalized {
case "duration", "durationseconds", "second", "seconds":
return relaycommon.MaxTaskDurationSeconds, true
@@ -1615,13 +1810,13 @@ func canonicalUsageLimit(key string) (int, bool) {
}
func (a *TaskAdaptor) logRejectedUsage(hook string, _ error) {
common.SysError(fmt.Sprintf("task plugin %s rejected invalid %s billing facts", a.plugin.Meta.Key, hook))
common.SysError(common.LogText("task plugin %s rejected invalid %s billing facts", a.plugin.Meta.Key, hook))
}
func (a *TaskAdaptor) hasHook(ctx context.Context, hook string) bool {
has, err := a.plugin.Engine.HasExport(ctx, hook)
return err == nil && has
func (a *TaskAdaptor) hasHook(_ context.Context, hook string) bool {
return a.plugin.Engine.HasExport(hook)
}
func convert(value any, target any) error {
data, err := common.Marshal(value)
if err != nil {
@@ -1645,8 +1840,45 @@ func jsonValue(value any) any {
return normalized
}
// Already-parsed JSON and Sobek's plain exported values do not need another
// encode/decode pass. Copy containers to isolate hooks, preserving the codec's
// isPlainJSONValue reports whether hooks can take value as is: it reaches
// JavaScript exactly as the codec's decoding of it would, and the engine never
// writes back into it. Integers count, since hooks see the same numbers; a nil
// slice does not, since it would reach hooks as [] instead of null.
func isPlainJSONValue(value any, depth int) bool {
if depth > 64 {
return false
}
switch typed := value.(type) {
case nil, bool, int64, int:
return true
case string:
return utf8.ValidString(typed)
case float64:
return !math.IsNaN(typed) && !math.IsInf(typed, 0)
case map[string]any:
for key, item := range typed {
if !utf8.ValidString(key) || !isPlainJSONValue(item, depth+1) {
return false
}
}
return true
case []any:
if typed == nil {
return false
}
for _, item := range typed {
if !isPlainJSONValue(item, depth+1) {
return false
}
}
return true
default:
return false
}
}
// Already-parsed JSON and the engine's plain exported values do not need another
// encode/decode pass. Copy containers for host-side edits, preserving the codec's
// float64 numbers and nulls. Other Go types still use the configured codec.
func cloneJSONValue(value any, depth int) (any, bool) {
if depth > 64 {
relay/channel/task/jsplugin/adaptor_test.go
+596 -36
View File
@@ -5,8 +5,10 @@ import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"io"
"math"
"mime"
"mime/multipart"
"net/http"
"net/http/httptest"
@@ -82,9 +84,10 @@ func TestTaskAdaptorRejectsDeprecatedClientResponse(t *testing.T) {
}
func TestTaskAdaptorBuildsMultipartFromOpaqueFileReference(t *testing.T) {
// The plugin forwards the client's Content-Type, whose boundary is not the one of the rebuilt body.
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 buildSubmitRequest(ctx) { return {url:ctx.baseUrl+"/submit",headers:{"Content-Type":ctx.requestHeaders["Content-Type"]},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{})
@@ -107,7 +110,13 @@ export function parseSubmitResponse(ctx,r){return {taskId:"1"}} export function
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="))
upstreamRequest := httptest.NewRequest(http.MethodPost, "https://provider.example/submit", bytes.NewReader(requestBytes))
require.NoError(t, adaptor.BuildRequestHeader(c, upstreamRequest, info))
mediaType, params, err := mime.ParseMediaType(upstreamRequest.Header.Get("Content-Type"))
require.NoError(t, err)
assert.Equal(t, "multipart/form-data", mediaType)
require.NotEmpty(t, params["boundary"])
reader := multipart.NewReader(bytes.NewReader(requestBytes), params["boundary"])
form, err := reader.ReadForm(1024)
require.NoError(t, err)
assert.Equal(t, []string{"m"}, form.Value["model"])
@@ -129,7 +138,8 @@ export function buildSubmitRequest(ctx) {
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"}
dataUrl:{__fileRef:ctx.files[0].ref,encoding:"dataUrl"},
refs:ctx.files.map(function(file){return file.ref;})
}};
}
export function parseSubmitResponse(){return {taskId:"1"}} export function buildQueryRequest(){return {url:"https://example.com"}} export function parseTaskResult(){return {status:"SUCCESS"}}
@@ -154,6 +164,57 @@ export function parseSubmitResponse(){return {taskId:"1"}} export function build
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"])
assert.Equal(t, []any{"request_file:input_reference"}, decoded["refs"])
t.Run("repeated field files are addressed by index", func(t *testing.T) {
second := base64.StdEncoding.EncodeToString([]byte("second-bytes"))
source := `
export const meta = {apiVersion:1,key:"json-inline-repeat",name:"JSON Inline Repeat",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task"};
export function buildSubmitRequest(ctx) {
return {url:ctx.baseUrl+"/submit",body:{
refs:ctx.files.map(function(file){return file.ref;}),
images:ctx.files.map(function(file){return {__fileRef:file.ref,encoding:"base64"};}),
missing:{__fileRef:"request_file:image[]#2",encoding:"base64"}
}};
}
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)
var input bytes.Buffer
writer := multipart.NewWriter(&input)
for _, content := range []string{fileBytes, "second-bytes"} {
file, createErr := writer.CreateFormFile("image[]", "ref.png")
require.NoError(t, createErr)
_, err = file.Write([]byte(content))
require.NoError(t, err)
}
require.NoError(t, writer.Close())
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(input.Bytes()))
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
c.Set("task_request", map[string]any{"prompt": "p"})
_, err = adaptor.BuildRequestBody(c, info)
require.ErrorContains(t, err, `unknown file reference "request_file:image[]#2"`, "an index beyond the uploaded files is rejected")
source = strings.Replace(source, `,
missing:{__fileRef:"request_file:image[]#2",encoding:"base64"}`, "", 1)
plugin, err = pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
require.NoError(t, err)
adaptor = New(plugin)
adaptor.Init(info)
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, []any{"request_file:image[]", "request_file:image[]#1"}, decoded["refs"])
assert.Equal(t, []any{encoded, second}, decoded["images"])
})
}
func TestTaskAdaptorJSONFilePlaceholderErrors(t *testing.T) {
@@ -351,6 +412,95 @@ export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit"}} expo
assert.Contains(t, taskErr.Message, "must not return renderer")
}
func TestTaskAdaptorMasksChannelAddressInBuildFailure(t *testing.T) {
source := `
export const meta = {apiVersion:1,key:"build-fail",name:"Build Fail",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task",allowedHosts:["assets.corp.internal"]};
export function buildSubmitRequest(ctx){throw new Error("utils.fetch to private.example.com failed: dial tcp: lookup private.example.com: no such host; utils.fetch to assets.corp.internal failed: dial tcp: lookup assets.corp.internal: no such host")}
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)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/video/generations", nil)
c.Set("task_request", map[string]any{"model": "m"})
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, "https://private.example.com")
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://private.example.com"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, OriginModelName: "m"}
adaptor := New(plugin)
adaptor.Init(info)
taskErr := adaptor.ValidateRequestAndSetAction(c, info)
require.NotNil(t, taskErr)
assert.Equal(t, "plugin_request_invalid", taskErr.Code)
assert.True(t, taskErr.LocalError)
assert.Contains(t, taskErr.Message, "no such host")
assert.NotContains(t, taskErr.Message, "private.example.com")
assert.NotContains(t, taskErr.Message, "assets.corp.internal")
}
const orderedProtocolPlugin = `
export const meta = {apiVersion:1,key:"ordered",name:"Ordered",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task",requiredCapabilities:["json-order@1"],protocols:[{name:"openai_responses",supports:["sync"]}]};
export const protocols = {openai_responses:{decodeRequest:function(ctx){const b = ctx.body.value; return {kind:"submit",model:ctx.model,requestBody:{model:b.model,state:b.state,questions:b.questions}};},renderFinal:function(){return {};}}};
export function buildSubmitRequest(ctx){const b = ctx.requestBody; return {url:ctx.baseUrl+"/submit",body:{model:b.model,state:b.state,questions:b.questions}}}
export function parseSubmitResponse(){return {taskId:"one"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}}
`
// A plugin declaring json-order@1 forwards the members in the order the client
// sent them, through the final decode, the request body hooks receive and the
// upstream body, while billing bounds still read the decoded value; other
// plugins keep the host codec's sorted encoding.
func TestTaskAdaptorJSONOrderCapability(t *testing.T) {
client := `{"model":"m","state":{"zeta":"1","alpha":"2","mid":"&<>"},"questions":{"q2":{"type":"noul"},"q1":{"type":"noul"}}}`
undeclared := strings.Replace(orderedProtocolPlugin, `requiredCapabilities:["json-order@1"],`, "", 1)
cases := []struct {
name string
source string
client string
want string
status int
}{
{name: "declared keeps the client's order", source: orderedProtocolPlugin, client: client,
want: `{"model":"m","state":{"zeta":"1","alpha":"2","mid":"&<>"},"questions":{"q2":{"type":"noul"},"q1":{"type":"noul"}}}`},
{name: "undeclared stays sorted", source: undeclared, client: client,
want: `{"model":"m","questions":{"q1":{"type":"noul"},"q2":{"type":"noul"}},"state":{"alpha":"2","mid":"\u0026\u003c\u003e","zeta":"1"}}`},
{name: "declared still bounds billing quantities", source: orderedProtocolPlugin,
client: `{"model":"m","state":{"duration":3601},"questions":{}}`, status: http.StatusBadRequest},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
plugin, err := pluginruntime.NewRegistry().Register(tc.source, pluginruntime.Options{})
require.NoError(t, err)
var value any
require.NoError(t, common.Unmarshal([]byte(tc.client), &value))
protocolContext := pluginruntime.ProtocolRequestContext{
RouteRequestContext: pluginruntime.RouteRequestContext{Body: map[string]any{"kind": "json", "value": value}, BodyText: json.RawMessage(tc.client)},
Protocol: "openai_responses", Model: "m",
}
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: "m"})
c.Set(pluginruntime.ContextKeyProtocolRequest, protocolContext)
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}, OriginModelName: "m"}
adaptor := New(plugin)
adaptor.Init(info)
taskErr := adaptor.ValidateRequestAndSetAction(c, info)
if tc.status != 0 {
require.NotNil(t, taskErr)
assert.Equal(t, tc.status, taskErr.StatusCode)
assert.Equal(t, "plugin_usage_invalid", taskErr.Code)
return
}
require.Nil(t, taskErr)
body, err := adaptor.BuildRequestBody(c, info)
require.NoError(t, err)
raw, err := io.ReadAll(body)
require.NoError(t, err)
assert.Equal(t, tc.want, string(raw))
})
}
}
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"}]; }
@@ -376,7 +526,7 @@ export function buildContentRequest(ctx) {
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})
descriptor, err := adaptor.BuildContentRequest(t.Context(), 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)
@@ -396,7 +546,7 @@ export function parseTaskResult(){return {status:"SUCCESS"};}
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})
descriptor, err = fallback.BuildContentRequest(t.Context(), &model.Task{}, "video", channel.TaskArtifactClientRequest{Method: http.MethodGet})
require.NoError(t, err)
assert.Nil(t, descriptor)
}
@@ -437,6 +587,7 @@ export function buildContentRequest(ctx) { return {url:"https://cdn.example/vide
adaptor := New(plugin)
adaptor.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"}})
descriptor, err := adaptor.BuildContentRequest(
t.Context(),
&model.Task{TaskID: "task", Data: []byte(`{}`)},
"video",
channel.TaskArtifactClientRequest{Method: http.MethodGet},
@@ -447,6 +598,59 @@ export function buildContentRequest(ctx) { return {url:"https://cdn.example/vide
assert.Equal(t, "https://cdn.example/video.mp4", descriptor.URL)
}
func TestTaskAdaptorHookFetchReachesRequestHostsOnly(t *testing.T) {
service.InitHttpClient()
extra := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(`{"from":"extra"}`))
}))
defer extra.Close()
base := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/info":
assert.Equal(t, "Bearer key", r.Header.Get("Authorization"))
_, _ = w.Write([]byte(`{"from":"base"}`))
case "/moved":
http.Redirect(w, r, extra.URL+"/elsewhere", http.StatusFound)
default:
t.Errorf("unexpected request %s", r.URL.Path)
}
}))
defer base.Close()
source := strings.Replace(mockPlugin, `fetchMode: "per_task",`, `fetchMode: "per_task", allowedHosts: ["`+strings.TrimPrefix(extra.URL, "http://")+`"],`, 1)
source = strings.Replace(source, `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) {
const base = utils.fetch({url: ctx.baseUrl + "/info", headers: {Authorization: "Bearer " + ctx.apiKey}});
const extra = utils.fetch({url: "http://" + meta.allowedHosts[0] + "/info"});
const moved = utils.fetch({url: ctx.baseUrl + "/moved"});
let denied = "";
try { utils.fetch({url: "https://other.example/info"}); } catch (e) { denied = e.message; }
return {url: ctx.baseUrl + "/content", method: "GET", headers: {
"X-Base": base.body.from, "X-Extra": extra.body.from, "X-Moved": moved.status + " " + moved.headers.Location, "X-Denied": denied,
}};
}`, 1)
source = strings.Replace(source, `render: function(ctx, task) { return {id: task.task_id, status: "completed"}; }`,
`render: function(ctx, task) { return utils.fetch({url: "https://provider.example/info"}).body; }`, 1)
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
require.NoError(t, err)
adaptor := New(plugin)
adaptor.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: base.URL, ApiKey: "key"}})
descriptor, err := adaptor.BuildContentRequest(t.Context(), &model.Task{TaskID: "task", Data: []byte(`{}`)}, "video", channel.TaskArtifactClientRequest{Method: http.MethodGet})
require.NoError(t, err)
require.NotNil(t, descriptor)
assert.Equal(t, map[string]string{
"X-Base": "base",
"X-Extra": "extra",
"X-Moved": "302 " + extra.URL + "/elsewhere",
"X-Denied": `utils.fetch to other.example failed: plugin request host "other.example" is not allowed`,
}, descriptor.Headers)
// Renderers run on every client read and may not fetch.
_, err = adaptor.ConvertToOpenAIVideo(&model.Task{TaskID: "task_public", Status: model.TaskStatusSuccess})
assert.ErrorContains(t, err, "utils.fetch is not available in protocols.openai_video.render")
}
func TestTaskAdaptorMapsJSContract(t *testing.T) {
gin.SetMode(gin.TestMode)
service.InitHttpClient()
@@ -610,13 +814,22 @@ func TestTaskAdaptorPreservesSoraVideoResponseFields(t *testing.T) {
}
func TestTaskAdaptorRejectsNonObjectOpenAIVideoRendererOutput(t *testing.T) {
for _, value := range []string{"null", "[]", `"video"`, "42", "false"} {
for value, message := range map[string]string{
"null": "plugin returned an invalid OpenAI video object",
"undefined": "plugin returned an invalid OpenAI video object",
"[]": "plugin returned an invalid OpenAI video object",
`"video"`: "plugin returned an invalid OpenAI video object",
"42": "plugin returned an invalid OpenAI video object",
"false": "plugin returned an invalid OpenAI video object",
"{seconds: NaN}": "json: unsupported value: NaN",
"[Infinity]": "json: unsupported value: +Inf",
} {
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")
require.EqualError(t, err, message)
})
}
}
@@ -808,6 +1021,53 @@ export function extractUsageOnComplete(task, result, body) { return (body || {})
assert.Nil(t, ratios)
})
t.Run("duration -1 passes only for plugins that declare duration-auto@1", func(t *testing.T) {
declaring, err := pluginruntime.NewRegistry().Register(strings.Replace(source, `fetchMode: "per_task",`, `fetchMode: "per_task", requiredCapabilities: ["duration-auto@1"],`, 1), pluginruntime.Options{})
require.NoError(t, err)
for _, tc := range []struct {
name string
declaring bool
body map[string]any
accepted bool
}{
{"undeclared plugins keep rejecting it", false, map[string]any{"metadata": map[string]any{"duration": -1}}, false},
{"vendor parameters and declared usage fields", true, map[string]any{
"duration": -1,
"metadata": map[string]any{"seconds": "-1", "parameters": map[string]any{"durationSeconds": -1}},
}, true},
{"other negative durations", true, map[string]any{"metadata": map[string]any{"duration": -2}}, false},
{"durations above the host limit", true, map[string]any{"metadata": map[string]any{"duration": relaycommon.MaxTaskDurationSeconds + 1}}, false},
{"negative counts", true, map[string]any{"metadata": map[string]any{"count": -1}}, false},
} {
t.Run(tc.name, func(t *testing.T) {
adaptor, context, info := newRequest(t, tc.body)
if tc.declaring {
adaptor = New(declaring)
adaptor.Init(info)
}
taskErr := adaptor.ValidateRequestAndSetAction(context, info)
if tc.accepted {
assert.Nil(t, taskErr)
return
}
require.NotNil(t, taskErr)
assert.Equal(t, "plugin_usage_invalid", taskErr.Code)
})
}
adaptor, context, info := newRequest(t, map[string]any{
"metadata": map[string]any{"duration": -1},
"hookUsageEntries": []any{map[string]any{"name": "duration", "value": -1}},
})
adaptor = New(declaring)
adaptor.Init(info)
require.Nil(t, adaptor.ValidateRequestAndSetAction(context, info))
_, err = adaptor.ExtractUsageFactsValidated(context, info)
assert.Error(t, err, "usage facts stay non-negative")
_, err = adaptor.EstimateBillingValidated(context, info)
assert.Error(t, err, "billing ratios stay non-negative")
})
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{
@@ -1100,6 +1360,26 @@ func TestSubmitContextOriginTasksNilDataOnInvalidJSON(t *testing.T) {
assert.Nil(t, originTasks[0]["data"])
}
func TestSubmitContextNormalizesRequestBodyOnEveryCall(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"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
c.Set("task_request", map[string]any{"prompt": string([]byte{0xff}), "count": int64(math.MaxInt64)})
normalized := map[string]any{"prompt": "�", "count": float64(math.MaxInt64)}
assert.Equal(t, normalized, adaptor.submitContext(c, info)["requestBody"])
assert.Equal(t, normalized, adaptor.submitContext(c, info)["requestBody"])
assert.Equal(t, normalized, adaptor.submitContext(nil, info)["requestBody"])
replaced := map[string]any{"prompt": "second", "count": int64(2)}
c.Set("task_request", replaced)
assert.Equal(t, replaced, adaptor.submitContext(c, info)["requestBody"])
}
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"})
@@ -1464,10 +1744,10 @@ 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) {
newRequest := func(t *testing.T, origin, upstream string, body map[string]any) (*TaskAdaptor, *gin.Context, *relaycommon.RelayInfo) {
t.Helper()
info := &relaycommon.RelayInfo{
OriginModelName: "public-alias",
OriginModelName: origin,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: upstream, ChannelBaseUrl: "https://provider.example",
},
@@ -1482,18 +1762,20 @@ export function extractUsageOnComplete(ctx,result,body){return body.usage;}
}
for _, tc := range []struct {
name, upstream, rewrite, mode string
units float64
wantError bool
name, origin, 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},
{"mapped image", "public-alias", "image", "", "image", 2, false},
{"mapped video token units", "public-alias", "video", "", "video", 500000, false},
{"rewritten model", "public-alias", "image", "video", "video", 500000, false},
{"image count ceiling", "public-alias", "image", "", "image", float64(dto.MaxImageN + 1), true},
{"profile enum", "public-alias", "image", "", "video", 2, true},
{"endpoint upstream keeps the declared origin profile", "image", "ep-endpoint", "", "image", 2, false},
{"endpoint upstream enforces the origin count ceiling", "image", "ep-endpoint", "", "image", float64(dto.MaxImageN + 1), true},
} {
t.Run(tc.name, func(t *testing.T) {
adaptor, c, info := newRequest(t, tc.upstream, map[string]any{
adaptor, c, info := newRequest(t, tc.origin, tc.upstream, map[string]any{
"hookUnits": tc.units, "hookMode": tc.mode, "rewriteTo": tc.rewrite,
})
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
@@ -1514,7 +1796,7 @@ export function extractUsageOnComplete(ctx,result,body){return body.usage;}
}
t.Run("request revalidated after model rewrite", func(t *testing.T) {
adaptor, c, info := newRequest(t, "video", map[string]any{
adaptor, c, info := newRequest(t, "public-alias", "video", map[string]any{
"rewriteTo": "image", "metadata": map[string]any{"units": dto.MaxImageN + 1},
})
taskErr := adaptor.ValidateRequestAndSetAction(c, info)
@@ -1523,7 +1805,7 @@ export function extractUsageOnComplete(ctx,result,body){return body.usage;}
})
t.Run("request enum belongs to the rewritten model", func(t *testing.T) {
adaptor, c, info := newRequest(t, "video", map[string]any{
adaptor, c, info := newRequest(t, "public-alias", "video", map[string]any{
"rewriteTo": "image", "mode": "image", "units": 2,
})
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
@@ -1531,7 +1813,7 @@ export function extractUsageOnComplete(ctx,result,body){return body.usage;}
})
t.Run("completion selects each persisted model including legacy origin fallback", func(t *testing.T) {
adaptor, _, _ := newRequest(t, "image", map[string]any{})
adaptor, _, _ := newRequest(t, "public-alias", "image", map[string]any{})
for _, tc := range []struct {
upstream, origin, mode string
units float64
@@ -1542,6 +1824,8 @@ export function extractUsageOnComplete(ctx,result,body){return body.usage;}
{"", "video", "video", 500000, true},
{"image", "public-alias", "image", float64(dto.MaxImageN + 1), false},
{"video", "public-alias", "image", 1, false},
{"ep-endpoint", "image", "image", 2, true},
{"ep-endpoint", "image", "image", float64(dto.MaxImageN + 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}})
@@ -1561,7 +1845,7 @@ export function extractUsageOnComplete(ctx,result,body){return body.usage;}
`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{})
_, _, info := newRequest(t, "public-alias", "image", map[string]any{})
adaptor := New(updated)
adaptor.Init(info)
result, err := adaptor.ParseTaskResult(&model.Task{Properties: model.Properties{UpstreamModelName: "image"}},
@@ -1574,20 +1858,23 @@ export function extractUsageOnComplete(ctx,result,body){return body.usage;}
})
t.Run("mixed batch selects schema by task rather than adaptor", func(t *testing.T) {
adaptor, _, _ := newRequest(t, "image", map[string]any{})
adaptor, _, _ := newRequest(t, "public-alias", "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"}},
{TaskID: "endpoint-image", Properties: model.Properties{OriginModelName: "image", UpstreamModelName: "ep-endpoint"}},
}
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"}}}
{"taskId":"invalid-image","status":"SUCCESS","data":{"usage":{"units":129,"mode":"image"}}},
{"taskId":"endpoint-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)
require.Len(t, results, 4)
assert.Nil(t, results["endpoint-image"].TaskInfo.UsageFacts, "endpoint upstream must keep the origin profile's count ceiling")
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)
@@ -1698,22 +1985,27 @@ func TestTaskSubmitStreamIdleTimeout(t *testing.T) {
}
func TestPluginJSONValuesPreserveCodecNormalizationAndIsolation(t *testing.T) {
for _, value := range []any{
map[string]any{"units": int64(3), "nested": []any{true, "<image> / 图像", 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 {
for _, tc := range []struct {
value any
plain bool
}{
{map[string]any{"units": int64(3), "nested": []any{true, "<image> / 图像", math.Copysign(0, -1), nil}}, true},
{map[string]any{"empty": []any{}, "object": map[string]any(nil)}, true},
{map[string]any{"null": []any(nil)}, false},
{map[string]any{"large": int64(math.MaxInt64), "invalid UTF-8": string([]byte{0xff, 0xfe})}, false},
{map[string]any{string([]byte{0xff}): "invalid key"}, false},
{map[string]any{"bytes": []byte{0, 1, 255}, "number": json.Number("9007199254740993")}, false},
{json.RawMessage(`{"units":2,"enabled":false}`), false},
{struct {
Units int `json:"units"`
}{Units: 0},
}{Units: 0}, false},
} {
encoded, err := common.Marshal(value)
encoded, err := common.Marshal(tc.value)
require.NoError(t, err)
var expected any
require.NoError(t, common.Unmarshal(encoded, &expected))
assert.Equal(t, expected, jsonValue(value))
assert.Equal(t, expected, jsonValue(tc.value))
assert.Equal(t, tc.plain, isPlainJSONValue(tc.value, 0))
}
source := map[string]any{"items": []any{map[string]any{"label": "original"}}}
copy := jsonValue(source).(map[string]any)
@@ -1721,6 +2013,86 @@ func TestPluginJSONValuesPreserveCodecNormalizationAndIsolation(t *testing.T) {
assert.Equal(t, "original", source["items"].([]any)[0].(map[string]any)["label"])
}
func TestRequestDescriptorDecodingMatchesCodec(t *testing.T) {
engine, err := pluginruntime.Compile(`
class Descriptor { constructor() { this.url = "u"; this.method = "PUT"; } }
let deep = "leaf";
for (let i = 0; i < 70; i++) deep = [deep];
const results = {
"plain body": () => ({url: "https://provider.example/submit", method: "POST", headers: {"X-Plugin": "submit"}, body: {
units: 3, large: 2 ** 60, zero: -0, text: "<image> & 图像", missing: undefined,
items: [null, true, 1.5, {}, [], undefined],
}}),
"argument body": (arg) => ({url: "u", body: arg}),
"argument headers": (arg) => ({url: "u", headers: arg.headers}),
"changed argument body": (arg) => { arg.units = 4; return {url: "u", body: arg}; },
"string body": () => ({url: "u", body: '{"raw":true}'}),
"null body": () => ({url: "u", body: null}),
"undefined body": () => ({url: "u", body: undefined}),
"no body": () => ({url: "u", credentialless: true}),
"class instance": () => new Descriptor(),
"null prototype": () => Object.assign(Object.create(null), {url: "u", body: {a: 1}}),
"getter body": () => ({url: "u", get body() { return {from: "getter"}; }}),
"date body": () => ({url: "u", body: {at: new Date(0)}}),
"typed array body": () => ({url: "u", body: new Uint8Array([0, 1, 255])}),
"deep body": () => ({url: "u", body: deep}),
"capitalized body key": () => ({url: "u", Body: {from: "Body"}}),
"body and capitalized body": () => ({url: "u", body: {from: "body"}, Body: {from: "Body"}, BODY: "upper"}),
"case-insensitive other field": () => ({URL: "u", Method: "PUT", body: "text"}),
"NaN body": () => ({url: "u", body: {ratio: NaN}}),
"infinite method": () => ({url: "u", method: Infinity, body: {}}),
"cyclic body": () => { const body = {}; body.self = body; return {url: "u", body}; },
"invalid header": () => ({url: "u", headers: {x: 1}, body: {}}),
"invalid parts": () => ({url: "u", parts: "none", body: {}}),
"not an object": () => "descriptor",
"throwing getter": () => ({url: "u", get body() { throw new Error("getter failed"); }}),
};
export function build(name, arg) { return results[name](arg); }
`, pluginruntime.Options{})
require.NoError(t, err)
argument := func() map[string]any {
return map[string]any{
"units": int64(3), "large": int64(math.MaxInt64), "zero": math.Copysign(0, -1),
"headers": map[string]any{"X-Plugin": "submit"},
"items": []any{nil, true, 1.5, map[string]any{"label": "original"}, []any{}},
}
}
for _, name := range []string{
"plain body", "argument body", "argument headers", "changed argument body", "string body", "null body",
"undefined body", "no body", "class instance", "null prototype", "getter body", "date body",
"typed array body", "deep body", "capitalized body key", "body and capitalized body",
"case-insensitive other field", "NaN body", "infinite method", "cyclic body", "invalid header",
"invalid parts", "not an object", "throwing getter",
} {
var expected requestDescriptor
value, expectedErr := engine.Call(t.Context(), "build", name, argument())
hookFailed := expectedErr != nil
if !hookFailed {
expectedErr = convert(value, &expected)
}
var decoded requestDescriptor
err := engine.CallInto(t.Context(), &decoded, "build", name, argument())
if expectedErr != nil {
require.EqualError(t, err, expectedErr.Error(), name)
var invalid *pluginruntime.ResultError
assert.Equal(t, !hookFailed, errors.As(err, &invalid), name)
continue
}
require.NoError(t, err, name)
assert.Equal(t, expected, decoded, name)
expectedJSON, err := common.Marshal(expected)
require.NoError(t, err, name)
decodedJSON, err := common.Marshal(decoded)
require.NoError(t, err, name)
assert.Equal(t, string(expectedJSON), string(decodedJSON), name)
}
source := argument()
var decoded requestDescriptor
require.NoError(t, engine.CallInto(t.Context(), &decoded, "build", "argument body", source))
decoded.Body.(map[string]any)["items"].([]any)[3].(map[string]any)["label"] = "changed"
assert.Equal(t, "original", source["items"].([]any)[3].(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"}}};
@@ -1752,6 +2124,39 @@ export function parseTaskResult(){return {status:"SUCCESS"};}
assert.Equal(t, "original", request["nested"].(map[string]any)["label"])
}
func TestTaskUsageValidationLeavesReturnedHostRequestUnchanged(t *testing.T) {
// Returning ctx.requestBody hands back the host request map itself, so
// normalized facts must be built in a new map.
source := `
export const meta={apiVersion:1,key:"request-usage",name:"Request usage",version:"1.0.0",author:{name:"Test"},models:["usage"],fetchMode:"per_task",usageSchema:{units:{type:"number",unit:"count"}}};
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit",body:ctx.requestBody};}
export function extractUsage(ctx){return ctx.requestBody;}
export function parseSubmitResponse(){return {taskId:"usage"};}
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)}
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/submit", nil)
c.Set("task_request", request)
info := &relaycommon.RelayInfo{OriginModelName: "usage", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "usage", 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)
assert.Equal(t, int64(2), request["units"])
estimateRequest := map[string]any{"units": int64(2)}
c.Set("task_request", estimateRequest)
ratios, err := adaptor.EstimateBillingValidated(c, info)
require.NoError(t, err)
assert.Equal(t, map[string]float64{"units": 2}, ratios)
assert.Equal(t, int64(2), estimateRequest["units"])
}
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) {
@@ -1808,6 +2213,26 @@ export function parseSubmitEventDelta(ctx,event,previous) {
assert.Equal(t, map[string]any{"units": float64(0)}, parsed.Immediate.UsageFacts)
})
}
t.Run("delta hook exported without the declaration", func(t *testing.T) {
undeclared := strings.Replace(source, `requiredCapabilities:["submit-sse-delta@1"],`, "", 1)
undeclared = strings.Replace(undeclared, "export function parseSubmitEvent(", "function unusedSnapshotHook(", 1)
plugin, err := pluginruntime.CompilePlugin(undeclared, pluginruntime.Options{})
require.NoError(t, err)
require.Empty(t, plugin.Meta.RequiredCapabilities)
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))
stream := "data: " + first + "\n\ndata: " + last + "\n\n"
response := &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(stream))}
parsed, taskErr := adaptor.ParseResponse(c, response, info)
require.Nil(t, taskErr)
assert.JSONEq(t, `{"document":"helloworld","units":0}`, string(parsed.TaskData))
})
}
func TestAlibabaSubmitDeltaDoesNotMutateControlState(t *testing.T) {
@@ -1841,3 +2266,138 @@ func TestAlibabaSubmitDeltaDoesNotMutateControlState(t *testing.T) {
require.NoError(t, err)
assert.JSONEq(t, `{"choices":[{"count":2,"lastText":false,"finishReason":"stop"}],"hasUsage":true}`, string(encoded))
}
// newAPIGateway emulates the upstream New API instance a type-60 channel
// points at: it answers only on the plugin's prefixed native routes, requires
// the gateway token as a Bearer header, and renders what its own presenters
// render (the public task id already substituted into the id fields).
func newAPIGateway(t *testing.T, routes map[string]string, seen *[]string) *httptest.Server {
t.Helper()
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
*seen = append(*seen, r.Method+" "+r.URL.Path)
if r.Header.Get("Authorization") != "Bearer sk-gateway" {
w.WriteHeader(http.StatusUnauthorized)
return
}
body, ok := routes[r.Method+" "+r.URL.Path]
if !ok {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(body))
}))
}
func TestTaskAdaptorChainsDoubaoThroughNewAPIUpstream(t *testing.T) {
gin.SetMode(gin.TestMode)
service.InitHttpClient()
var seen []string
gateway := newAPIGateway(t, map[string]string{
"POST /doubao/api/v3/contents/generations/tasks": `{"id":"task_up_public"}`,
"GET /doubao/api/v3/contents/generations/tasks/task_up_public": `{"id":"task_up_public","status":"succeeded","content":{"video_url":"https://cdn.example/v.mp4"},"usage":{"completion_tokens":1200,"total_tokens":1200}}`,
}, &seen)
defer gateway.Close()
source, err := plugins.Source("doubao")
require.NoError(t, err)
plugin, err := pluginruntime.NewRegistry().RegisterFactory(source, pluginruntime.Options{Key: "doubao"})
require.NoError(t, err)
adaptor := New(plugin)
const modelName = "doubao-seedance-1-0-pro-250528"
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeNewAPI, ChannelBaseUrl: gateway.URL, ApiKey: "sk-gateway", UpstreamModelName: modelName},
OriginModelName: modelName,
TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "task_local"},
}
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
c.Set("task_request", relaycommon.TaskSubmitReq{Prompt: "a cat", Model: modelName})
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
requestBody, err := adaptor.BuildRequestBody(c, info)
require.NoError(t, err)
submitURL, err := adaptor.BuildRequestURL(info)
require.NoError(t, err)
assert.Equal(t, gateway.URL+"/doubao/api/v3/contents/generations/tasks", submitURL)
resp, err := adaptor.DoRequest(c, info, requestBody)
require.NoError(t, err)
parsed, taskErr := adaptor.ParseResponse(c, resp, info)
require.Nil(t, taskErr)
assert.Equal(t, "task_up_public", parsed.UpstreamTaskID, "the gateway's presenter emits its public id in the id field doubao reads")
task := &model.Task{
Action: info.Action,
Properties: model.Properties{OriginModelName: modelName, UpstreamModelName: modelName},
PrivateData: model.TaskPrivateData{UpstreamTaskID: parsed.UpstreamTaskID},
}
queryResp, err := adaptor.FetchTask(gateway.URL, "sk-gateway", task, "")
require.NoError(t, err)
queryBody, err := io.ReadAll(queryResp.Body)
require.NoError(t, err)
require.NoError(t, queryResp.Body.Close())
result, err := adaptor.ParseTaskResult(task, queryResp, queryBody)
require.NoError(t, err)
assert.Equal(t, "SUCCESS", result.Status)
assert.Equal(t, "https://cdn.example/v.mp4", result.Url)
assert.Equal(t, 1200, result.TotalTokens)
assert.Equal(t, []string{
"POST /doubao/api/v3/contents/generations/tasks",
"GET /doubao/api/v3/contents/generations/tasks/task_up_public",
}, seen)
}
func TestTaskAdaptorChainsSunoBatchFetchThroughNewAPIUpstream(t *testing.T) {
gin.SetMode(gin.TestMode)
service.InitHttpClient()
var seen []string
gateway := newAPIGateway(t, map[string]string{
"POST /suno/submit/MUSIC": `{"code":"success","message":"","data":"task_up_public"}`,
"POST /suno/fetch": `{"code":"success","message":"","data":[{"task_id":"task_up_public","status":"SUCCESS","progress":"100%","data":[{"id":"clip-1","status":"complete","audio_url":"https://cdn.example/a.mp3"}]}]}`,
}, &seen)
defer gateway.Close()
source, err := plugins.Source("sunoapi")
require.NoError(t, err)
plugin, err := pluginruntime.NewRegistry().RegisterFactory(source, pluginruntime.Options{Key: "sunoapi"})
require.NoError(t, err)
adaptor := New(plugin)
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeNewAPI, ChannelBaseUrl: gateway.URL, ApiKey: "sk-gateway", UpstreamModelName: "suno_music"},
OriginModelName: "suno_music",
TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "task_local", Action: "MUSIC"},
}
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/suno/submit/music", nil)
c.Set("task_request", relaycommon.TaskSubmitReq{Prompt: "a song about cats"})
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
requestBody, err := adaptor.BuildRequestBody(c, info)
require.NoError(t, err)
submitURL, err := adaptor.BuildRequestURL(info)
require.NoError(t, err)
assert.Equal(t, gateway.URL+"/suno/submit/MUSIC", submitURL)
resp, err := adaptor.DoRequest(c, info, requestBody)
require.NoError(t, err)
parsed, taskErr := adaptor.ParseResponse(c, resp, info)
require.Nil(t, taskErr)
assert.Equal(t, "task_up_public", parsed.UpstreamTaskID)
task := &model.Task{
Action: "MUSIC",
Properties: model.Properties{OriginModelName: "suno_music", UpstreamModelName: "suno_music"},
PrivateData: model.TaskPrivateData{UpstreamTaskID: parsed.UpstreamTaskID},
}
batchResp, err := adaptor.FetchBatchTasks(gateway.URL, "sk-gateway", []*model.Task{task}, "")
require.NoError(t, err)
batchBody, err := io.ReadAll(batchResp.Body)
require.NoError(t, err)
require.NoError(t, batchResp.Body.Close())
results, err := adaptor.ParseBatchResult([]*model.Task{task}, batchResp, batchBody)
require.NoError(t, err)
require.Contains(t, results, "task_up_public")
assert.Equal(t, "SUCCESS", results["task_up_public"].TaskInfo.Status)
assert.Equal(t, []string{"POST /suno/submit/MUSIC", "POST /suno/fetch"}, seen)
}
relay/channel/task/jsplugin/auth_test.go
+78
View File
@@ -9,7 +9,10 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay/channel"
vertexcore "github.com/QuantumNous/new-api/relay/channel/vertex"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/gin-gonic/gin"
@@ -75,3 +78,78 @@ export function parseTaskResult(){return {status:"SUCCESS"}}
require.NoError(t, adaptor.BuildRequestHeader(c, req, info))
assert.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
}
func TestUpstreamContextsUseGatewayBearerOnNewAPIChannels(t *testing.T) {
pluginAuthCache = sync.Map{}
original := acquireAccessToken
exchanges := 0
acquireAccessToken = func(_ vertexcore.Credentials, _ string) (string, error) {
exchanges++
return "access-token", nil
}
t.Cleanup(func() { acquireAccessToken = original; pluginAuthCache = sync.Map{} })
credentials, err := common.Marshal(vertexcore.Credentials{ProjectID: "project", ClientEmail: "a@example.com", PrivateKey: "secret"})
require.NoError(t, err)
const base = "https://upstream.example"
testCases := []struct {
name string
auth string
channelType int
key string
wantKind string
wantAuthHeader string
wantAPIKey any // nil when the raw key must stay hidden from the plugin
}{
{"api key plugin on a New API channel", "", constant.ChannelTypeNewAPI, "sk-gateway", pluginruntime.UpstreamKindNewAPI, "Bearer sk-gateway", "sk-gateway"},
{"api key plugin on a Task Plugin channel", "", constant.ChannelTypeTaskPlugin, "vendor-key", pluginruntime.UpstreamKindVendor, "vendor-key", "vendor-key"},
{"api key plugin on a legacy channel", "", 1002, "vendor-key", pluginruntime.UpstreamKindVendor, "vendor-key", "vendor-key"},
{"oauth2_jwt plugin on a New API channel", `,auth:{type:"oauth2_jwt"}`, constant.ChannelTypeNewAPI, "sk-gateway", pluginruntime.UpstreamKindNewAPI, "Bearer sk-gateway", "sk-gateway"},
{"oauth2_jwt plugin on a Task Plugin channel", `,auth:{type:"oauth2_jwt"}`, constant.ChannelTypeTaskPlugin, string(credentials), pluginruntime.UpstreamKindVendor, "Bearer access-token", nil},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
source := `
export const meta = {apiVersion:1,key:"upstream-probe",name:"Upstream Probe",version:"1.0.0",author:{name:"Test"},channelTypes:[1002],models:["probe"],fetchMode:"per_task",upstreams:["vendor","new_api"]` + tc.auth + `};
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.taskId}}
export function parseTaskResult(){return {status:"SUCCESS"}}
export function listArtifacts(){return [{key:"video",type:"video"}]}
export function buildContentRequest(ctx){return {url:ctx.baseUrl+"/content/"+ctx.upstream.kind,headers:{Authorization:ctx.authHeader,"X-Api-Key":String(ctx.apiKey)}}}
`
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
require.NoError(t, err)
adaptor := New(plugin)
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelType: tc.channelType, ChannelBaseUrl: base, ApiKey: tc.key}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
task := &model.Task{PrivateData: model.TaskPrivateData{UpstreamTaskID: "t1"}}
submit := adaptor.submitContext(c, info)
query, err := adaptor.queryContext(task, tc.key, base, "")
require.NoError(t, err)
batch, err := adaptor.batchQueryContext(tc.key, base, "", nil)
require.NoError(t, err)
assert.NotContains(t, submit, "authError")
for name, ctx := range map[string]map[string]any{"submit": submit, "query": query, "batch": batch} {
assert.Equal(t, map[string]any{"kind": tc.wantKind}, ctx["upstream"], name)
assert.Equal(t, tc.wantAuthHeader, ctx["authHeader"], name)
assert.Equal(t, tc.wantAPIKey, ctx["apiKey"], name)
assert.Equal(t, tc.wantAuthHeader, ctx["auth"].(map[string]any)["authHeader"], name)
}
content, err := adaptor.BuildContentRequest(t.Context(), task, "video", channel.TaskArtifactClientRequest{Method: http.MethodGet})
require.NoError(t, err)
assert.Equal(t, base+"/content/"+tc.wantKind, content.URL)
assert.Equal(t, tc.wantAuthHeader, content.Headers["Authorization"])
if tc.wantAPIKey == nil {
assert.Equal(t, "undefined", content.Headers["X-Api-Key"])
} else {
assert.Equal(t, tc.wantAPIKey, content.Headers["X-Api-Key"])
}
})
}
assert.Equal(t, 1, exchanges, "the vendor token exchange runs only for the oauth2_jwt vendor channel, never for a New API channel")
}
relay/channel/task/jsplugin/performance_test.go
+6 -2
View File
@@ -60,9 +60,13 @@ func BenchmarkTaskPluginRuntime(b *testing.B) {
var err error
switch name {
case "HasExport":
_, err = plugin.Engine.HasExport(ctx, "extractUsage")
if !plugin.Engine.HasExport("extractUsage") {
b.Fatal("extractUsage not exported")
}
case "MissingExport":
_, err = plugin.Engine.HasExport(ctx, "absentHook")
if plugin.Engine.HasExport("absentHook") {
b.Fatal("absentHook exported")
}
default:
_, err = plugin.Engine.Call(ctx, "extractUsage", input)
}
relay/channel/task/jsplugin/submit_stream.go
+10 -9
View File
@@ -5,7 +5,6 @@ import (
"fmt"
"mime"
"net/http"
"slices"
"strings"
"time"
@@ -37,7 +36,7 @@ func (a *TaskAdaptor) readSubmitEvents(parent context.Context, resp *http.Respon
scanner := helper.NewStreamScanner(resp.Body, maxTaskPluginPersistedJSONBytes+1)
hook := "parseSubmitEvent"
var accumulated *pluginruntime.JSONState
if slices.Contains(a.plugin.Meta.RequiredCapabilities, pluginruntime.CapabilitySubmitSSEDelta) {
if pluginruntime.UsesSubmitEventDelta(a.plugin.Meta, a.plugin.Engine) {
hook = "parseSubmitEventDelta"
accumulated = pluginruntime.NewJSONState(maxTaskPluginPersistedJSONBytes)
}
@@ -117,15 +116,17 @@ func (a *TaskAdaptor) readSubmitEvents(parent context.Context, resp *http.Respon
return nil, fmt.Errorf("invalid submit stream changes: %w", err)
}
}
// Keep the exact encoded-byte limit above. Plain JSON state can be
// isolated without parsing large strings again; codec-specific values
// (for example exported typed arrays) retain the old normalization.
var plainJSON bool
state, plainJSON = cloneJSONValue(nextState, 0)
if !plainJSON {
if err = common.Unmarshal(encoded, &state); err != nil {
// Keep the exact encoded-byte limit above. Plain JSON state goes back
// to the next event as is; codec-specific values (for example exported
// typed arrays) retain the old normalization.
if isPlainJSONValue(nextState, 0) {
state = nextState
} else {
var decoded any
if err = common.Unmarshal(encoded, &decoded); err != nil {
return nil, err
}
state = decoded
}
if done {
if accumulated != nil {