feat: sync deployed v1.0.0-rc.38-custom.42
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
@@ -2,6 +2,7 @@ package jsplugin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"cmp"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
@@ -18,6 +19,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
@@ -33,6 +35,7 @@ import (
|
||||
)
|
||||
|
||||
type requestDescriptor struct {
|
||||
ResponseType string `json:"responseType"`
|
||||
URL string `json:"url"`
|
||||
Method string `json:"method"`
|
||||
Headers map[string]string `json:"headers"`
|
||||
@@ -120,14 +123,24 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
|
||||
return taskErr
|
||||
}
|
||||
}
|
||||
if request, exists := c.Get("task_request"); exists {
|
||||
if err := a.validateResolvedUsageRequest(request); err != nil {
|
||||
request, hasRequest := c.Get("task_request")
|
||||
hasUsageProfiles := len(a.plugin.Meta.UsageProfiles) > 0
|
||||
if hasRequest && !hasUsageProfiles {
|
||||
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)
|
||||
}
|
||||
// The descriptor may rewrite the model. Validate profiled requests against
|
||||
// that final model before any quota calculation or upstream submission.
|
||||
if hasRequest && hasUsageProfiles {
|
||||
usageModel := cmp.Or(info.UpstreamModelName, info.OriginModelName)
|
||||
if err := a.validateResolvedUsageRequest(request, usageModel); err != nil {
|
||||
return service.TaskErrorWrapperLocal(err, "plugin_usage_invalid", http.StatusBadRequest)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -143,7 +156,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(), "extractUsage", usageContext)
|
||||
return a.usageRatios(c.Request.Context(), cmp.Or(info.UpstreamModelName, info.OriginModelName), "extractUsage", usageContext)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ExtractUsageFacts(c *gin.Context, info *relaycommon.RelayInfo) map[string]any {
|
||||
@@ -172,7 +185,7 @@ 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); err != nil {
|
||||
if _, err = a.validatedUsageRatios(facts, cmp.Or(info.UpstreamModelName, info.OriginModelName)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return facts, nil
|
||||
@@ -183,7 +196,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(), "extractUsageOnSubmit", a.submitContext(nil, info), data)
|
||||
ratios, err := a.usageRatios(context.Background(), cmp.Or(info.UpstreamModelName, info.OriginModelName), "extractUsageOnSubmit", a.submitContext(nil, info), data)
|
||||
if err != nil {
|
||||
a.logRejectedUsage("extractUsageOnSubmit", err)
|
||||
return nil
|
||||
@@ -199,7 +212,11 @@ func (a *TaskAdaptor) AdjustBillingOnComplete(task *model.Task, result *relaycom
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
a.applyCompletionUsageFacts(result, value)
|
||||
usageModel := ""
|
||||
if task != nil {
|
||||
usageModel = cmp.Or(task.Properties.UpstreamModelName, task.Properties.OriginModelName)
|
||||
}
|
||||
a.applyCompletionUsageFacts(result, value, usageModel)
|
||||
return 0
|
||||
}
|
||||
|
||||
@@ -423,6 +440,9 @@ func encodeFilePlaceholder(placeholder map[string]any, form *multipart.Form, lim
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, body io.Reader) (*http.Response, error) {
|
||||
wasStream := info.IsStream
|
||||
info.IsStream = false
|
||||
defer func() { info.IsStream = wasStream }()
|
||||
if a.submit != nil && strings.TrimSpace(a.submit.Method) != "" {
|
||||
originalMethod := c.Request.Method
|
||||
c.Request.Method = strings.ToUpper(a.submit.Method)
|
||||
@@ -431,24 +451,42 @@ func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, bod
|
||||
return channel.DoTaskApiRequest(a, c, info, body)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*channel.TaskSubmitResponse, *dto.TaskError) {
|
||||
func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (_ *channel.TaskSubmitResponse, taskErr *dto.TaskError) {
|
||||
started := time.Now()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
logger.LogDebug(c, "task_plugin subsystem=adaptor event=parse_submit_failed plugin=%q stage=read_response reason=read_failed status=%d", a.plugin.Meta.Key, resp.StatusCode)
|
||||
return nil, service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
|
||||
var responseBody any
|
||||
var err error
|
||||
streaming := a.submit != nil && a.submit.ResponseType == "sse"
|
||||
mediaType, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
|
||||
acceptedStream := streaming || mediaType == "text/event-stream"
|
||||
defer func() {
|
||||
if acceptedStream && taskErr != nil {
|
||||
taskErr.NoRetry = true
|
||||
}
|
||||
}()
|
||||
if !streaming && acceptedStream {
|
||||
return nil, service.TaskErrorWrapperLocal(fmt.Errorf("unexpected SSE response for a JSON submission"), "plugin_submit_response_invalid", http.StatusBadGateway)
|
||||
}
|
||||
logger.LogDebug(
|
||||
c,
|
||||
"task_plugin subsystem=adaptor event=submit_response_received plugin=%q status=%d body_bytes=%d",
|
||||
a.plugin.Meta.Key,
|
||||
resp.StatusCode,
|
||||
len(body),
|
||||
)
|
||||
responseBody := any(string(body))
|
||||
var decoded any
|
||||
if common.Unmarshal(body, &decoded) == nil {
|
||||
responseBody = decoded
|
||||
if streaming {
|
||||
// Every failure after accepting this stream is non-retryable: the
|
||||
// upstream may already have performed billable work.
|
||||
defer resp.Body.Close()
|
||||
responseBody, err = a.readSubmitEvents(c.Request.Context(), resp, a.submitContext(c, info))
|
||||
} else {
|
||||
var body []byte
|
||||
body, err = io.ReadAll(io.LimitReader(resp.Body, maxTaskPluginPersistedJSONBytes+1))
|
||||
if err == nil && len(body) > maxTaskPluginPersistedJSONBytes {
|
||||
err = fmt.Errorf("task submit response exceeds size limit")
|
||||
}
|
||||
responseBody = string(body)
|
||||
var decoded any
|
||||
if err == nil && common.Unmarshal(body, &decoded) == nil {
|
||||
responseBody = decoded
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
failure := service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusBadGateway)
|
||||
failure.NoRetry = streaming
|
||||
return nil, failure
|
||||
}
|
||||
headers := make(map[string][]string, len(resp.Header))
|
||||
maps.Copy(headers, resp.Header)
|
||||
@@ -485,34 +523,57 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r
|
||||
var taskData []byte
|
||||
if parsed.TaskData != nil {
|
||||
taskData, err = common.Marshal(parsed.TaskData)
|
||||
if err != nil {
|
||||
if err != nil || len(taskData) > maxTaskPluginPersistedJSONBytes {
|
||||
if err == nil {
|
||||
err = fmt.Errorf("task data exceeds size limit")
|
||||
}
|
||||
return nil, service.TaskErrorWrapper(err, "plugin_submit_response_invalid", http.StatusBadGateway)
|
||||
}
|
||||
}
|
||||
var immediate *relaycommon.TaskInfo
|
||||
if parsed.Immediate != nil {
|
||||
immediate = &relaycommon.TaskInfo{Code: parsed.Immediate.Code, TaskID: parsed.Immediate.TaskID, Status: parsed.Immediate.Status, Progress: parsed.Immediate.Progress, Reason: parsed.Immediate.Reason, Url: parsed.Immediate.URL, RemoteUrl: parsed.Immediate.RemoteURL}
|
||||
}
|
||||
logger.LogDebug(
|
||||
c,
|
||||
"task_plugin subsystem=adaptor event=parse_submit_complete plugin=%q status=%d task_data_bytes=%d immediate=%t elapsed_ms=%d",
|
||||
a.plugin.Meta.Key,
|
||||
resp.StatusCode,
|
||||
len(taskData),
|
||||
immediate != nil,
|
||||
time.Since(started).Milliseconds(),
|
||||
)
|
||||
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)))
|
||||
pluginState = nil
|
||||
}
|
||||
return &channel.TaskSubmitResponse{
|
||||
response := &channel.TaskSubmitResponse{
|
||||
UpstreamTaskID: parsed.TaskID,
|
||||
TaskData: taskData,
|
||||
Immediate: immediate,
|
||||
PluginState: pluginState,
|
||||
}, nil
|
||||
}
|
||||
logger.LogDebug(c, "task_plugin subsystem=adaptor event=parse_submit_complete plugin=%q status=%d task_data_bytes=%d immediate=%t elapsed_ms=%d",
|
||||
a.plugin.Meta.Key, resp.StatusCode, len(taskData), parsed.Immediate != nil, time.Since(started).Milliseconds())
|
||||
if parsed.Immediate == nil {
|
||||
return response, nil
|
||||
}
|
||||
immediate := &relaycommon.TaskInfo{
|
||||
Code: parsed.Immediate.Code, TaskID: parsed.Immediate.TaskID,
|
||||
Status: parsed.Immediate.Status, Progress: parsed.Immediate.Progress,
|
||||
Reason: parsed.Immediate.Reason, Url: parsed.Immediate.URL, RemoteUrl: parsed.Immediate.RemoteURL,
|
||||
CompletionTokens: positiveInt(parsed.Immediate.CompletionTokens), TotalTokens: positiveInt(parsed.Immediate.TotalTokens),
|
||||
}
|
||||
response.Immediate = immediate
|
||||
if immediate.Status != model.TaskStatusSuccess || !a.hasHook(c.Request.Context(), "extractUsageOnComplete") {
|
||||
return response, nil
|
||||
}
|
||||
// Match the polling hook context without persisting a temporary task.
|
||||
task := &model.Task{TaskID: info.PublicTaskID, Action: info.Action, Data: taskData,
|
||||
Properties: model.Properties{OriginModelName: info.OriginModelName, UpstreamModelName: info.UpstreamModelName}}
|
||||
task.PrivateData.UpstreamTaskID = parsed.TaskID
|
||||
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))
|
||||
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))
|
||||
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))
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) GetModelList() []string { return append([]string(nil), a.plugin.Meta.Models...) }
|
||||
@@ -680,7 +741,8 @@ 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 {
|
||||
a.applyCompletionUsageFacts(&info, facts)
|
||||
usageModel, _ := itemCtx["upstreamModel"].(string)
|
||||
a.applyCompletionUsageFacts(&info, facts, usageModel)
|
||||
}
|
||||
}
|
||||
results[item.TaskID] = &service.BatchTaskResult{TaskInfo: info, Action: item.Action, SubmitTime: item.SubmitTime, StartTime: item.StartTime, FinishTime: item.FinishTime, Data: item.Data}
|
||||
@@ -745,7 +807,8 @@ 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 {
|
||||
a.applyCompletionUsageFacts(result, facts)
|
||||
usageModel, _ := ctx["upstreamModel"].(string)
|
||||
a.applyCompletionUsageFacts(result, facts, usageModel)
|
||||
}
|
||||
}
|
||||
taskStatus := model.TaskStatus(result.Status)
|
||||
@@ -760,20 +823,20 @@ func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, bod
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, facts any) {
|
||||
values, err := a.validatedCompletionUsageFacts(facts)
|
||||
func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, facts any, modelName string) error {
|
||||
values, err := a.validatedCompletionUsageFacts(facts, modelName)
|
||||
if err != nil {
|
||||
a.logRejectedUsage("extractUsageOnComplete", err)
|
||||
return
|
||||
return err
|
||||
}
|
||||
if len(values) == 0 {
|
||||
return
|
||||
return nil
|
||||
}
|
||||
result.UsageFacts = values
|
||||
if units := positiveInt(values["upstreamUnits"]); units > 0 {
|
||||
result.CompletionTokens = units
|
||||
result.TotalTokens = units
|
||||
return
|
||||
return nil
|
||||
}
|
||||
if completionTokens, exists := values["completionTokens"]; exists {
|
||||
result.CompletionTokens = positiveInt(completionTokens)
|
||||
@@ -781,6 +844,7 @@ func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, fa
|
||||
if totalTokens, exists := values["totalTokens"]; exists {
|
||||
result.TotalTokens = positiveInt(totalTokens)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
@@ -805,26 +869,24 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rendered := kitdto.NewOpenAIVideo()
|
||||
if err = common.Unmarshal(encoded, rendered); err != nil {
|
||||
var rendered map[string]any
|
||||
if err = common.Unmarshal(encoded, &rendered); err != nil || 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
|
||||
// identity and lifecycle, including completion timestamps after settlement.
|
||||
host := task.ToOpenAIVideo()
|
||||
rendered.ID = host.ID
|
||||
rendered.Object = host.Object
|
||||
rendered.TaskID = ""
|
||||
rendered.Status = host.Status
|
||||
rendered.Progress = host.Progress
|
||||
rendered.CreatedAt = host.CreatedAt
|
||||
rendered.Model = host.Model
|
||||
rendered.CompletedAt = host.CompletedAt
|
||||
for key := range rendered.Metadata {
|
||||
if strings.EqualFold(key, "url") {
|
||||
delete(rendered.Metadata, key)
|
||||
}
|
||||
}
|
||||
if len(rendered.Metadata) == 0 {
|
||||
rendered.Metadata = nil
|
||||
rendered["id"] = host.ID
|
||||
rendered["object"] = host.Object
|
||||
delete(rendered, "task_id")
|
||||
rendered["status"] = host.Status
|
||||
rendered["progress"] = host.Progress
|
||||
rendered["created_at"] = host.CreatedAt
|
||||
rendered["model"] = host.Model
|
||||
if host.CompletedAt != 0 {
|
||||
rendered["completed_at"] = host.CompletedAt
|
||||
} else {
|
||||
delete(rendered, "completed_at")
|
||||
}
|
||||
return common.Marshal(rendered)
|
||||
}
|
||||
@@ -1139,6 +1201,16 @@ func (a *TaskAdaptor) buildSubmit(c *gin.Context, info *relaycommon.RelayInfo) (
|
||||
)
|
||||
return nil, err
|
||||
}
|
||||
if descriptor.ResponseType == "" {
|
||||
descriptor.ResponseType = "json"
|
||||
}
|
||||
allowed := a.plugin.Meta.SubmitResponseTypes
|
||||
if len(allowed) == 0 {
|
||||
allowed = []string{"json"}
|
||||
}
|
||||
if !slices.Contains(allowed, descriptor.ResponseType) {
|
||||
return nil, fmt.Errorf("plugin does not support submit response type %q", descriptor.ResponseType)
|
||||
}
|
||||
if strings.TrimSpace(descriptor.URL) == "" {
|
||||
logger.LogDebug(c, "task_plugin subsystem=adaptor event=build_submit_failed plugin=%q stage=validate_url reason=empty_url", a.plugin.Meta.Key)
|
||||
return nil, fmt.Errorf("plugin returned an empty submit URL")
|
||||
@@ -1294,7 +1366,7 @@ func (a *TaskAdaptor) submitContext(c *gin.Context, info *relaycommon.RelayInfo)
|
||||
return ctx
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) usageRatios(ctx context.Context, hook string, args ...any) (map[string]float64, error) {
|
||||
func (a *TaskAdaptor) usageRatios(ctx context.Context, modelName, hook string, args ...any) (map[string]float64, error) {
|
||||
if !a.hasHook(ctx, hook) {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -1312,7 +1384,7 @@ func (a *TaskAdaptor) usageRatios(ctx context.Context, hook string, args ...any)
|
||||
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)
|
||||
ratios, err := a.validatedUsageRatios(facts, modelName)
|
||||
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
|
||||
@@ -1329,15 +1401,16 @@ func (a *TaskAdaptor) usageRatios(ctx context.Context, hook string, args ...any)
|
||||
return ratios, nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) validateResolvedUsageRequest(request any) error {
|
||||
return a.validateResolvedUsageValue(jsonValue(request))
|
||||
func (a *TaskAdaptor) validateResolvedUsageRequest(request any, modelName string) error {
|
||||
schema, _ := a.plugin.Meta.UsageForModel(modelName)
|
||||
return a.validateResolvedUsageValue(jsonValue(request), schema)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) validateResolvedUsageValue(value any) error {
|
||||
func (a *TaskAdaptor) validateResolvedUsageValue(value any, usageSchema map[string]pluginruntime.UsageFieldSchema) error {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
for key, item := range typed {
|
||||
if schema, declared := a.plugin.Meta.UsageSchema[key]; declared {
|
||||
if schema, declared := usageSchema[key]; declared {
|
||||
if _, err := validateUsageValue(item, schema, true); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -1346,13 +1419,13 @@ func (a *TaskAdaptor) validateResolvedUsageValue(value any) error {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := a.validateResolvedUsageValue(item); err != nil {
|
||||
if err := a.validateResolvedUsageValue(item, usageSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, item := range typed {
|
||||
if err := a.validateResolvedUsageValue(item); err != nil {
|
||||
if err := a.validateResolvedUsageValue(item, usageSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
@@ -1360,10 +1433,11 @@ func (a *TaskAdaptor) validateResolvedUsageValue(value any) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any) (map[string]float64, error) {
|
||||
func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any, modelName string) (map[string]float64, error) {
|
||||
usageSchema, _ := a.plugin.Meta.UsageForModel(modelName)
|
||||
ratios := make(map[string]float64)
|
||||
for key, value := range facts {
|
||||
if schema, declared := a.plugin.Meta.UsageSchema[key]; declared {
|
||||
if schema, declared := usageSchema[key]; declared {
|
||||
number, err := validateUsageValue(value, schema, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -1390,6 +1464,9 @@ func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any) (map[string]flo
|
||||
if err := validateUsageNumberLimit(number, limit); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Keep numeric facts usable by existing expressions even when the
|
||||
// selected profile no longer declares that legacy extension field.
|
||||
facts[key] = number
|
||||
if number > 0 {
|
||||
ratios[key] = number
|
||||
}
|
||||
@@ -1397,7 +1474,8 @@ func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any) (map[string]flo
|
||||
return ratios, nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any) (map[string]any, error) {
|
||||
func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any, modelName string) (map[string]any, error) {
|
||||
usageSchema, _ := a.plugin.Meta.UsageForModel(modelName)
|
||||
if facts == nil {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -1408,7 +1486,7 @@ func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any) (map[string]any,
|
||||
validated := make(map[string]any, len(values))
|
||||
for key, value := range values {
|
||||
validated[key] = value
|
||||
if schema, declared := a.plugin.Meta.UsageSchema[key]; declared {
|
||||
if schema, declared := usageSchema[key]; declared {
|
||||
number, err := validateUsageValue(value, schema, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -1436,6 +1514,10 @@ func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any) (map[string]any,
|
||||
return nil, fmt.Errorf("plugin usage value must be a finite non-negative number")
|
||||
}
|
||||
validated[key] = float64(common.QuotaFromFloat(number))
|
||||
default:
|
||||
if number, numeric := usageNumber(value, false); numeric {
|
||||
validated[key] = number
|
||||
}
|
||||
}
|
||||
}
|
||||
return validated, nil
|
||||
@@ -1549,6 +1631,9 @@ func convert(value any, target any) error {
|
||||
}
|
||||
|
||||
func jsonValue(value any) any {
|
||||
if normalized, ok := cloneJSONValue(value, 0); ok {
|
||||
return normalized
|
||||
}
|
||||
data, err := common.Marshal(value)
|
||||
if err != nil {
|
||||
return value
|
||||
@@ -1559,6 +1644,58 @@ 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
|
||||
// float64 numbers and nulls. Other Go types still use the configured codec.
|
||||
func cloneJSONValue(value any, depth int) (any, bool) {
|
||||
if depth > 64 {
|
||||
return nil, false
|
||||
}
|
||||
switch typed := value.(type) {
|
||||
case nil, bool:
|
||||
return typed, true
|
||||
case string:
|
||||
return typed, utf8.ValidString(typed)
|
||||
case float64:
|
||||
return typed, !math.IsNaN(typed) && !math.IsInf(typed, 0)
|
||||
case int64:
|
||||
return float64(typed), true
|
||||
case int:
|
||||
return float64(typed), true
|
||||
case map[string]any:
|
||||
if typed == nil {
|
||||
return nil, true
|
||||
}
|
||||
cloned := make(map[string]any, len(typed))
|
||||
for key, item := range typed {
|
||||
if !utf8.ValidString(key) {
|
||||
return nil, false
|
||||
}
|
||||
child, ok := cloneJSONValue(item, depth+1)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
cloned[key] = child
|
||||
}
|
||||
return cloned, true
|
||||
case []any:
|
||||
if typed == nil {
|
||||
return nil, true
|
||||
}
|
||||
cloned := make([]any, len(typed))
|
||||
for index, item := range typed {
|
||||
child, ok := cloneJSONValue(item, depth+1)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
cloned[index] = child
|
||||
}
|
||||
return cloned, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
func positiveInt(value any) int {
|
||||
switch number := value.(type) {
|
||||
case int64:
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"math"
|
||||
"mime/multipart"
|
||||
@@ -16,7 +17,9 @@ import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/plugins"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
@@ -532,69 +535,90 @@ func TestTaskAdaptorMapsJSContract(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestTaskAdaptorSanitizesOpenAIVideoRendererOutput(t *testing.T) {
|
||||
source := `
|
||||
export const meta = {
|
||||
apiVersion: 1, key: "safe-video", name: "Safe Video", version: "1.0.0",
|
||||
author: {name: "Test"}, models: ["model"], fetchMode: "per_task", protocols: ["openai_video"],
|
||||
};
|
||||
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl + "/submit"}; }
|
||||
export function parseSubmitResponse() { return {taskId: "upstream"}; }
|
||||
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl + "/query"}; }
|
||||
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
||||
export function listArtifacts() { return []; }
|
||||
export function buildContentRequest() { throw new Error("artifact_not_found"); }
|
||||
export const protocols = {openai_video: {
|
||||
decodeRequest: function(ctx) { return {kind:"submit", model:ctx.model, requestBody:ctx.body.value}; },
|
||||
render: function() {
|
||||
return {
|
||||
id: "upstream-id",
|
||||
task_id: "upstream-task-id",
|
||||
object: "provider-video",
|
||||
model: "model",
|
||||
status: "completed",
|
||||
progress: 100,
|
||||
created_at: 10,
|
||||
completed_at: 20,
|
||||
metadata: {
|
||||
url: "https://upstream.example/video.mp4",
|
||||
URL: "https://upstream.example/uppercase.mp4",
|
||||
label: "safe",
|
||||
},
|
||||
url: "https://upstream.example/top-level.mp4",
|
||||
upstream_url: "https://upstream.example/unknown.mp4",
|
||||
provider_payload: {task_id: "upstream-task-id"},
|
||||
};
|
||||
}}};
|
||||
`
|
||||
func TestTaskAdaptorPreservesSoraVideoResponseFields(t *testing.T) {
|
||||
source, err := plugins.Source("sora")
|
||||
require.NoError(t, err)
|
||||
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
adaptor := New(plugin)
|
||||
|
||||
rendered, err := adaptor.ConvertToOpenAIVideo(&model.Task{
|
||||
TaskID: "task_public",
|
||||
Status: model.TaskStatusInProgress,
|
||||
Properties: model.Properties{OriginModelName: "origin-model"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
for _, tc := range []struct {
|
||||
status model.TaskStatus
|
||||
want string
|
||||
}{
|
||||
{model.TaskStatusInProgress, "in_progress"},
|
||||
{model.TaskStatusSuccess, "completed"},
|
||||
{model.TaskStatusFailure, "failed"},
|
||||
} {
|
||||
t.Run(string(tc.status), func(t *testing.T) {
|
||||
task := &model.Task{
|
||||
TaskID: "task_public",
|
||||
Status: tc.status,
|
||||
Progress: "42%",
|
||||
CreatedAt: 100,
|
||||
FinishTime: 200,
|
||||
Properties: model.Properties{OriginModelName: "origin-model"},
|
||||
PrivateData: model.TaskPrivateData{UpstreamTaskID: "upstream-task-id"},
|
||||
Data: []byte(`{
|
||||
"id":"upstream-task-id","task_id":"upstream-task-id",
|
||||
"object":"provider-video","model":"provider-model",
|
||||
"status":"completed","progress":100,"created_at":10,"completed_at":20,
|
||||
"url":"https://cdn.example/video.mp4",
|
||||
"metadata":{"url":"https://cdn.example/video.mp4","URL":"https://cdn.example/uppercase.mp4"},
|
||||
"provider_payload":{"task_id":"upstream-task-id","items":[{"enabled":false,"count":0,"value":null}]},
|
||||
"seconds":8,"resolution":"720p","aspect_ratio":"16:9",
|
||||
"reference_images":["https://cdn.example/reference.png"],
|
||||
"error":{"code":"provider_error","message":"provider rejected request","detail":{"retryable":false}}
|
||||
}`),
|
||||
}
|
||||
rendered, err := adaptor.ConvertToOpenAIVideo(task)
|
||||
require.NoError(t, err)
|
||||
|
||||
var video dto.OpenAIVideo
|
||||
require.NoError(t, common.Unmarshal(rendered, &video))
|
||||
assert.Equal(t, "task_public", video.ID)
|
||||
assert.Equal(t, "video", video.Object)
|
||||
assert.Empty(t, video.TaskID)
|
||||
assert.Equal(t, "origin-model", video.Model)
|
||||
assert.Zero(t, video.CompletedAt)
|
||||
assert.Equal(t, map[string]any{"label": "safe"}, video.Metadata)
|
||||
var fields map[string]any
|
||||
require.NoError(t, common.Unmarshal(rendered, &fields))
|
||||
assert.Equal(t, "task_public", fields["id"])
|
||||
assert.NotContains(t, fields, "task_id")
|
||||
assert.Equal(t, "video", fields["object"])
|
||||
assert.Equal(t, "origin-model", fields["model"])
|
||||
assert.Equal(t, tc.want, fields["status"])
|
||||
assert.Equal(t, float64(42), fields["progress"])
|
||||
assert.Equal(t, float64(100), fields["created_at"])
|
||||
if tc.status == model.TaskStatusSuccess {
|
||||
assert.Equal(t, float64(200), fields["completed_at"])
|
||||
} else {
|
||||
assert.NotContains(t, fields, "completed_at")
|
||||
}
|
||||
assert.Equal(t, "https://cdn.example/video.mp4", fields["url"])
|
||||
assert.Equal(t, map[string]any{
|
||||
"url": "https://cdn.example/video.mp4",
|
||||
"URL": "https://cdn.example/uppercase.mp4",
|
||||
}, fields["metadata"])
|
||||
assert.Equal(t, map[string]any{
|
||||
"task_id": "task_public",
|
||||
"items": []any{map[string]any{"enabled": false, "count": float64(0), "value": nil}},
|
||||
}, fields["provider_payload"])
|
||||
assert.Equal(t, float64(8), fields["seconds"])
|
||||
assert.Equal(t, "720p", fields["resolution"])
|
||||
assert.Equal(t, "16:9", fields["aspect_ratio"])
|
||||
assert.Equal(t, []any{"https://cdn.example/reference.png"}, fields["reference_images"])
|
||||
assert.Equal(t, map[string]any{
|
||||
"code": "provider_error", "message": "provider rejected request",
|
||||
"detail": map[string]any{"retryable": false},
|
||||
}, fields["error"])
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
var fields map[string]any
|
||||
require.NoError(t, common.Unmarshal(rendered, &fields))
|
||||
assert.NotContains(t, fields, "url")
|
||||
assert.NotContains(t, fields, "upstream_url")
|
||||
assert.NotContains(t, fields, "provider_payload")
|
||||
assert.NotContains(t, fields, "completed_at")
|
||||
assert.NotContains(t, string(rendered), "upstream.example")
|
||||
assert.NotContains(t, string(rendered), "upstream-task-id")
|
||||
func TestTaskAdaptorRejectsNonObjectOpenAIVideoRendererOutput(t *testing.T) {
|
||||
for _, value := range []string{"null", "[]", `"video"`, "42", "false"} {
|
||||
t.Run(value, func(t *testing.T) {
|
||||
source := strings.Replace(mockPlugin, `return {id: task.task_id, status: "completed"};`, "return "+value+";", 1)
|
||||
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
_, err = New(plugin).ConvertToOpenAIVideo(&model.Task{TaskID: "task_public"})
|
||||
require.ErrorContains(t, err, "invalid OpenAI video object")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskAdaptorPreservesOpenAIVideoFailureSlotsAndOwnsLifecycle(t *testing.T) {
|
||||
@@ -1413,3 +1437,407 @@ export function parseBatchResult(){return [];}
|
||||
assert.Equal(t, []any{"model-a", "model-b"}, captured["models"])
|
||||
assert.Equal(t, false, captured["hasRequestBody"])
|
||||
}
|
||||
|
||||
func TestTaskAdaptorUsageProfilesFollowExecutionModel(t *testing.T) {
|
||||
const source = `
|
||||
export const meta = {
|
||||
apiVersion:1, key:"profile-usage", name:"Profile Usage", version:"1.0.0", author:{name:"Test"},
|
||||
models:["image", "video"], fetchMode:"batch",
|
||||
usageSchema:{units:{type:"number",unit:"second"}},
|
||||
usageProfiles:[
|
||||
{models:["image"],schema:{units:{type:"number",unit:"count"},mode:{enum:["image"]}}},
|
||||
{models:["video"],schema:{units:{type:"number",unit:"token"},mode:{enum:["video"]}},
|
||||
examples:[{label:"video",facts:{units:1,mode:"video"}}]}
|
||||
]
|
||||
};
|
||||
export function buildSubmitRequest(ctx) {
|
||||
return {url:ctx.baseUrl+"/submit",rewriteModel:ctx.requestBody.rewriteTo || ""};
|
||||
}
|
||||
export function parseSubmitResponse(){return {taskId:"task"};}
|
||||
export function buildQueryRequest(){return {url:"https://provider.example/task"};}
|
||||
export function buildBatchQueryRequest(){return {url:"https://provider.example/tasks"};}
|
||||
export function parseTaskResult(){return {status:"SUCCESS"};}
|
||||
export function parseBatchResult(ctx,body){return body.items;}
|
||||
export function extractUsage(ctx){return {units:ctx.requestBody.hookUnits,mode:ctx.requestBody.hookMode,legacyRatio:2};}
|
||||
export function extractUsageOnSubmit(ctx,body){return body.usage;}
|
||||
export function extractUsageOnComplete(ctx,result,body){return body.usage;}
|
||||
`
|
||||
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
newRequest := func(t *testing.T, upstream string, body map[string]any) (*TaskAdaptor, *gin.Context, *relaycommon.RelayInfo) {
|
||||
t.Helper()
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "public-alias",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
UpstreamModelName: upstream, ChannelBaseUrl: "https://provider.example",
|
||||
},
|
||||
TaskRelayInfo: &relaycommon.TaskRelayInfo{},
|
||||
}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/native/submit", nil)
|
||||
c.Set("task_request", body)
|
||||
return adaptor, c, info
|
||||
}
|
||||
|
||||
for _, tc := range []struct {
|
||||
name, upstream, rewrite, mode string
|
||||
units float64
|
||||
wantError bool
|
||||
}{
|
||||
{"mapped image", "image", "", "image", 2, false},
|
||||
{"mapped video token units", "video", "", "video", 500000, false},
|
||||
{"rewritten model", "image", "video", "video", 500000, false},
|
||||
{"image count ceiling", "image", "", "image", float64(dto.MaxImageN + 1), true},
|
||||
{"profile enum", "image", "", "video", 2, true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
adaptor, c, info := newRequest(t, tc.upstream, map[string]any{
|
||||
"hookUnits": tc.units, "hookMode": tc.mode, "rewriteTo": tc.rewrite,
|
||||
})
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
facts, err := adaptor.ExtractUsageFactsValidated(c, info)
|
||||
if tc.wantError {
|
||||
require.Error(t, err)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]any{"units": tc.units, "mode": tc.mode, "legacyRatio": 2.0}, facts)
|
||||
ratios, err := adaptor.EstimateBillingValidated(c, info)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]float64{"units": tc.units, "legacyRatio": 2}, ratios)
|
||||
body, err := common.Marshal(map[string]any{"usage": facts})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, ratios, adaptor.AdjustBillingOnSubmit(info, body))
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("request revalidated after model rewrite", func(t *testing.T) {
|
||||
adaptor, c, info := newRequest(t, "video", map[string]any{
|
||||
"rewriteTo": "image", "metadata": map[string]any{"units": dto.MaxImageN + 1},
|
||||
})
|
||||
taskErr := adaptor.ValidateRequestAndSetAction(c, info)
|
||||
require.NotNil(t, taskErr)
|
||||
assert.Equal(t, "plugin_usage_invalid", taskErr.Code)
|
||||
})
|
||||
|
||||
t.Run("request enum belongs to the rewritten model", func(t *testing.T) {
|
||||
adaptor, c, info := newRequest(t, "video", map[string]any{
|
||||
"rewriteTo": "image", "mode": "image", "units": 2,
|
||||
})
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
assert.Equal(t, "image", info.UpstreamModelName)
|
||||
})
|
||||
|
||||
t.Run("completion selects each persisted model including legacy origin fallback", func(t *testing.T) {
|
||||
adaptor, _, _ := newRequest(t, "image", map[string]any{})
|
||||
for _, tc := range []struct {
|
||||
upstream, origin, mode string
|
||||
units float64
|
||||
accepted bool
|
||||
}{
|
||||
{"video", "public-alias", "video", 500000, true},
|
||||
{"image", "public-alias", "image", 2, true},
|
||||
{"", "video", "video", 500000, true},
|
||||
{"image", "public-alias", "image", float64(dto.MaxImageN + 1), false},
|
||||
{"video", "public-alias", "image", 1, false},
|
||||
} {
|
||||
task := &model.Task{Properties: model.Properties{OriginModelName: tc.origin, UpstreamModelName: tc.upstream}}
|
||||
body, err := common.Marshal(map[string]any{"usage": map[string]any{"units": tc.units, "mode": tc.mode}})
|
||||
require.NoError(t, err)
|
||||
result, err := adaptor.ParseTaskResult(task, &http.Response{StatusCode: http.StatusOK}, body)
|
||||
require.NoError(t, err)
|
||||
if tc.accepted {
|
||||
assert.Equal(t, map[string]any{"units": tc.units, "mode": tc.mode}, result.UsageFacts)
|
||||
} else {
|
||||
assert.Nil(t, result.UsageFacts)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("legacy completion facts remain usable by stored expressions", func(t *testing.T) {
|
||||
updated, err := pluginruntime.NewRegistry().Register(strings.Replace(source,
|
||||
`export function extractUsageOnComplete(ctx,result,body){return body.usage;}`,
|
||||
`export function extractUsageOnComplete(){return {legacyRatio:3};}`, 1), pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
_, _, info := newRequest(t, "image", map[string]any{})
|
||||
adaptor := New(updated)
|
||||
adaptor.Init(info)
|
||||
result, err := adaptor.ParseTaskResult(&model.Task{Properties: model.Properties{UpstreamModelName: "image"}},
|
||||
&http.Response{StatusCode: http.StatusOK}, []byte(`{}`))
|
||||
require.NoError(t, err)
|
||||
cost, _, err := billingexpr.RunExprWithRequest(`tier("legacy", u("legacyRatio") * 2)`,
|
||||
billingexpr.TokenParams{}, billingexpr.RequestInput{Usage: result.UsageFacts})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 6.0, cost)
|
||||
})
|
||||
|
||||
t.Run("mixed batch selects schema by task rather than adaptor", func(t *testing.T) {
|
||||
adaptor, _, _ := newRequest(t, "image", map[string]any{})
|
||||
tasks := []*model.Task{
|
||||
{TaskID: "video-task", Properties: model.Properties{UpstreamModelName: "video"}},
|
||||
{TaskID: "image-task", Properties: model.Properties{UpstreamModelName: "image"}},
|
||||
{TaskID: "invalid-image", Properties: model.Properties{UpstreamModelName: "image"}},
|
||||
}
|
||||
body := []byte(`{"items":[
|
||||
{"taskId":"image-task","status":"SUCCESS","data":{"usage":{"units":2,"mode":"image"}}},
|
||||
{"taskId":"video-task","status":"SUCCESS","data":{"usage":{"units":500000,"mode":"video"}}},
|
||||
{"taskId":"invalid-image","status":"SUCCESS","data":{"usage":{"units":129,"mode":"image"}}}
|
||||
]}`)
|
||||
results, err := adaptor.ParseBatchResult(tasks, &http.Response{StatusCode: http.StatusOK}, body)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 3)
|
||||
assert.Equal(t, map[string]any{"units": 500000.0, "mode": "video"}, results["video-task"].TaskInfo.UsageFacts)
|
||||
assert.Equal(t, map[string]any{"units": 2.0, "mode": "image"}, results["image-task"].TaskInfo.UsageFacts)
|
||||
assert.Nil(t, results["invalid-image"].TaskInfo.UsageFacts)
|
||||
})
|
||||
}
|
||||
|
||||
const documentStreamPlugin = `
|
||||
export const meta = {apiVersion:1,key:"document-stream",name:"Document stream",version:"1.0.0",author:{name:"Test"},models:["document"],fetchMode:"per_task",submitResponseTypes:["json","sse"],usageSchema:{units:{type:"number",unit:"count"}}};
|
||||
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/compile",responseType:ctx.requestBody.responseType || "sse"};}
|
||||
export function parseSubmitEvent(ctx,event,previous) {
|
||||
const chunk = JSON.parse(event.data);
|
||||
if (chunk.error) throw new Error("provider stream failure");
|
||||
if (chunk.badState) return {state:null};
|
||||
if (chunk.largeState) return {state:{document:"x".repeat(1048577)},done:chunk.complete === true};
|
||||
if (chunk.escapedState) return {state:{document:"<".repeat(200000)},done:true};
|
||||
if (chunk.resetState) return {state:{document:"",units:0},done:true};
|
||||
const state = Object.assign({}, previous || {document:"",units:0});
|
||||
state.document += chunk.part || "";
|
||||
if (chunk.units !== undefined) state.units = chunk.units;
|
||||
state.event = event.event; state.id = event.id;
|
||||
return {state:state,done:chunk.complete === true};
|
||||
}
|
||||
export function parseSubmitResponse(ctx,response){return {taskId:"vendor-document",taskData:response.body,immediate:{status:"SUCCESS"},state:{revision:1}};}
|
||||
export function parseTaskResult(){return {status:"SUCCESS"};}
|
||||
export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"};}
|
||||
export function extractUsageOnComplete(ctx,result,body){
|
||||
if(ctx.model!=="alias" || ctx.upstreamModel!=="document" || ctx.taskId!=="vendor-document" || ctx.publicTaskId!=="public-document" || ctx.state.revision!==1 || ctx.data.document!==body.document) throw new Error("invalid completion context");
|
||||
return {units:body.units};
|
||||
}
|
||||
`
|
||||
|
||||
func TestTaskSubmitStreamContract(t *testing.T) {
|
||||
plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
for _, tc := range []struct {
|
||||
name, body, contentType string
|
||||
valid bool
|
||||
}{
|
||||
{"multiline and CRLF", ": heartbeat\r\nid: doc-1\r\nevent: update\r\ndata: {\"part\":\r\ndata: \"hello\",\"units\":2}\r\n\r\ndata: {\"part\":\"world\",\"units\":3,\"complete\":true}\n\n", "text/event-stream; charset=utf-8", true},
|
||||
{"zero actual units", "data: {\"part\":\"free\",\"units\":0,\"complete\":true}\n\n", "text/event-stream", true},
|
||||
{"premature EOF", "data: {\"part\":\"partial\"}\n\n", "text/event-stream", false},
|
||||
{"unterminated event", "data: {\"complete\":true}", "text/event-stream", false},
|
||||
{"provider error", "data: {\"error\":true}\n\n", "text/event-stream", false},
|
||||
{"invalid hook result", "data: {\"badState\":true}\n\n", "text/event-stream", false},
|
||||
{"state limit", "data: {\"largeState\":true}\n\n", "text/event-stream", false},
|
||||
{"intermediate state limit", "data: {\"largeState\":true}\n\ndata: {\"resetState\":true}\n\n", "text/event-stream", false},
|
||||
{"encoded state limit", "data: {\"escapedState\":true}\n\n", "text/event-stream", false},
|
||||
{"event limit", "data: " + strings.Repeat("x", maxTaskPluginPersistedJSONBytes) + "\n\n", "text/event-stream", false},
|
||||
{"wrong response type", "{}", "application/json", false},
|
||||
{"undeclared stream", "data: {}\n\n", "text/event-stream", false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "alias", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "document", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "public-document"}}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
c, recorder := gin.CreateTestContext(httptest.NewRecorder())
|
||||
_ = recorder
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/compile", nil)
|
||||
requestBody := map[string]any{}
|
||||
if tc.name == "undeclared stream" {
|
||||
requestBody["responseType"] = "json"
|
||||
}
|
||||
c.Set("task_request", requestBody)
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
response := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": {tc.contentType}}, Body: io.NopCloser(strings.NewReader(tc.body))}
|
||||
parsed, taskErr := adaptor.ParseResponse(c, response, info)
|
||||
assert.False(t, c.Writer.Written())
|
||||
if !tc.valid {
|
||||
require.NotNil(t, taskErr)
|
||||
assert.True(t, taskErr.NoRetry)
|
||||
return
|
||||
}
|
||||
require.Nil(t, taskErr)
|
||||
require.NotNil(t, parsed.Immediate)
|
||||
if tc.name == "zero actual units" {
|
||||
assert.Equal(t, map[string]any{"units": float64(0)}, parsed.Immediate.UsageFacts)
|
||||
return
|
||||
}
|
||||
assert.Equal(t, map[string]any{"units": float64(3)}, parsed.Immediate.UsageFacts)
|
||||
assert.JSONEq(t, `{"document":"helloworld","units":3,"event":"message","id":"doc-1"}`, string(parsed.TaskData))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskSubmitStreamCancellationClosesReader(t *testing.T) {
|
||||
plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
reader, writer := io.Pipe()
|
||||
defer writer.Close()
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
cancel()
|
||||
_, err = New(plugin).readSubmitEvents(ctx, &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: reader}, nil)
|
||||
require.ErrorIs(t, err, context.Canceled)
|
||||
_, err = writer.Write([]byte("data: {}\n\n"))
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestTaskSubmitStreamIdleTimeout(t *testing.T) {
|
||||
previous := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 1
|
||||
t.Cleanup(func() { constant.StreamingTimeout = previous })
|
||||
plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
reader, writer := io.Pipe()
|
||||
defer writer.Close()
|
||||
_, err = New(plugin).readSubmitEvents(t.Context(), &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: reader}, nil)
|
||||
require.ErrorContains(t, err, "idle timeout")
|
||||
}
|
||||
|
||||
func TestPluginJSONValuesPreserveCodecNormalizationAndIsolation(t *testing.T) {
|
||||
for _, value := range []any{
|
||||
map[string]any{"units": int64(3), "nested": []any{true, "<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 {
|
||||
Units int `json:"units"`
|
||||
}{Units: 0},
|
||||
} {
|
||||
encoded, err := common.Marshal(value)
|
||||
require.NoError(t, err)
|
||||
var expected any
|
||||
require.NoError(t, common.Unmarshal(encoded, &expected))
|
||||
assert.Equal(t, expected, jsonValue(value))
|
||||
}
|
||||
source := map[string]any{"items": []any{map[string]any{"label": "original"}}}
|
||||
copy := jsonValue(source).(map[string]any)
|
||||
copy["items"].([]any)[0].(map[string]any)["label"] = "changed"
|
||||
assert.Equal(t, "original", source["items"].([]any)[0].(map[string]any)["label"])
|
||||
}
|
||||
|
||||
func TestTaskSubmitHooksReceiveIndependentRequestSnapshots(t *testing.T) {
|
||||
source := `
|
||||
export const meta={apiVersion:1,key:"request-copy",name:"Request copy",version:"1.0.0",author:{name:"Test"},models:["copy"],fetchMode:"per_task",usageSchema:{units:{type:"number",unit:"count"}}};
|
||||
export function buildSubmitRequest(ctx){ctx.requestBody.units=3;ctx.requestBody.nested.label="built";return {url:ctx.baseUrl+"/submit",body:ctx.requestBody};}
|
||||
export function extractUsage(ctx){if(ctx.requestBody.nested.label!=="original")throw new Error("request changed");return {units:ctx.requestBody.units};}
|
||||
export function parseSubmitResponse(){return {taskId:"copy"};}
|
||||
export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"};}
|
||||
export function parseTaskResult(){return {status:"SUCCESS"};}
|
||||
`
|
||||
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
request := map[string]any{"units": int64(2), "nested": map[string]any{"label": "original"}}
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/submit", nil)
|
||||
c.Set("task_request", request)
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "copy", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "copy", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
facts, err := adaptor.ExtractUsageFactsValidated(c, info)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]any{"units": float64(2)}, facts)
|
||||
body, err := adaptor.BuildRequestBody(c, info)
|
||||
require.NoError(t, err)
|
||||
encoded, err := io.ReadAll(body)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, `{"units":3,"nested":{"label":"built"}}`, string(encoded))
|
||||
assert.Equal(t, int64(2), request["units"])
|
||||
assert.Equal(t, "original", request["nested"].(map[string]any)["label"])
|
||||
}
|
||||
|
||||
func TestTaskSubmitDeltaStreamContract(t *testing.T) {
|
||||
source := strings.Replace(documentStreamPlugin, `submitResponseTypes:["json","sse"],`, `submitResponseTypes:["json","sse"],requiredCapabilities:["submit-sse-delta@1"],`, 1) + `
|
||||
export function parseSubmitEventDelta(ctx,event,previous) {
|
||||
if(previous && previous.document !== undefined) throw new Error("full result leaked into control state");
|
||||
const chunk=JSON.parse(event.data);
|
||||
let changes=chunk.changes;
|
||||
if(chunk.largeResult) changes=[{op:"set",path:[],value:{document:"<".repeat(200000)}}];
|
||||
if(chunk.resetAfterLarge) changes.push({op:"set",path:[],value:{document:"",units:0}});
|
||||
const result={changes:changes,state:{events:(previous ? previous.events : 0)+1},done:chunk.complete===true};
|
||||
if(chunk.largeControl) result.state={text:"x".repeat(65537)};
|
||||
if(chunk.extra) result.extra=true;
|
||||
return result;
|
||||
}`
|
||||
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
const first = `{"changes":[{"op":"set","path":[],"value":{"document":"hello","units":2}}]}`
|
||||
const last = `{"changes":[{"op":"appendText","path":["document"],"value":"world"},{"op":"set","path":["units"],"value":0}],"complete":true}`
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
frames []string
|
||||
valid bool
|
||||
}{
|
||||
{"control state and zero usage", []string{first, last}, true},
|
||||
{"oversized intermediate result", []string{`{"largeResult":true}`, last}, false},
|
||||
{"oversized operation before reset", []string{`{"largeResult":true,"resetAfterLarge":true,"complete":true}`}, false},
|
||||
{"control state limit", []string{first, `{"changes":[],"largeControl":true,"complete":true}`}, false},
|
||||
{"missing changes", []string{`{"complete":true}`}, false},
|
||||
{"extra result fields", []string{first, `{"changes":[],"extra":true,"complete":true}`}, false},
|
||||
{"unfinished delta stream", []string{first}, false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var stream strings.Builder
|
||||
for _, frame := range tc.frames {
|
||||
stream.WriteString("data: " + frame + "\n\n")
|
||||
}
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "alias", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "document", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "public-document"}}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/compile", nil)
|
||||
c.Set("task_request", map[string]any{})
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
response := &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(stream.String()))}
|
||||
parsed, taskErr := adaptor.ParseResponse(c, response, info)
|
||||
assert.False(t, c.Writer.Written())
|
||||
if !tc.valid {
|
||||
require.NotNil(t, taskErr)
|
||||
assert.True(t, taskErr.NoRetry)
|
||||
return
|
||||
}
|
||||
require.Nil(t, taskErr)
|
||||
require.NotNil(t, parsed.Immediate)
|
||||
assert.JSONEq(t, `{"document":"helloworld","units":0}`, string(parsed.TaskData))
|
||||
assert.Equal(t, map[string]any{"units": float64(0)}, parsed.Immediate.UsageFacts)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAlibabaSubmitDeltaDoesNotMutateControlState(t *testing.T) {
|
||||
source, err := plugins.Source("alibaba")
|
||||
require.NoError(t, err)
|
||||
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
const before = `{"request_id":"old","usage":{"image_count":1},"output":{"finished":false,"choices":[{"message":{"role":"assistant","content":[{"text":"hello"}]}}]}}`
|
||||
const control = `{"choices":[{"count":1,"lastText":true,"finishReason":""}],"hasUsage":true}`
|
||||
var previous, initial any
|
||||
require.NoError(t, common.UnmarshalJsonStr(control, &previous))
|
||||
require.NoError(t, common.UnmarshalJsonStr(before, &initial))
|
||||
accumulated := pluginruntime.NewJSONState(maxTaskPluginPersistedJSONBytes)
|
||||
require.NoError(t, accumulated.Apply(t.Context(), []any{map[string]any{"op": "set", "path": []any{}, "value": initial}}))
|
||||
value, err := plugin.Engine.Call(t.Context(), "parseSubmitEventDelta", map[string]any{}, map[string]any{
|
||||
"event": "message", "data": `{"request_id":"new","usage":{"image_count":2},"output":{"choices":[{"finish_reason":"stop","message":{"content":[{"text":" world"},{"image":"https://cdn.example/image.png"}]}}]}}`,
|
||||
}, previous)
|
||||
require.NoError(t, err)
|
||||
unchanged, err := common.Marshal(previous)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, control, string(unchanged))
|
||||
result := value.(map[string]any)
|
||||
assert.Equal(t, true, result["done"])
|
||||
require.NoError(t, accumulated.Apply(t.Context(), result["changes"]))
|
||||
final, err := accumulated.Value()
|
||||
require.NoError(t, err)
|
||||
encoded, err := common.Marshal(final)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, `{"request_id":"new","usage":{"image_count":2},"output":{"finished":true,"choices":[{"finish_reason":"stop","message":{"role":"assistant","content":[{"text":"hello world"},{"image":"https://cdn.example/image.png"}]}}]}}`, string(encoded))
|
||||
encoded, err = common.Marshal(result["state"])
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, `{"choices":[{"count":2,"lastText":false,"finishReason":"stop"}],"hasUsage":true}`, string(encoded))
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/plugins"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func benchmarkAlibaba(b *testing.B) (*pluginruntime.LoadedPlugin, *atomic.Int64) {
|
||||
b.Helper()
|
||||
source, err := plugins.Source("alibaba")
|
||||
require.NoError(b, err)
|
||||
initializations := &atomic.Int64{}
|
||||
plugin, err := pluginruntime.CompilePlugin(source+"\n"+`console.log("benchmark_runtime_initialized");`, pluginruntime.Options{
|
||||
Log: func(message string) {
|
||||
if strings.HasSuffix(message, "benchmark_runtime_initialized") {
|
||||
initializations.Add(1)
|
||||
}
|
||||
},
|
||||
})
|
||||
require.NoError(b, err)
|
||||
return plugin, initializations
|
||||
}
|
||||
|
||||
func BenchmarkTaskPluginRuntime(b *testing.B) {
|
||||
for _, name := range []string{"HasExport", "MissingExport", "HookCall", "ParallelHookCall"} {
|
||||
b.Run(name, func(b *testing.B) {
|
||||
plugin, initializations := benchmarkAlibaba(b)
|
||||
ctx := context.Background()
|
||||
input := map[string]any{"model": "wan2.5-i2v-preview", "upstreamModel": "wan2.5-i2v-preview", "requestBody": map[string]any{
|
||||
"prompt": "animate the landscape", "image": "https://cdn.example/image.png", "duration": 5, "size": "720P",
|
||||
}}
|
||||
before := initializations.Load()
|
||||
b.ReportAllocs()
|
||||
if name == "ParallelHookCall" {
|
||||
b.ResetTimer()
|
||||
b.RunParallel(func(pb *testing.PB) {
|
||||
for pb.Next() {
|
||||
if _, err := plugin.Engine.Call(ctx, "extractUsage", input); err != nil {
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
} else {
|
||||
for b.Loop() {
|
||||
var err error
|
||||
switch name {
|
||||
case "HasExport":
|
||||
_, err = plugin.Engine.HasExport(ctx, "extractUsage")
|
||||
case "MissingExport":
|
||||
_, err = plugin.Engine.HasExport(ctx, "absentHook")
|
||||
default:
|
||||
_, err = plugin.Engine.Call(ctx, "extractUsage", input)
|
||||
}
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Run with -benchtime=50x: GC is excluded so this measures the first hook after
|
||||
// reclamation, not the runtime.GC operation or an artificial throughput workload.
|
||||
func BenchmarkTaskPluginAfterGC(b *testing.B) {
|
||||
plugin, initializations := benchmarkAlibaba(b)
|
||||
input := map[string]any{"model": "wan2.5-i2v-preview", "upstreamModel": "wan2.5-i2v-preview", "requestBody": map[string]any{
|
||||
"prompt": "animate the landscape", "image": "https://cdn.example/image.png", "duration": 5, "size": "720P",
|
||||
}}
|
||||
before := initializations.Load()
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for range b.N {
|
||||
b.StopTimer()
|
||||
runtime.GC()
|
||||
runtime.GC()
|
||||
b.StartTimer()
|
||||
if _, err := plugin.Engine.Call(b.Context(), "extractUsage", input); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
|
||||
}
|
||||
|
||||
func BenchmarkTaskPluginSubmit(b *testing.B) {
|
||||
for _, size := range []int{4 << 10, 1 << 20, 8 << 20} {
|
||||
b.Run(fmt.Sprintf("%dKiB", size>>10), func(b *testing.B) {
|
||||
plugin, initializations := benchmarkAlibaba(b)
|
||||
const modelName = "wan2.5-i2v-preview"
|
||||
wire, err := common.Marshal(map[string]any{
|
||||
"model": modelName, "duration": 5, "size": "720P",
|
||||
"input": []any{map[string]any{"type": "input_text", "text": "animate the landscape"}, map[string]any{"type": "input_image", "image_url": "data:image/png;base64," + strings.Repeat("A", size)}},
|
||||
})
|
||||
require.NoError(b, err)
|
||||
before := initializations.Load()
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for b.Loop() {
|
||||
var body any
|
||||
if err := common.Unmarshal(wire, &body); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
protocol := pluginruntime.ProtocolRequestContext{
|
||||
RouteRequestContext: pluginruntime.RouteRequestContext{Path: "/v1/responses", Method: http.MethodPost, Body: map[string]any{"kind": "json", "value": body}},
|
||||
Protocol: "openai_responses", Operation: "create", Model: modelName,
|
||||
}
|
||||
// The middleware decoder runs before channel selection; the adaptor
|
||||
// decodes again after the executing provider has been pinned.
|
||||
decoded, err := plugin.Engine.CallPath(b.Context(), "protocols", []string{"openai_responses", "decodeRequest"}, protocol.JSValue())
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set("task_request", decoded.(map[string]any)["requestBody"])
|
||||
c.Set(pluginruntime.ContextKeyRouteRequest, protocol.RouteRequestContext)
|
||||
c.Set(pluginruntime.ContextKeyProtocolRequest, protocol)
|
||||
c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{Plugin: plugin, Protocol: protocol.Protocol, Model: modelName})
|
||||
info := &relaycommon.RelayInfo{OriginModelName: modelName, ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: modelName, ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "benchmark-task"}}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
if taskErr := adaptor.ValidateRequestAndSetAction(c, info); taskErr != nil {
|
||||
b.Fatal(taskErr)
|
||||
}
|
||||
if _, err := adaptor.ExtractUsageFactsValidated(c, info); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
upstream, err := adaptor.BuildRequestBody(c, info)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
if _, err := io.Copy(io.Discard, upstream); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkTaskPluginSSE(b *testing.B) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
events, textBytes int
|
||||
}{{"small_state", 32, 48}, {"growing_state", 64, 8 << 10}} {
|
||||
b.Run(tc.name, func(b *testing.B) {
|
||||
plugin, initializations := benchmarkAlibaba(b)
|
||||
var stream strings.Builder
|
||||
for index := range tc.events {
|
||||
choice := map[string]any{"message": map[string]any{"role": "assistant", "content": []any{map[string]any{"text": strings.Repeat("a", tc.textBytes)}}}}
|
||||
if index == tc.events-1 {
|
||||
choice["finish_reason"] = "stop"
|
||||
}
|
||||
frame, err := common.Marshal(map[string]any{"request_id": "benchmark-sse", "output": map[string]any{"choices": []any{choice}}, "usage": map[string]any{"image_count": 1}})
|
||||
require.NoError(b, err)
|
||||
stream.WriteString("data: ")
|
||||
stream.Write(frame)
|
||||
stream.WriteString("\n\n")
|
||||
}
|
||||
wire := stream.String()
|
||||
input := map[string]any{"model": "wan2.6-image"}
|
||||
before := initializations.Load()
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for b.Loop() {
|
||||
response := &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(wire))}
|
||||
if _, err := New(plugin).readSubmitEvents(b.Context(), response, input); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"mime"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
)
|
||||
|
||||
// readSubmitEvents owns only SSE framing and bounded JSON state. The plugin
|
||||
// interprets data, accumulates provider output and decides when it is complete.
|
||||
// No response bytes are written to the client at this boundary.
|
||||
func (a *TaskAdaptor) readSubmitEvents(parent context.Context, resp *http.Response, driverContext map[string]any) (any, error) {
|
||||
mediaType, _, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
|
||||
if err != nil || mediaType != "text/event-stream" {
|
||||
return nil, fmt.Errorf("expected a text/event-stream submit response")
|
||||
}
|
||||
ctx, cancel := context.WithCancelCause(parent)
|
||||
defer cancel(nil)
|
||||
stopClose := context.AfterFunc(ctx, func() { _ = resp.Body.Close() })
|
||||
defer stopClose()
|
||||
idleTimeout := time.Duration(constant.StreamingTimeout) * time.Second
|
||||
if idleTimeout <= 0 {
|
||||
idleTimeout = 30 * time.Second
|
||||
}
|
||||
timer := time.AfterFunc(idleTimeout, func() { cancel(fmt.Errorf("task submit stream idle timeout")) })
|
||||
defer timer.Stop()
|
||||
|
||||
scanner := helper.NewStreamScanner(resp.Body, maxTaskPluginPersistedJSONBytes+1)
|
||||
hook := "parseSubmitEvent"
|
||||
var accumulated *pluginruntime.JSONState
|
||||
if slices.Contains(a.plugin.Meta.RequiredCapabilities, pluginruntime.CapabilitySubmitSSEDelta) {
|
||||
hook = "parseSubmitEventDelta"
|
||||
accumulated = pluginruntime.NewJSONState(maxTaskPluginPersistedJSONBytes)
|
||||
}
|
||||
var state any
|
||||
var data []string
|
||||
var eventName, lastID string
|
||||
frameBytes := 0
|
||||
firstLine := true
|
||||
for scanner.Scan() {
|
||||
if err := context.Cause(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
timer.Reset(idleTimeout)
|
||||
line := scanner.Text()
|
||||
if firstLine {
|
||||
line = strings.TrimPrefix(line, "\ufeff")
|
||||
firstLine = false
|
||||
}
|
||||
frameBytes += len(line) + 1
|
||||
if frameBytes > maxTaskPluginPersistedJSONBytes {
|
||||
return nil, fmt.Errorf("task submit SSE event exceeds size limit")
|
||||
}
|
||||
if line != "" {
|
||||
field, value, _ := strings.Cut(line, ":")
|
||||
value, _ = strings.CutPrefix(value, " ")
|
||||
switch field {
|
||||
case "data":
|
||||
data = append(data, value)
|
||||
case "event":
|
||||
eventName = value
|
||||
case "id":
|
||||
if !strings.ContainsRune(value, '\x00') {
|
||||
lastID = value
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
frameBytes = 0
|
||||
if len(data) == 0 {
|
||||
eventName = ""
|
||||
continue
|
||||
}
|
||||
if eventName == "" {
|
||||
eventName = "message"
|
||||
}
|
||||
event := map[string]any{"event": eventName, "id": lastID, "data": strings.Join(data, "\n")}
|
||||
value, err := a.plugin.Engine.Call(ctx, hook, driverContext, event, state)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result, ok := value.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%s must return an object", hook)
|
||||
}
|
||||
nextState, hasState := result["state"]
|
||||
done, hasDone := result["done"].(bool)
|
||||
if !hasState || !hasDone {
|
||||
return nil, fmt.Errorf("%s must return state and a boolean done", hook)
|
||||
}
|
||||
encoded, err := common.Marshal(nextState)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid submit stream state: %w", err)
|
||||
}
|
||||
if len(encoded) > maxTaskPluginPersistedJSONBytes {
|
||||
return nil, fmt.Errorf("task submit stream state exceeds size limit")
|
||||
}
|
||||
if accumulated != nil {
|
||||
if len(result) != 3 {
|
||||
return nil, fmt.Errorf("parseSubmitEventDelta must return only changes, state and done")
|
||||
}
|
||||
// Control state is separate from the full result and is the only
|
||||
// snapshot passed back to JS on the next event.
|
||||
if len(encoded) > 64<<10 {
|
||||
return nil, fmt.Errorf("task submit stream control state exceeds size limit")
|
||||
}
|
||||
if err = accumulated.Apply(ctx, result["changes"]); err != nil {
|
||||
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 {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if done {
|
||||
if accumulated != nil {
|
||||
return accumulated.Value()
|
||||
}
|
||||
return state, nil
|
||||
}
|
||||
data = nil
|
||||
eventName = ""
|
||||
}
|
||||
if err := context.Cause(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, fmt.Errorf("read task submit stream: %w", err)
|
||||
}
|
||||
return nil, fmt.Errorf("task submit stream ended before the plugin reported completion")
|
||||
}
|
||||
Reference in new issue
Block a user