feat: publish sanitized custom.97 source
This commit is contained in:
1022 files changed
+69903
-19207
No files matched your search
@@ -0,0 +1,260 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// serveTaskPluginImageProtocol answers an OpenAI Images API request that a
|
||||
// task plugin claimed. The API is synchronous: the durable submission runs
|
||||
// first on a context that survives a client disconnect, an asynchronous
|
||||
// upstream task is then polled inside the request until it is terminal or the
|
||||
// protocol timeout elapses, and the plugin's render hook finally shapes the
|
||||
// OpenAI image response from the terminal task.
|
||||
func serveTaskPluginImageProtocol(c *gin.Context, pinned pluginruntime.PinnedEndpoint, deps pluginProtocolBridgeDeps) {
|
||||
deps = deps.withDefaults()
|
||||
requestValue, exists := c.Get(pluginruntime.ContextKeyProtocolRequest)
|
||||
protocolRequest, ok := requestValue.(pluginruntime.ProtocolRequestContext)
|
||||
if !exists || !ok || protocolRequest.Protocol != pinned.Protocol || pinned.Plugin == nil {
|
||||
respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed")
|
||||
return
|
||||
}
|
||||
clientRequest := c.Request
|
||||
var relayInfo *relaycommon.RelayInfo
|
||||
var relayInfoErr error
|
||||
var outcome *taskSubmissionOutcome
|
||||
var taskErr *dto.TaskError
|
||||
// DashScope bills the images as soon as it accepts the upstream call. A
|
||||
// client that gives up while the submission is in flight must not cancel
|
||||
// that call, the task row, its settlement, or its consume log, so the
|
||||
// submission runs on its own bounded context exactly like the Responses
|
||||
// bridge. Only the in-request wait and the response write observe the
|
||||
// client connection.
|
||||
func() {
|
||||
submissionContext, cancelSubmission := context.WithTimeout(
|
||||
context.WithoutCancel(clientRequest.Context()),
|
||||
deps.submissionTimeout,
|
||||
)
|
||||
c.Request = clientRequest.Clone(submissionContext)
|
||||
defer func() {
|
||||
c.Request = clientRequest
|
||||
cancelSubmission()
|
||||
}()
|
||||
|
||||
relayInfo, relayInfoErr = relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
|
||||
if relayInfoErr != nil {
|
||||
return
|
||||
}
|
||||
relayInfo.IsStream = false
|
||||
relayInfo.OriginModelName = c.GetString("resolved_task_model")
|
||||
if action := c.GetString("task_action"); action != "" {
|
||||
relayInfo.Action = action
|
||||
}
|
||||
if taskErr = relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
|
||||
return
|
||||
}
|
||||
if taskErr = relay.ApplyOriginTaskAffinity(c, relayInfo); taskErr != nil {
|
||||
return
|
||||
}
|
||||
outcome, taskErr = deps.submit(c, relayInfo)
|
||||
}()
|
||||
if clientRequest.Context().Err() != nil {
|
||||
// The client is gone. A durable task keeps its settlement and consume
|
||||
// log; an asynchronous one is finished by the background poller.
|
||||
return
|
||||
}
|
||||
if relayInfoErr != nil {
|
||||
logger.LogError(c, common.LogText("build task protocol relay info failed: %s", relayInfoErr.Error()))
|
||||
respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed")
|
||||
return
|
||||
}
|
||||
if taskErr != nil {
|
||||
respondTaskPluginImageError(c, taskErr)
|
||||
return
|
||||
}
|
||||
if outcome == nil || outcome.Task == nil || outcome.Task.Platform != constant.TaskPlatform(pinned.Plugin.Meta.Key) {
|
||||
logger.LogError(c, common.LogText("task protocol submission returned an invalid durable outcome"))
|
||||
respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed")
|
||||
return
|
||||
}
|
||||
task := outcome.Task
|
||||
if task.Status != model.TaskStatusSuccess && task.Status != model.TaskStatusFailure {
|
||||
if taskErr = waitTaskPluginImageTask(c, task, deps); taskErr != nil {
|
||||
if c.Request.Context().Err() == nil {
|
||||
respondTaskPluginImageError(c, taskErr)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
if task.Status == model.TaskStatusFailure {
|
||||
reason := strings.TrimSpace(task.FailReason)
|
||||
if reason == "" {
|
||||
reason = "image generation failed"
|
||||
}
|
||||
// The fail reason comes from the upstream, so mask the channel address in it.
|
||||
taskErr = service.TaskErrorWrapper(errors.New(reason), "image_generation_failed", http.StatusBadRequest, common.UpstreamAddresses(c)...)
|
||||
taskErr.LocalError = true
|
||||
respondTaskPluginImageError(c, taskErr)
|
||||
return
|
||||
}
|
||||
view, err := service.BuildTaskPluginView(task)
|
||||
if err != nil {
|
||||
logger.LogError(c, common.LogText("build task protocol view failed: %s", err.Error()))
|
||||
respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed")
|
||||
return
|
||||
}
|
||||
viewValue, err := taskPluginProtocolJSONValue(view)
|
||||
if err != nil {
|
||||
logger.LogError(c, common.LogText("encode task protocol view failed: %s", err.Error()))
|
||||
respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed")
|
||||
return
|
||||
}
|
||||
payload, err := pinned.Plugin.Engine.CallPathWithAdmissionTimeout(c.Request.Context(), deps.admissionTimeout, "protocols", []string{pinned.Protocol, "render"}, protocolRequest.JSValue(), viewValue)
|
||||
if err != nil {
|
||||
logger.LogError(c, common.LogText("task protocol image render hook failed: %s", err.Error()))
|
||||
respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed")
|
||||
return
|
||||
}
|
||||
response, ok := payload.(map[string]any)
|
||||
data, hasData := response["data"].([]any)
|
||||
if !ok || !hasData {
|
||||
logger.LogError(c, common.LogText("task protocol image render hook must return an object with a data array"))
|
||||
respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed")
|
||||
return
|
||||
}
|
||||
if _, present := response["created"]; !present {
|
||||
createdAt := task.CreatedAt
|
||||
if createdAt == 0 {
|
||||
createdAt = task.SubmitTime
|
||||
}
|
||||
response["created"] = createdAt
|
||||
}
|
||||
// response_format is host-owned: the plugin renders upstream URLs and the
|
||||
// host inlines each image when the client asked for Base64.
|
||||
if taskPluginImageResponseFormat(protocolRequest) == "b64_json" {
|
||||
for _, entry := range data {
|
||||
item, isObject := entry.(map[string]any)
|
||||
if !isObject {
|
||||
continue
|
||||
}
|
||||
url, _ := item["url"].(string)
|
||||
if existing, _ := item["b64_json"].(string); strings.TrimSpace(url) == "" || existing != "" {
|
||||
continue
|
||||
}
|
||||
_, encoded, downloadErr := deps.downloadImage(url)
|
||||
if downloadErr != nil {
|
||||
logger.LogWarn(c, common.LogText("task protocol image download failed; returning the upstream URL instead of b64_json: %s", common.MaskSensitiveInfo(downloadErr.Error(), common.UpstreamAddresses(c)...)))
|
||||
continue
|
||||
}
|
||||
item["b64_json"] = encoded
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusOK, response)
|
||||
}
|
||||
|
||||
// waitTaskPluginImageTask polls an asynchronous image task inside the client
|
||||
// request through the ordinary task polling path, so completion settles
|
||||
// billing exactly as the background poller would, until the task is terminal
|
||||
// or the protocol timeout elapses. A client disconnect stops waiting; the
|
||||
// background poller then finishes the durable task.
|
||||
func waitTaskPluginImageTask(c *gin.Context, task *model.Task, deps pluginProtocolBridgeDeps) *dto.TaskError {
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), deps.submissionTimeout)
|
||||
defer cancel()
|
||||
for {
|
||||
timer := time.NewTimer(deps.imagePollInterval)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
if errors.Is(ctx.Err(), context.DeadlineExceeded) && c.Request.Context().Err() == nil {
|
||||
return service.TaskErrorWrapperLocal(fmt.Errorf("image generation is still running; task %s settles when it completes", task.TaskID), "task_timeout", http.StatusGatewayTimeout)
|
||||
}
|
||||
return service.TaskErrorWrapperLocal(ctx.Err(), "request_cancelled", http.StatusRequestTimeout)
|
||||
case <-timer.C:
|
||||
}
|
||||
if err := deps.pollTask(ctx, task); err != nil {
|
||||
logger.LogWarn(c, common.LogText("task protocol image poll failed for task %s: %s", task.TaskID, common.MaskSensitiveInfo(err.Error(), common.UpstreamAddresses(c)...)))
|
||||
continue
|
||||
}
|
||||
if task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pollTaskPluginImageTask runs one poll round for a single task through the
|
||||
// same code path as the scheduled task poller, including status CAS and
|
||||
// completion settlement.
|
||||
func pollTaskPluginImageTask(ctx context.Context, task *model.Task) error {
|
||||
upstreamID := task.GetUpstreamTaskID()
|
||||
if upstreamID == "" {
|
||||
return errors.New("task has no upstream id")
|
||||
}
|
||||
return service.UpdateVideoTasks(ctx, task.Platform, map[int][]string{task.ChannelId: {upstreamID}}, map[string]*model.Task{upstreamID: task})
|
||||
}
|
||||
|
||||
func taskPluginImageResponseFormat(request pluginruntime.ProtocolRequestContext) string {
|
||||
body, ok := request.Body.(map[string]any)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
switch body["kind"] {
|
||||
case string(pluginruntime.BodyJSON):
|
||||
value, _ := body["value"].(map[string]any)
|
||||
format, _ := value["response_format"].(string)
|
||||
return strings.TrimSpace(format)
|
||||
case string(pluginruntime.BodyMultipart), string(pluginruntime.BodyForm):
|
||||
fields, _ := body["fields"].(map[string][]string)
|
||||
if values := fields["response_format"]; len(values) > 0 {
|
||||
return strings.TrimSpace(values[0])
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// respondTaskPluginImageError writes an OpenAI error envelope for the Images
|
||||
// API. DashScope failure bodies arrive as JSON {code, message}; their message
|
||||
// is surfaced so a client sees the vendor reason instead of raw JSON.
|
||||
func respondTaskPluginImageError(c *gin.Context, taskErr *dto.TaskError) {
|
||||
service.RecordRequestPolicyTermination(c, taskSubmissionAPIError(taskErr))
|
||||
status := http.StatusInternalServerError
|
||||
if taskErr.StatusCode >= 400 && taskErr.StatusCode <= 599 {
|
||||
status = taskErr.StatusCode
|
||||
}
|
||||
code := taskErr.Code
|
||||
message := taskErr.Message
|
||||
var upstream struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
if common.UnmarshalJsonStr(message, &upstream) == nil && strings.TrimSpace(upstream.Message) != "" {
|
||||
message = upstream.Message
|
||||
if strings.TrimSpace(upstream.Code) != "" {
|
||||
code = upstream.Code
|
||||
}
|
||||
}
|
||||
errorType := "new_api_error"
|
||||
if status < http.StatusInternalServerError {
|
||||
errorType = "invalid_request_error"
|
||||
}
|
||||
c.JSON(status, gin.H{"error": gin.H{
|
||||
"message": common.MessageWithRequestId(message, c.GetString(common.RequestIdKey)),
|
||||
"type": errorType,
|
||||
"code": code,
|
||||
}})
|
||||
}
|
||||
Reference in new issue
Block a user