Skip to content

phsl / new-api

Go to a project
Commit 9df450fe by CaIon
Options

feat(task): give polling hooks a real query context, host HTTP classification,… 

feat(task): give polling hooks a real query context, host HTTP classification, and bounded poll failures

Plugin polling hooks previously ran against a hollow context: parseTaskResult
and parseBatchResult received {} / nil, buildQueryRequest received a
{task_id, action} map under the misleading name requestBody, and batch hooks
saw only bare task ids. The per-task poller also never looked at the upstream
HTTP status, and every built-in plugin papered over unrecognized bodies with
`|| "IN_PROGRESS"`, so a 404, a revoked key, or a shape the plugin did not
know would sit in IN_PROGRESS for the full 24h TASK_TIMEOUT_MINUTES while
holding the user's pre-charged quota.

Contract (docs/plugin-api v1.d.ts, v1.md, v1.schema.json):
- TaskQueryContext is declared separately from DriverContext and rebuilt from
  the persisted Task row: taskId, publicTaskId, action, model, upstreamModel,
  baseUrl, apiKey, authHeader, auth, data, state. Query-side requestBody is
  removed; the original request is not persisted and hooks that need a
  request-derived value must save it into state at submit time.
- parseTaskResult / parseBatchResult receive a third {status, headers}
  argument. Batch hooks receive tasks[] with one TaskQueryContext per task.
- NormalizedTaskResult accepts status "UNKNOWN" meaning "I do not recognize
  this body". Falling back to IN_PROGRESS for unknown shapes is forbidden;
  `plugin lint` warns on the literal.
- parseSubmitResponse / parseTaskResult / parseBatchResult may return `state`.
  Task.Data remains a per-round snapshot overwritten on every valid parse;
  state is plugin-owned, persisted in TaskPrivateData.PluginState, preserved
  when a hook omits it, byte-capped like taskData, and never exposed through
  presenter views.

Host (service/task_polling.go, relay/channel/task/jsplugin/adaptor.go):
- TaskPollingAdaptor / BatchTaskPollingAdaptor take *model.Task and the
  *http.Response so the adaptor can build the full context; jsplugin is the
  only implementation.
- HTTP classification before the plugin sees the body: 2xx -> plugin;
  404/410 -> FAILURE and refund; 401/403 -> poll failure plus a channel-scoped
  warning, no auto-disable; 429/5xx/transport -> poll failure; other 4xx ->
  plugin with the status visible, counted as unrecognized if the plugin still
  reports a non-terminal state.
- TaskPrivateData.PollFailures counts consecutive poll failures (transient
  HTTP, auth, transport, hook error, UNKNOWN). It is persisted through the
  existing UpdateWithStatus CAS so a concurrent terminal transition on another
  instance is never clobbered, and reset on any valid 2xx non-terminal parse.
  Reaching TASK_POLL_MAX_FAILURES (default 20, <= 0 disables) fails the task
  with the last classification and HTTP code in fail_reason and runs the
  existing settle/refund chain exactly once. sweepTimedOutTasks and its
  1440-minute default are unchanged as the outer backstop.
- Unrecognized bodies are logged at WARN with a bounded redacted copy since
  Task.Data is intentionally not overwritten on that path.

Plugins (all ten bumped one patch version):
- jimeng persists the outbound req_key in state and reads it back in
  buildQueryRequest, replacing dead reads of ctx.data / ctx.requestBody that
  never resolved.
- sunoapi batch hooks read tasks[] instead of the removed requestBody.
- hailuo treats base_resp.status_code != 0 as FAILURE before the status table.
- kling, vidu, sora, alibaba, doubao, hailuo, jimeng return UNKNOWN with the
  raw upstream status in reason on table miss.
- google and vertex-ai treat a missing `done` as in-progress: Google
  long-running operations omit proto3 default fields, so a running Veo
  operation has no `done` key at all. Only a body without an operation name is
  UNKNOWN. plugins/veo_poll_test.go locks this so the poll-failure cutoff can
  never fail a rendering Veo task.

Tests cover the classification table end to end against a real DB (404
immediate refund, 429xN refund, 401 increments without status change, 2xx
reset, UNKNOWN increments, state preserved vs replaced, PollFailures survives
the CAS write), the query-context shape, UNKNOWN on unrecognized bodies, and
the absence of PluginState/PollFailures from TaskView. Controller tests derive
the kling factory version from the embedded manifest instead of hardcoding it.
parent 9f506dd7
Showing with 1319 additions and 181 deletions
.env.example
...@@ -56,6 +56,10 @@ ...@@ -56,6 +56,10 @@
# 任务和功能配置 # 任务和功能配置
# 更新任务启用 # 更新任务启用
# UPDATE_TASK=true # UPDATE_TASK=true
# 异步任务硬超时（分钟），按提交时间计算，超时未完成的任务标记失败并退款；0 表示禁用
# TASK_TIMEOUT_MINUTES=1440
# 异步任务连续轮询失败阈值（上游 429/5xx/401/403、网络错误、无法识别的响应），达到后任务标记失败并退款；正常轮询成功一次即归零
# TASK_POLL_MAX_FAILURES=20
# 对话超时设置 # 对话超时设置
# 所有请求超时时间，单位秒，默认为0，表示不限制 # 所有请求超时时间，单位秒，默认为0，表示不限制
......
common/init.go
...@@ -202,6 +202,8 @@ func initConstantEnv() { ...@@ -202,6 +202,8 @@ func initConstantEnv() {
constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000) constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000)
// 异步任务超时时间（分钟），超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。 // 异步任务超时时间（分钟），超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。
constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 1440) constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 1440)
// Consecutive unrecognized/transient poll failures before the task is failed and refunded.
constant.TaskPollMaxFailures = GetEnvOrDefault("TASK_POLL_MAX_FAILURES", 20)
// 声明式任务协议桥只观察数据库；这些值控制一次客户端观察连接， // 声明式任务协议桥只观察数据库；这些值控制一次客户端观察连接，
// 不改变后台轮询或结算生命周期。 // 不改变后台轮询或结算生命周期。
constant.TaskPluginProtocolTimeoutSeconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_TIMEOUT_SECONDS", 600) constant.TaskPluginProtocolTimeoutSeconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_TIMEOUT_SECONDS", 600)
......
constant/env.go
...@@ -18,6 +18,7 @@ var GenerateDefaultToken bool ...@@ -18,6 +18,7 @@ var GenerateDefaultToken bool
var ErrorLogEnabled bool var ErrorLogEnabled bool
var TaskQueryLimit int var TaskQueryLimit int
var TaskTimeoutMinutes int var TaskTimeoutMinutes int
var TaskPollMaxFailures = 20
var TaskPluginProtocolTimeoutSeconds int var TaskPluginProtocolTimeoutSeconds int
var TaskPluginProtocolTickMilliseconds int var TaskPluginProtocolTickMilliseconds int
var TaskPluginProtocolTickJitterMilliseconds int var TaskPluginProtocolTickJitterMilliseconds int
......
controller/plugin_protocol_test.go
...@@ -366,14 +366,14 @@ type terminalSettlementPollingAdaptor struct { ...@@ -366,14 +366,14 @@ type terminalSettlementPollingAdaptor struct {
func (a *terminalSettlementPollingAdaptor) Init(*relaycommon.RelayInfo) {} func (a *terminalSettlementPollingAdaptor) Init(*relaycommon.RelayInfo) {}
func (a *terminalSettlementPollingAdaptor) FetchTask(string, string, map[string]any, string) (*http.Response, error) { func (a *terminalSettlementPollingAdaptor) FetchTask(string, string, *model.Task, string) (*http.Response, error) {
return &http.Response{ return &http.Response{
StatusCode: http.StatusOK, StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(`{}`)), Body: io.NopCloser(strings.NewReader(`{}`)),
}, nil }, nil
} }
func (a *terminalSettlementPollingAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) { func (a *terminalSettlementPollingAdaptor) ParseTaskResult(*model.Task, *http.Response, []byte) (*relaycommon.TaskInfo, error) {
return &relaycommon.TaskInfo{ return &relaycommon.TaskInfo{
Status: model.TaskStatusSuccess, Status: model.TaskStatusSuccess,
Progress: "100%", Progress: "100%",
......
controller/relay.go
...@@ -744,6 +744,9 @@ func executeTaskSubmissionWith( ...@@ -744,6 +744,9 @@ func executeTaskSubmissionWith(
} }
task.Quota = result.Quota task.Quota = result.Quota
task.Data = result.TaskData task.Data = result.TaskData
if len(result.PluginState) > 0 {
task.PrivateData.PluginState = result.PluginState
}
task.Action = relayInfo.Action task.Action = relayInfo.Action
if immediate := result.Immediate; immediate != nil { if immediate := result.Immediate; immediate != nil {
task.Status = model.TaskStatus(immediate.Status) task.Status = model.TaskStatus(immediate.Status)
......
controller/task_plugin_test.go
...@@ -5,6 +5,7 @@ import ( ...@@ -5,6 +5,7 @@ import (
"fmt" "fmt"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"regexp"
"strings" "strings"
"testing" "testing"
...@@ -117,6 +118,17 @@ func TestDisableThirdPartyPluginSupportsCascadeAndForce(t *testing.T) { ...@@ -117,6 +118,17 @@ func TestDisableThirdPartyPluginSupportsCascadeAndForce(t *testing.T) {
assert.Equal(t, common.ChannelStatusManuallyDisabled, updated.Status) assert.Equal(t, common.ChannelStatusManuallyDisabled, updated.Status)
} }
// klingFactoryVersion returns the version declared in the embedded kling factory
// manifest so tests do not hardcode a value that moves with every plugin release.
func klingFactoryVersion(t *testing.T) string {
t.Helper()
factorySource, err := plugins.Source("kling")
require.NoError(t, err)
match := regexp.MustCompile(`version:\s*"([^"]+)"`).FindStringSubmatch(factorySource)
require.Len(t, match, 2, "kling factory manifest must declare a version")
return match[1]
}
func setupTaskPluginFactoryDisableTest(t *testing.T) { func setupTaskPluginFactoryDisableTest(t *testing.T) {
t.Helper() t.Helper()
setupTaskPluginControllerTest(t) setupTaskPluginControllerTest(t)
...@@ -240,7 +252,9 @@ func TestDisableFactoryOverrideRowKeepsEnabledFlagPath(t *testing.T) { ...@@ -240,7 +252,9 @@ func TestDisableFactoryOverrideRowKeepsEnabledFlagPath(t *testing.T) {
setupTaskPluginFactoryDisableTest(t) setupTaskPluginFactoryDisableTest(t)
factorySource, err := plugins.Source("kling") factorySource, err := plugins.Source("kling")
require.NoError(t, err) require.NoError(t, err)
overrideSource := strings.Replace(factorySource, `version: "1.0.0"`, `version: "1.0.0-test-factory-status"`, 1) factoryVersion := klingFactoryVersion(t)
overrideSource := strings.Replace(factorySource, `version: "`+factoryVersion+`"`, `version: "`+factoryVersion+`-test-factory-status"`, 1)
require.NotEqual(t, factorySource, overrideSource, "factory version marker must be found in kling source")
loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{}) loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{})
require.NoError(t, err) require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("kling") }) t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("kling") })
...@@ -266,7 +280,7 @@ func TestDisableFactoryOverrideRowKeepsEnabledFlagPath(t *testing.T) { ...@@ -266,7 +280,7 @@ func TestDisableFactoryOverrideRowKeepsEnabledFlagPath(t *testing.T) {
assert.True(t, taskPluginOptionsHasKey(t, "kling")) assert.True(t, taskPluginOptionsHasKey(t, "kling"))
got, ok := jsplugin.DefaultRegistry.Get("kling") got, ok := jsplugin.DefaultRegistry.Get("kling")
require.True(t, ok) require.True(t, ok)
assert.Equal(t, "1.0.0", got.Meta.Version) assert.Equal(t, factoryVersion, got.Meta.Version)
} }
func TestListTaskPluginsIncludesFactoryWithoutDatabaseRows(t *testing.T) { func TestListTaskPluginsIncludesFactoryWithoutDatabaseRows(t *testing.T) {
...@@ -358,7 +372,9 @@ func TestListTaskPluginsShowsDisabledFallbackWhenOverridesAreDisabled(t *testing ...@@ -358,7 +372,9 @@ func TestListTaskPluginsShowsDisabledFallbackWhenOverridesAreDisabled(t *testing
setupTaskPluginControllerTest(t) setupTaskPluginControllerTest(t)
factorySource, err := plugins.Source("kling") factorySource, err := plugins.Source("kling")
require.NoError(t, err) require.NoError(t, err)
overrideSource := strings.Replace(factorySource, `version: "1.0.0"`, `version: "1.0.0-test-disabled-override"`, 1) factoryVersion := klingFactoryVersion(t)
overrideSource := strings.Replace(factorySource, `version: "`+factoryVersion+`"`, `version: "`+factoryVersion+`-test-disabled-override"`, 1)
require.NotEqual(t, factorySource, overrideSource, "factory version marker must be found in kling source")
loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{}) loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{})
require.NoError(t, err) require.NoError(t, err)
plugin := model.TaskPlugin{ plugin := model.TaskPlugin{
...@@ -399,7 +415,9 @@ func TestDeleteActiveOverrideFallsBackToFactoryAndDeletesRecord(t *testing.T) { ...@@ -399,7 +415,9 @@ func TestDeleteActiveOverrideFallsBackToFactoryAndDeletesRecord(t *testing.T) {
setupTaskPluginControllerTest(t) setupTaskPluginControllerTest(t)
factorySource, err := plugins.Source("kling") factorySource, err := plugins.Source("kling")
require.NoError(t, err) require.NoError(t, err)
overrideSource := strings.Replace(factorySource, `version: "1.0.0"`, `version: "1.0.0-test-override"`, 1) factoryVersion := klingFactoryVersion(t)
overrideSource := strings.Replace(factorySource, `version: "`+factoryVersion+`"`, `version: "`+factoryVersion+`-test-override"`, 1)
require.NotEqual(t, factorySource, overrideSource, "factory version marker must be found in kling source")
loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{Key: "kling", Version: "test-override"}) loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{Key: "kling", Version: "test-override"})
require.NoError(t, err) require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("kling") }) t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("kling") })
......
docs/plugin-api/v1.d.ts
...@@ -25,6 +25,9 @@ export type UsageExample = {label: string; facts: Readonly<Record<string, string ...@@ -25,6 +25,9 @@ export type UsageExample = {label: string; facts: Readonly<Record<string, string
export interface Meta {apiVersion: 1; key: string; name: string; icon?: string; description?: LocalizedText; version: string; author: {name: string; url?: string}; channelTypes?: readonly number[]; models: readonly string[]; fetchMode: "per_task" | "batch"; allowedHosts?: readonly string[]; routes?: readonly NativeRoute[]; protocols?: readonly ProtocolClaim[]; usageSchema?: Readonly<Record<string, UsageFieldSchema>>; usageExamples?: readonly UsageExample[]; auth?: "none" | "api_key" | "vertex_oauth" | {type: "none" | "api_key" | "oauth2_jwt"}} export interface Meta {apiVersion: 1; key: string; name: string; icon?: string; description?: LocalizedText; version: string; author: {name: string; url?: string}; channelTypes?: readonly number[]; models: readonly string[]; fetchMode: "per_task" | "batch"; allowedHosts?: readonly string[]; routes?: readonly NativeRoute[]; protocols?: readonly ProtocolClaim[]; usageSchema?: Readonly<Record<string, UsageFieldSchema>>; usageExamples?: readonly UsageExample[]; auth?: "none" | "api_key" | "vertex_oauth" | {type: "none" | "api_key" | "oauth2_jwt"}}
export interface TaskView {task_id: string; status: string; progress?: string; fail_reason?: string; created_at?: number; updated_at?: number; data?: unknown; properties?: Record<string, unknown>} export interface TaskView {task_id: string; status: string; progress?: string; fail_reason?: string; created_at?: number; updated_at?: number; data?: unknown; properties?: Record<string, unknown>}
export interface DriverContext {requestBody: unknown; requestHeaders: Readonly<Record<string, string>>; action: string; model: string; upstreamModel: string; baseUrl: string; apiKey?: string; authHeader: string; files: readonly FileReference[]; publicTaskId: string; originTasks?: readonly {taskId: string; upstreamTaskId: string; action: string; status: string; data: unknown}[]} export interface DriverContext {requestBody: unknown; requestHeaders: Readonly<Record<string, string>>; action: string; model: string; upstreamModel: string; baseUrl: string; apiKey?: string; authHeader: string; files: readonly FileReference[]; publicTaskId: string; originTasks?: readonly {taskId: string; upstreamTaskId: string; action: string; status: string; data: unknown}[]}
export interface TaskQueryContext {taskId: string; publicTaskId: string; action: string; model: string; upstreamModel: string; baseUrl: string; apiKey?: string; authHeader: string; auth?: unknown; data: unknown; state: unknown}
export interface BatchQueryContext {baseUrl: string; apiKey?: string; authHeader: string; auth?: unknown; tasks: readonly TaskQueryContext[]}
export type HookHTTPResponse = {readonly status: number; readonly headers: Readonly<Record<string, string>>}
export interface RequestDescriptor {url: string; method?: string; headers?: Record<string, string>; /** JSON body may contain FilePlaceholder objects at any depth; the host replaces each with a Base64 or data-URL string. */ body?: unknown; credentialless?: boolean; action?: string; model?: string; rewriteModel?: string; bodyType?: "json" | "multipart"; parts?: readonly {name: string; value?: unknown; fileRef?: string; filename?: string}[]} export interface RequestDescriptor {url: string; method?: string; headers?: Record<string, string>; /** JSON body may contain FilePlaceholder objects at any depth; the host replaces each with a Base64 or data-URL string. */ body?: unknown; credentialless?: boolean; action?: string; model?: string; rewriteModel?: string; bodyType?: "json" | "multipart"; parts?: readonly {name: string; value?: unknown; fileRef?: string; filename?: string}[]}
export interface UpstreamResponse {statusCode: number; headers: Readonly<Record<string, readonly string[]>>; body: unknown} export interface UpstreamResponse {statusCode: number; headers: Readonly<Record<string, readonly string[]>>; body: unknown}
export interface NormalizedTaskResult {taskId?: string; status: "NOT_START" | "SUBMITTED" | "QUEUED" | "IN_PROGRESS" | "SUCCESS" | "FAILURE" | "UNKNOWN"; progress?: string; reason?: string; url?: string; remoteUrl?: string; completionTokens?: number; totalTokens?: number} export interface NormalizedTaskResult {taskId?: string; status: "NOT_START" | "SUBMITTED" | "QUEUED" | "IN_PROGRESS" | "SUCCESS" | "FAILURE" | "UNKNOWN"; progress?: string; reason?: string; url?: string; remoteUrl?: string; completionTokens?: number; totalTokens?: number}
...@@ -36,13 +39,13 @@ export declare const protocols: { ...@@ -36,13 +39,13 @@ export declare const protocols: {
openai_video?: {decodeRequest(ctx: ProtocolDecodeContext): SubmitIntent; render(ctx: unknown, task: TaskView): unknown}; openai_video?: {decodeRequest(ctx: ProtocolDecodeContext): SubmitIntent; render(ctx: unknown, task: TaskView): unknown};
}; };
export declare function buildSubmitRequest(ctx: DriverContext): RequestDescriptor; export declare function buildSubmitRequest(ctx: DriverContext): RequestDescriptor;
export declare function parseSubmitResponse(ctx: DriverContext, response: UpstreamResponse): {taskId: string; taskData?: unknown; immediate?: NormalizedTaskResult}; export declare function parseSubmitResponse(ctx: DriverContext, response: UpstreamResponse): {taskId: string; taskData?: unknown; immediate?: NormalizedTaskResult; state?: unknown};
export declare function buildQueryRequest(ctx: DriverContext & {taskId: string}): RequestDescriptor; export declare function buildQueryRequest(ctx: TaskQueryContext): RequestDescriptor;
export declare function buildBatchQueryRequest(ctx: DriverContext, taskIds: readonly string[]): RequestDescriptor; export declare function buildBatchQueryRequest(ctx: BatchQueryContext, tasks: readonly TaskQueryContext[]): RequestDescriptor;
export declare function parseTaskResult(ctx: DriverContext, body: unknown): NormalizedTaskResult; export declare function parseTaskResult(ctx: TaskQueryContext, body: unknown, response: HookHTTPResponse): NormalizedTaskResult;
export declare function parseBatchResult(ctx: DriverContext, body: unknown): readonly (NormalizedTaskResult & {taskId: string; data?: unknown})[]; export declare function parseBatchResult(ctx: BatchQueryContext, body: unknown, response: HookHTTPResponse): readonly (NormalizedTaskResult & {taskId: string; data?: unknown; state?: unknown})[];
export declare function extractUsage(ctx: DriverContext & {usagePurpose?: "facts" | "billing_ratios"}): Readonly<Record<string, string | number | boolean>> | null; export declare function extractUsage(ctx: DriverContext & {usagePurpose?: "facts" | "billing_ratios"}): Readonly<Record<string, string | number | boolean>> | null;
export declare function extractUsageOnSubmit(ctx: DriverContext, taskData: unknown): Readonly<Record<string, string | number | boolean>> | null; export declare function extractUsageOnSubmit(ctx: DriverContext, taskData: unknown): Readonly<Record<string, string | number | boolean>> | null;
export declare function extractUsageOnComplete(task: TaskView, result: NormalizedTaskResult, data: unknown): Readonly<Record<string, string | number | boolean>> | null; export declare function extractUsageOnComplete(task: TaskView, result: NormalizedTaskResult, data: unknown): Readonly<Record<string, string | number | boolean>> | null;
export declare function listArtifacts(task: {taskId: string; status: string; action: string; data: unknown; producerVersion: string}): readonly TaskArtifact[]; export declare function listArtifacts(task: {taskId: string; status: string; action: string; data: unknown; producerVersion: string}): readonly TaskArtifact[];
export declare function buildContentRequest(ctx: DriverContext & {artifactKey: string; data: unknown; upstreamTaskId: string; clientRequest: {method: "GET" | "HEAD"; headers: Readonly<Record<string, string>>}}): RequestDescriptor; export declare function buildContentRequest(ctx: DriverContext & {artifactKey: string; data: unknown; state?: unknown; upstreamTaskId: string; clientRequest: {method: "GET" | "HEAD"; headers: Readonly<Record<string, string>>}}): RequestDescriptor;
docs/plugin-api/v1.md
...@@ -32,7 +32,7 @@ Each `protocols` entry claims a host protocol. A protocol that defines modes mus ...@@ -32,7 +32,7 @@ Each `protocols` entry claims a host protocol. A protocol that defines modes mus
Enabled uploads pre-flight the candidate against the live routing generation and reject the first channel-type, native-route, or protocol-model conflict (the error names the counterpart plugin). Set `force: true` or `enabled: false` to store the plugin anyway. Enabled uploads pre-flight the candidate against the live routing generation and reject the first channel-type, native-route, or protocol-model conflict (the error names the counterpart plugin). Set `force: true` or `enabled: false` to store the plugin anyway.
`endpoints`, `routes[].renderer`, global `resolveRequest`, global `renderError`, and global `renderers` are rejected. `parseSubmitResponse` returns only `{taskId, taskData}` (plus the documented lifecycle fields); `clientResponse` is rejected. `endpoints`, `routes[].renderer`, global `resolveRequest`, global `renderError`, and global `renderers` are rejected. `parseSubmitResponse` returns only `{taskId, taskData, immediate?, state?}`; `clientResponse` is rejected.
`icon` is an optional LobeHub icon name string (for example `Sora.Color`). The values `text` and `text:<label>` request a generated text avatar instead (label defaults to the first two characters of `name`). It is display-only and does not participate in routing, billing, or admission beyond type and length checks. `icon` is an optional LobeHub icon name string (for example `Sora.Color`). The values `text` and `text:<label>` request a generated text avatar instead (label defaults to the first two characters of `name`). It is display-only and does not participate in routing, billing, or admission beyond type and length checks.
...@@ -140,3 +140,38 @@ Protocol media uses host-injected `ctx.artifacts[key].url`. Provider URLs from ` ...@@ -140,3 +140,38 @@ Protocol media uses host-injected `ctx.artifacts[key].url`. Provider URLs from `
The persisted field remains `task.data`; there is no `task.raw` alias. Driver hooks (`buildSubmitRequest`, `parseSubmitResponse`, query/result, usage, artifact, and content hooks) stay flat and must not branch on the client path or protocol. The persisted field remains `task.data`; there is no `task.raw` alias. Driver hooks (`buildSubmitRequest`, `parseSubmitResponse`, query/result, usage, artifact, and content hooks) stay flat and must not branch on the client path or protocol.
`ctx.model` is the billing and display identity (the origin name the client sent, including a channel-mapping alias). `ctx.upstreamModel` is the machine identity after channel `model_mapping`. Rate tables and model-keyed usage facts must use `ctx.upstreamModel || ctx.model`. Decode and render hooks that echo the client model must keep `ctx.model`. `buildSubmitRequest` must not set descriptor top-level `model` on a mapped pin; the host requires the plugin to echo the alias verbatim. Background polling has no relay info, so query hooks receive both identities from the persisted task properties, and `ctx.upstreamModel` falls back to `ctx.model` when the task was submitted without a channel mapping. `ctx.model` is the billing and display identity (the origin name the client sent, including a channel-mapping alias). `ctx.upstreamModel` is the machine identity after channel `model_mapping`. Rate tables and model-keyed usage facts must use `ctx.upstreamModel || ctx.model`. Decode and render hooks that echo the client model must keep `ctx.model`. `buildSubmitRequest` must not set descriptor top-level `model` on a mapped pin; the host requires the plugin to echo the alias verbatim. Background polling has no relay info, so query hooks receive both identities from the persisted task properties, and `ctx.upstreamModel` falls back to `ctx.model` when the task was submitted without a channel mapping.
## Polling contract
Query and parse hooks use `TaskQueryContext`, not `DriverContext`. The host rebuilds that context from the persisted task row. There is no query-side `requestBody`.
| Field | Source |
|-------|--------|
| `taskId` | Upstream task id (`PrivateData.UpstreamTaskID`, else `TaskID`) |
| `publicTaskId` | Gateway task id |
| `action` | Normalized persisted action |
| `model` | `Properties.OriginModelName` |
| `upstreamModel` | `Properties.UpstreamModelName`, falling back to `model` |
| `baseUrl` / `apiKey` / `authHeader` / `auth` | Channel credentials |
| `data` | Current `Task.Data` snapshot |
| `state` | Plugin-owned `PrivateData.PluginState` |
`Task.Data` is the latest upstream response snapshot for presenters and artifacts. The host overwrites it on every successful parse. Values that must survive across poll rounds belong in `state`.
`parseSubmitResponse`, `parseTaskResult`, and each `parseBatchResult` item may return optional `state`. The host writes it only when the hook returns it. Omitting `state` preserves the previous value. Oversized state is rejected with a warning, not truncated.
`buildBatchQueryRequest(ctx, tasks)` and `parseBatchResult` receive `tasks: TaskQueryContext[]`. `parseTaskResult` / `parseBatchResult` also receive `{status, headers}` for the upstream HTTP response.
`status: "UNKNOWN"` means the plugin does not recognize the response. Do not write `|| "IN_PROGRESS"` (or equivalent) for a missing table entry. The host treats `UNKNOWN`, hook errors, empty status, and unrecognized status strings as consecutive poll failures.
The host classifies the HTTP status before trusting a non-terminal parse:
| Upstream HTTP | Host action |
|---------------|-------------|
| 2xx | Call the parse hook |
| 404 / 410 | Immediate `FAILURE` and refund |
| 401 / 403 | Leave task status unchanged; increment `PollFailures`; `LogWarn` with channel id. Channels are not auto-disabled. |
| 429 / 5xx / transport error | Increment `PollFailures` |
| Other 4xx | Call the parse hook with `response.status`. A still-non-terminal result is unrecognized and increments `PollFailures`. |
A valid 2xx non-terminal parse resets `PollFailures` to 0. After `TASK_POLL_MAX_FAILURES` (default 20) consecutive failures the task becomes `FAILURE` and follows the existing refund chain. The 24h `TASK_TIMEOUT_MINUTES` sweep remains the outer deadline.
docs/plugin-api/v1.schema.json
...@@ -41,6 +41,45 @@ ...@@ -41,6 +41,45 @@
"maxBytes": {"type": "integer", "exclusiveMinimum": 0} "maxBytes": {"type": "integer", "exclusiveMinimum": 0}
} }
}, },
"route": {"type": "object", "additionalProperties": false, "required": ["method", "path", "type", "render"], "properties": {"method": {"enum": ["GET", "POST", "PUT", "PATCH", "DELETE"]}, "path": {"type": "string", "pattern": "^/"}, "type": {"enum": ["submit", "query", "dynamic"]}, "action": {"type": "string"}, "taskIdParam": {"type": "string"}, "decode": {"type": "string"}, "render": {"type": "string"}, "models": {"type": "array", "minItems": 1, "uniqueItems": true, "items": {"type": "string", "minLength": 1}}}, "allOf": [{"if": {"properties": {"type": {"const": "query"}}}, "then": {"allOf": [{"not": {"required": ["decode"]}}, {"not": {"required": ["models"]}}]}}, {"if": {"properties": {"type": {"enum": ["submit", "dynamic"]}}}, "then": {"required": ["decode"]}}]} "route": {"type": "object", "additionalProperties": false, "required": ["method", "path", "type", "render"], "properties": {"method": {"enum": ["GET", "POST", "PUT", "PATCH", "DELETE"]}, "path": {"type": "string", "pattern": "^/"}, "type": {"enum": ["submit", "query", "dynamic"]}, "action": {"type": "string"}, "taskIdParam": {"type": "string"}, "decode": {"type": "string"}, "render": {"type": "string"}, "models": {"type": "array", "minItems": 1, "uniqueItems": true, "items": {"type": "string", "minLength": 1}}}, "allOf": [{"if": {"properties": {"type": {"const": "query"}}}, "then": {"allOf": [{"not": {"required": ["decode"]}}, {"not": {"required": ["models"]}}]}}, {"if": {"properties": {"type": {"enum": ["submit", "dynamic"]}}}, "then": {"required": ["decode"]}}]},
"taskQueryContext": {
"type": "object",
"additionalProperties": false,
"required": ["taskId", "publicTaskId", "action", "model", "upstreamModel", "baseUrl", "authHeader", "data", "state"],
"properties": {
"taskId": {"type": "string"},
"publicTaskId": {"type": "string"},
"action": {"type": "string"},
"model": {"type": "string"},
"upstreamModel": {"type": "string"},
"baseUrl": {"type": "string"},
"apiKey": {"type": "string"},
"authHeader": {"type": "string"},
"auth": true,
"data": true,
"state": true
}
},
"batchQueryContext": {
"type": "object",
"additionalProperties": false,
"required": ["baseUrl", "authHeader", "tasks"],
"properties": {
"baseUrl": {"type": "string"},
"apiKey": {"type": "string"},
"authHeader": {"type": "string"},
"auth": true,
"tasks": {"type": "array", "items": {"$ref": "#/$defs/taskQueryContext"}}
}
},
"hookHTTPResponse": {
"type": "object",
"additionalProperties": false,
"required": ["status", "headers"],
"properties": {
"status": {"type": "integer"},
"headers": {"type": "object", "additionalProperties": {"type": "string"}}
}
}
} }
} }
model/log_other.go
...@@ -140,9 +140,7 @@ func (o *LogOther) toMap() map[string]any { ...@@ -140,9 +140,7 @@ func (o *LogOther) toMap() map[string]any {
return result return result
} }
for key, value := range o.public { maps.Copy(result, o.public)
result[key] = value
}
if adminInfo := copyLogOtherMap(o.adminInfo); len(adminInfo) > 0 { if adminInfo := copyLogOtherMap(o.adminInfo); len(adminInfo) > 0 {
result[logOtherAdminInfoKey] = adminInfo result[logOtherAdminInfoKey] = adminInfo
} }
......
model/task.go
...@@ -126,6 +126,11 @@ type TaskPrivateData struct { ...@@ -126,6 +126,11 @@ type TaskPrivateData struct {
// disconnect regardless; this only echoes the protocol-level request // disconnect regardless; this only echoes the protocol-level request
// attribute back on retrieval snapshots. // attribute back on retrieval snapshots.
ResponsesBackground bool `json:"responses_background,omitempty"` ResponsesBackground bool `json:"responses_background,omitempty"`
// PluginState is plugin-owned cross-round data. Unlike Task.Data it is
// only replaced when a hook explicitly returns state.
PluginState json.RawMessage `json:"plugin_state,omitempty"`
// PollFailures counts consecutive unrecognized or transient poll outcomes.
PollFailures int `json:"poll_failures,omitempty"`
} }
type TaskExecutionSnapshot struct { type TaskExecutionSnapshot struct {
...@@ -194,7 +199,10 @@ func (p *TaskPrivateData) Scan(val interface{}) error { ...@@ -194,7 +199,10 @@ func (p *TaskPrivateData) Scan(val interface{}) error {
} }
func (p TaskPrivateData) Value() (driver.Value, error) { func (p TaskPrivateData) Value() (driver.Value, error) {
if (p == TaskPrivateData{}) { if p.Key == "" && p.UpstreamTaskID == "" && p.ResultURL == "" &&
p.Execution == nil && p.BillingSource == "" && p.SubscriptionId == 0 &&
p.TokenId == 0 && p.NodeName == "" && p.BillingContext == nil &&
!p.ResponsesBackground && len(p.PluginState) == 0 && p.PollFailures == 0 {
return nil, nil return nil, nil
} }
// 同 Properties.Value:string 避免 PG simple protocol 的 bytea 编码。 // 同 Properties.Value:string 避免 PG simple protocol 的 bytea 编码。
...@@ -473,6 +481,8 @@ type taskSnapshot struct { ...@@ -473,6 +481,8 @@ type taskSnapshot struct {
FailReason string FailReason string
ResultURL string ResultURL string
Data json.RawMessage Data json.RawMessage
PluginState json.RawMessage
PollFailures int
} }
func (s taskSnapshot) Equal(other taskSnapshot) bool { func (s taskSnapshot) Equal(other taskSnapshot) bool {
...@@ -482,7 +492,9 @@ func (s taskSnapshot) Equal(other taskSnapshot) bool { ...@@ -482,7 +492,9 @@ func (s taskSnapshot) Equal(other taskSnapshot) bool {
s.FinishTime == other.FinishTime && s.FinishTime == other.FinishTime &&
s.FailReason == other.FailReason && s.FailReason == other.FailReason &&
s.ResultURL == other.ResultURL && s.ResultURL == other.ResultURL &&
bytes.Equal(s.Data, other.Data) bytes.Equal(s.Data, other.Data) &&
bytes.Equal(s.PluginState, other.PluginState) &&
s.PollFailures == other.PollFailures
} }
func (t *Task) Snapshot() taskSnapshot { func (t *Task) Snapshot() taskSnapshot {
...@@ -494,6 +506,8 @@ func (t *Task) Snapshot() taskSnapshot { ...@@ -494,6 +506,8 @@ func (t *Task) Snapshot() taskSnapshot {
FailReason: t.FailReason, FailReason: t.FailReason,
ResultURL: t.PrivateData.ResultURL, ResultURL: t.PrivateData.ResultURL,
Data: t.Data, Data: t.Data,
PluginState: t.PrivateData.PluginState,
PollFailures: t.PrivateData.PollFailures,
} }
} }
......
model/task_cas_test.go
...@@ -177,6 +177,29 @@ func TestSnapshotEqual_NilVsEmpty(t *testing.T) { ...@@ -177,6 +177,29 @@ func TestSnapshotEqual_NilVsEmpty(t *testing.T) {
assert.True(t, a.Equal(b)) assert.True(t, a.Equal(b))
} }
func TestSnapshotEqual_PluginStateAndPollFailures(t *testing.T) {
base := taskSnapshot{
Status: TaskStatusInProgress,
PluginState: json.RawMessage(`{"req_key":"a"}`),
PollFailures: 2,
}
assert.True(t, base.Equal(taskSnapshot{
Status: TaskStatusInProgress,
PluginState: json.RawMessage(`{"req_key":"a"}`),
PollFailures: 2,
}))
assert.False(t, base.Equal(taskSnapshot{
Status: TaskStatusInProgress,
PluginState: json.RawMessage(`{"req_key":"b"}`),
PollFailures: 2,
}))
assert.False(t, base.Equal(taskSnapshot{
Status: TaskStatusInProgress,
PluginState: json.RawMessage(`{"req_key":"a"}`),
PollFailures: 3,
}))
}
func TestSnapshot_Roundtrip(t *testing.T) { func TestSnapshot_Roundtrip(t *testing.T) {
task := &Task{ task := &Task{
Status: TaskStatusInProgress, Status: TaskStatusInProgress,
...@@ -186,6 +209,8 @@ func TestSnapshot_Roundtrip(t *testing.T) { ...@@ -186,6 +209,8 @@ func TestSnapshot_Roundtrip(t *testing.T) {
FailReason: "timeout", FailReason: "timeout",
PrivateData: TaskPrivateData{ PrivateData: TaskPrivateData{
ResultURL: "https://example.com/result.mp4", ResultURL: "https://example.com/result.mp4",
PluginState: json.RawMessage(`{"req_key":"keep"}`),
PollFailures: 3,
}, },
Data: json.RawMessage(`{"model":"test-model"}`), Data: json.RawMessage(`{"model":"test-model"}`),
} }
...@@ -197,6 +222,8 @@ func TestSnapshot_Roundtrip(t *testing.T) { ...@@ -197,6 +222,8 @@ func TestSnapshot_Roundtrip(t *testing.T) {
assert.Equal(t, task.FailReason, snap.FailReason) assert.Equal(t, task.FailReason, snap.FailReason)
assert.Equal(t, task.PrivateData.ResultURL, snap.ResultURL) assert.Equal(t, task.PrivateData.ResultURL, snap.ResultURL)
assert.JSONEq(t, string(task.Data), string(snap.Data)) assert.JSONEq(t, string(task.Data), string(snap.Data))
assert.Equal(t, task.PrivateData.PluginState, snap.PluginState)
assert.Equal(t, task.PrivateData.PollFailures, snap.PollFailures)
} }
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
...@@ -292,3 +319,30 @@ func TestUpdateWithStatus_ConcurrentWinner(t *testing.T) { ...@@ -292,3 +319,30 @@ func TestUpdateWithStatus_ConcurrentWinner(t *testing.T) {
} }
assert.Equal(t, 1, winCount, "exactly one goroutine should win the CAS") assert.Equal(t, 1, winCount, "exactly one goroutine should win the CAS")
} }
func TestUpdateWithStatus_PersistsPluginStateAndPollFailures(t *testing.T) {
truncateTables(t)
task := &Task{
TaskID: "task_cas_plugin_state",
Status: TaskStatusInProgress,
Data: json.RawMessage(`{}`),
PrivateData: TaskPrivateData{
PluginState: json.RawMessage(`{"req_key":"old"}`),
PollFailures: 1,
},
}
insertTask(t, task)
task.PrivateData.PluginState = json.RawMessage(`{"req_key":"new"}`)
task.PrivateData.PollFailures = 4
won, err := task.UpdateWithStatus(TaskStatusInProgress)
require.NoError(t, err)
require.True(t, won)
var reloaded Task
require.NoError(t, DB.First(&reloaded, task.ID).Error)
assert.EqualValues(t, TaskStatusInProgress, reloaded.Status)
assert.JSONEq(t, `{"req_key":"new"}`, string(reloaded.PrivateData.PluginState))
assert.Equal(t, 4, reloaded.PrivateData.PollFailures)
}
pkg/jsplugin/cli.go
...@@ -33,6 +33,7 @@ func RunCLI(args []string, stdout, stderr io.Writer) int { ...@@ -33,6 +33,7 @@ func RunCLI(args []string, stdout, stderr io.Writer) int {
fmt.Fprintf(stderr, "plugin lint failed: %v\n", compileErr) fmt.Fprintf(stderr, "plugin lint failed: %v\n", compileErr)
return 1 return 1
} }
warnParseTaskResultInProgressFallback(string(source), stderr)
fmt.Fprintf(stdout, "plugin %s@%s is valid\n", plugin.Meta.Key, plugin.Meta.Version) fmt.Fprintf(stdout, "plugin %s@%s is valid\n", plugin.Meta.Key, plugin.Meta.Version)
return 0 return 0
case "test": case "test":
...@@ -57,3 +58,35 @@ func RunCLI(args []string, stdout, stderr io.Writer) int { ...@@ -57,3 +58,35 @@ func RunCLI(args []string, stdout, stderr io.Writer) int {
return 2 return 2
} }
} }
func warnParseTaskResultInProgressFallback(source string, stderr io.Writer) {
body := parseTaskResultFunctionBody(source)
if strings.Contains(body, `|| "IN_PROGRESS"`) || strings.Contains(body, `|| 'IN_PROGRESS'`) {
fmt.Fprintln(stderr, `warning: parseTaskResult uses || "IN_PROGRESS" fallback; return UNKNOWN for unrecognized statuses`)
}
}
func parseTaskResultFunctionBody(source string) string {
marker := strings.Index(source, "function parseTaskResult")
if marker < 0 {
return ""
}
brace := strings.Index(source[marker:], "{")
if brace < 0 {
return ""
}
start := marker + brace
depth := 0
for i := start; i < len(source); i++ {
switch source[i] {
case '{':
depth++
case '}':
depth--
if depth == 0 {
return source[start : i+1]
}
}
}
return ""
}
pkg/jsplugin/cli_test.go
...@@ -28,6 +28,25 @@ func TestPluginCLI(t *testing.T) { ...@@ -28,6 +28,25 @@ func TestPluginCLI(t *testing.T) {
assert.Contains(t, stdout.String(), "1/1 cases") assert.Contains(t, stdout.String(), "1/1 cases")
} }
func TestPluginCLIWarnsOnParseTaskResultInProgressFallback(t *testing.T) {
tempDir := t.TempDir()
pluginPath := filepath.Join(tempDir, "fallback.js")
require.NoError(t, os.WriteFile(pluginPath, []byte(`
export const meta = { apiVersion: 1, key: "fallback", name: "Fallback", version: "1.0.0", author: {name: "Test"}, models: ["m"], fetchMode: "per_task" };
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl}; }
export function parseSubmitResponse() { return {taskId: "task"}; }
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl}; }
export function parseTaskResult(ctx, body) { return {status: statuses[body.status] || "IN_PROGRESS"}; }
const statuses = { done: "SUCCESS" };
`), 0o600))
var stdout bytes.Buffer
var stderr bytes.Buffer
assert.Equal(t, 0, RunCLI([]string{"lint", pluginPath}, &stdout, &stderr))
assert.Contains(t, stdout.String(), "plugin fallback@1.0.0 is valid")
assert.Contains(t, stderr.String(), `|| "IN_PROGRESS"`)
}
const cliFixturePluginSource = ` const cliFixturePluginSource = `
export const meta = { apiVersion: 1, key: "cli-fixture", name: "CLI Fixture", version: "1.0.0", author: {name: "Test"}, channelTypes: [1003], models: ["fixture-model"], fetchMode: "per_task" }; export const meta = { apiVersion: 1, key: "cli-fixture", name: "CLI Fixture", version: "1.0.0", author: {name: "Test"}, channelTypes: [1003], models: ["fixture-model"], fetchMode: "per_task" };
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl}; } export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl}; }
......
plugins/hailuo_responses_test.go
...@@ -346,6 +346,9 @@ func TestHailuoParseTaskResult(t *testing.T) { ...@@ -346,6 +346,9 @@ func TestHailuoParseTaskResult(t *testing.T) {
{"H3 permanent query error", `{"type":"error","error":{"type":"authorized_error","message":"login failed","http_code":"401"}}`, "FAILURE", "", "login failed"}, {"H3 permanent query error", `{"type":"error","error":{"type":"authorized_error","message":"login failed","http_code":"401"}}`, "FAILURE", "", "login failed"},
{"legacy success", `{"task_id":"1","status":"Success","file_id":"f1","base_resp":{"status_code":0}}`, "SUCCESS", "", ""}, {"legacy success", `{"task_id":"1","status":"Success","file_id":"f1","base_resp":{"status_code":0}}`, "SUCCESS", "", ""},
{"legacy processing", `{"task_id":"1","status":"Processing","base_resp":{"status_code":0}}`, "IN_PROGRESS", "", ""}, {"legacy processing", `{"task_id":"1","status":"Processing","base_resp":{"status_code":0}}`, "IN_PROGRESS", "", ""},
{"H3 unrecognized", `{"task":{"id":"1","status":"weird"}}`, "UNKNOWN", "", "unrecognized status: weird"},
{"legacy unrecognized", `{"task_id":"1","status":"Weird","base_resp":{"status_code":0}}`, "UNKNOWN", "", "unrecognized status: Weird"},
{"legacy base_resp failure", `{"task_id":"1","status":"Success","base_resp":{"status_code":1001,"status_msg":"upstream down"}}`, "FAILURE", "", "upstream down"},
} }
for _, testCase := range testCases { for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) { t.Run(testCase.name, func(t *testing.T) {
...@@ -451,7 +454,7 @@ func TestHailuoH3CompletionUsageFacts(t *testing.T) { ...@@ -451,7 +454,7 @@ func TestHailuoH3CompletionUsageFacts(t *testing.T) {
t.Run("polling adaptor carries actual facts into task settlement", func(t *testing.T) { t.Run("polling adaptor carries actual facts into task settlement", func(t *testing.T) {
adaptor := taskplugin.New(plugin) adaptor := taskplugin.New(plugin)
result, err := adaptor.ParseTaskResult([]byte( result, err := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, []byte(
`{"task":{"id":"1","status":"succeeded","resolution":"2K","usage":{"output_seconds":5,"input_seconds":7.5,"input_image_count":6}}}`, `{"task":{"id":"1","status":"succeeded","resolution":"2K","usage":{"output_seconds":5,"input_seconds":7.5,"input_image_count":6}}}`,
)) ))
require.NoError(t, err) require.NoError(t, err)
......
plugins/jimeng_responses_test.go
package plugins_test package plugins_test
import "testing" import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/pkg/jsplugin"
builtinplugins "github.com/QuantumNous/new-api/plugins"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestJimengResponsesProtocol(t *testing.T) { func TestJimengResponsesProtocol(t *testing.T) {
testVideoResponsesProtocol(t, videoResponsesTestCase{ testVideoResponsesProtocol(t, videoResponsesTestCase{
...@@ -25,3 +33,60 @@ func TestJimengResponsesProtocol(t *testing.T) { ...@@ -25,3 +33,60 @@ func TestJimengResponsesProtocol(t *testing.T) {
wantVendorName: "jimeng", wantVendorName: "jimeng",
}) })
} }
func loadJimengPlugin(t *testing.T) *jsplugin.LoadedPlugin {
t.Helper()
source, err := builtinplugins.Source("jimeng")
require.NoError(t, err)
plugin, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: "jimeng"})
require.NoError(t, err)
return plugin
}
func TestJimengSubmitStateDrivesQueryReqKey(t *testing.T) {
plugin := loadJimengPlugin(t)
submitValue, err := plugin.Engine.Call(t.Context(), "parseSubmitResponse", map[string]any{
"upstreamModel": "jimeng_vgfm_i2v_l20",
"requestBody": map[string]any{"images": []any{"https://cdn.example/frame.png"}},
}, map[string]any{"body": map[string]any{"code": 10000, "data": map[string]any{"task_id": "t1"}}})
require.NoError(t, err)
encoded, err := common.Marshal(submitValue)
require.NoError(t, err)
var submit map[string]any
require.NoError(t, common.Unmarshal(encoded, &submit))
state, ok := submit["state"].(map[string]any)
require.True(t, ok)
assert.Equal(t, "jimeng_vgfm_i2v_l20", state["req_key"])
queryValue, err := plugin.Engine.Call(t.Context(), "buildQueryRequest", map[string]any{
"taskId": "t1",
"action": "text_to_video",
"baseUrl": "https://jimeng.example",
"apiKey": "sk-test",
"state": map[string]any{"req_key": "custom_req_key"},
})
require.NoError(t, err)
queryEncoded, err := common.Marshal(queryValue)
require.NoError(t, err)
var query map[string]any
require.NoError(t, common.Unmarshal(queryEncoded, &query))
var body map[string]any
require.NoError(t, common.UnmarshalJsonStr(common.Interface2String(query["body"]), &body))
assert.Equal(t, "custom_req_key", body["req_key"])
assert.Equal(t, "t1", body["task_id"])
}
func TestJimengParseTaskResultUnknownStatus(t *testing.T) {
plugin := loadJimengPlugin(t)
value, err := plugin.Engine.Call(t.Context(), "parseTaskResult", map[string]any{}, map[string]any{
"code": 10000,
"data": map[string]any{"status": "weird"},
})
require.NoError(t, err)
encoded, err := common.Marshal(value)
require.NoError(t, err)
var result map[string]any
require.NoError(t, common.Unmarshal(encoded, &result))
assert.Equal(t, "UNKNOWN", result["status"])
assert.Contains(t, common.Interface2String(result["reason"]), "weird")
}
plugins/tasks/alibaba/plugin.js
...@@ -7,7 +7,7 @@ export const meta = { ...@@ -7,7 +7,7 @@ export const meta = {
en: "Alibaba Cloud Bailian Wanxiang video generation (text-to-video and image-to-video)", en: "Alibaba Cloud Bailian Wanxiang video generation (text-to-video and image-to-video)",
zh: "阿里云百炼万相视频生成（文生视频、图生视频）", zh: "阿里云百炼万相视频生成（文生视频、图生视频）",
}, },
version: "1.0.0", version: "1.0.1",
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
channelTypes: [17], channelTypes: [17],
models: [ models: [
...@@ -270,7 +270,7 @@ export function parseTaskResult(ctx, body) { ...@@ -270,7 +270,7 @@ export function parseTaskResult(ctx, body) {
if (!reason) reason = "task failed"; if (!reason) reason = "task failed";
return { status: "FAILURE", reason: reason }; return { status: "FAILURE", reason: reason };
} }
return { status: "QUEUED" }; return { status: "UNKNOWN", reason: "unrecognized status: " + String(output.task_status || "") };
} }
function artifactData(ctx) { function artifactData(ctx) {
......
plugins/tasks/doubao/plugin.js
...@@ -7,7 +7,7 @@ export const meta = { ...@@ -7,7 +7,7 @@ export const meta = {
en: "Volcengine Doubao Seedance video generation (text-to-video, image-to-video, and video-to-video)", en: "Volcengine Doubao Seedance video generation (text-to-video, image-to-video, and video-to-video)",
zh: "火山引擎豆包 Seedance 视频生成（文生视频、图生视频、视频生视频）", zh: "火山引擎豆包 Seedance 视频生成（文生视频、图生视频、视频生视频）",
}, },
version: "1.0.0", version: "1.0.1",
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
channelTypes: [54, 45], // VolcEngine-type channels serve Ark video models with the same wire format channelTypes: [54, 45], // VolcEngine-type channels serve Ark video models with the same wire format
models: [ models: [
...@@ -324,7 +324,7 @@ export function parseTaskResult(ctx, body) { ...@@ -324,7 +324,7 @@ export function parseTaskResult(ctx, body) {
const reason = body.error && body.error.message ? body.error.message : body.status; const reason = body.error && body.error.message ? body.error.message : body.status;
return { status: "FAILURE", progress: "100%", reason: reason }; return { status: "FAILURE", progress: "100%", reason: reason };
} }
return { status: "IN_PROGRESS", progress: "30%" }; return { status: "UNKNOWN", reason: "unrecognized status: " + String(body.status || "") };
} }
function artifactData(ctx) { function artifactData(ctx) {
......
plugins/tasks/google/plugin.js
...@@ -7,7 +7,7 @@ export const meta = { ...@@ -7,7 +7,7 @@ export const meta = {
en: "Google Veo video generation on the Gemini API (text-to-video and image-to-video)", en: "Google Veo video generation on the Gemini API (text-to-video and image-to-video)",
zh: "Google Veo 视频生成（文生视频、图生视频），Gemini API 版本", zh: "Google Veo 视频生成（文生视频、图生视频），Gemini API 版本",
}, },
version: "1.0.0", version: "1.0.1",
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
channelTypes: [24], channelTypes: [24],
models: ["veo-3.0-generate-001", "veo-3.0-fast-generate-001", "veo-3.1-generate-preview", "veo-3.1-fast-generate-preview"], models: ["veo-3.0-generate-001", "veo-3.0-fast-generate-001", "veo-3.1-generate-preview", "veo-3.1-fast-generate-preview"],
...@@ -209,7 +209,11 @@ export function buildQueryRequest(ctx) { ...@@ -209,7 +209,11 @@ export function buildQueryRequest(ctx) {
export function parseTaskResult(ctx, body) { export function parseTaskResult(ctx, body) {
if (body.error && body.error.message) return { status: "FAILURE", progress: "100%", reason: body.error.message }; if (body.error && body.error.message) return { status: "FAILURE", progress: "100%", reason: body.error.message };
if (!body.done) return { status: "IN_PROGRESS", progress: "50%" }; // Google long-running operations omit `done` (proto3 default) while still
// running, so a missing key means in-progress; only a non-operation shape is
// unrecognized.
if (!body || typeof body !== "object" || !String(body.name || "").trim()) return { status: "UNKNOWN", reason: "unrecognized operation state" };
if (body.done !== true) return { status: "IN_PROGRESS", progress: "50%" };
const videos = ((body.response || {}).generateVideoResponse || {}).generatedVideos || []; const videos = ((body.response || {}).generateVideoResponse || {}).generatedVideos || [];
const uri = videos.length && videos[0].video ? videos[0].video.uri || "" : ""; const uri = videos.length && videos[0].video ? videos[0].video.uri || "" : "";
return { taskId: utils.base64URL(body.name || ""), status: "SUCCESS", progress: "100%", remoteUrl: uri }; return { taskId: utils.base64URL(body.name || ""), status: "SUCCESS", progress: "100%", remoteUrl: uri };
......
plugins/tasks/hailuo/plugin.js
...@@ -7,7 +7,7 @@ export const meta = { ...@@ -7,7 +7,7 @@ export const meta = {
en: "MiniMax Hailuo video generation (text-to-video, image-to-video, and MiniMax-H3 multimodal reference)", en: "MiniMax Hailuo video generation (text-to-video, image-to-video, and MiniMax-H3 multimodal reference)",
zh: "MiniMax 海螺视频生成（文生视频、图生视频、MiniMax-H3 多模态参考生视频）", zh: "MiniMax 海螺视频生成（文生视频、图生视频、MiniMax-H3 多模态参考生视频）",
}, },
version: "1.1.1", version: "1.1.2",
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
channelTypes: [35], channelTypes: [35],
models: [ models: [
...@@ -465,8 +465,6 @@ export function buildQueryRequest(ctx) { ...@@ -465,8 +465,6 @@ export function buildQueryRequest(ctx) {
} }
export function parseTaskResult(ctx, body) { export function parseTaskResult(ctx, body) {
// The host calls this hook with an empty context, so the response envelope is
// the only way to tell a /v2 result from a /v1 one.
const apiError = h3APIError(body); const apiError = h3APIError(body);
if (apiError) { if (apiError) {
if (apiError.statusCode === 408 || apiError.statusCode === 429 || apiError.statusCode >= 500) throw new Error(apiError.message); if (apiError.statusCode === 408 || apiError.statusCode === 429 || apiError.statusCode >= 500) throw new Error(apiError.message);
...@@ -475,7 +473,10 @@ export function parseTaskResult(ctx, body) { ...@@ -475,7 +473,10 @@ export function parseTaskResult(ctx, body) {
const h3Task = h3QueryTask(body); const h3Task = h3QueryTask(body);
if (h3Task) { if (h3Task) {
const h3Statuses = { queued: "QUEUED", running: "IN_PROGRESS", succeeded: "SUCCESS", failed: "FAILURE", cancelled: "FAILURE" }; const h3Statuses = { queued: "QUEUED", running: "IN_PROGRESS", succeeded: "SUCCESS", failed: "FAILURE", cancelled: "FAILURE" };
const h3Status = h3Statuses[h3Task.status] || "IN_PROGRESS"; const h3Status = h3Statuses[h3Task.status];
if (!h3Status) {
return { status: "UNKNOWN", reason: "unrecognized status: " + String(h3Task.status || "") };
}
const h3Result = { code: 0, status: h3Status, progress: h3Status === "QUEUED" ? "30%" : h3Status === "IN_PROGRESS" ? "50%" : "100%" }; const h3Result = { code: 0, status: h3Status, progress: h3Status === "QUEUED" ? "30%" : h3Status === "IN_PROGRESS" ? "50%" : "100%" };
if (h3Status === "SUCCESS") { if (h3Status === "SUCCESS") {
const url = trimmed(h3Task.content && h3Task.content.url); const url = trimmed(h3Task.content && h3Task.content.url);
...@@ -486,11 +487,17 @@ export function parseTaskResult(ctx, body) { ...@@ -486,11 +487,17 @@ export function parseTaskResult(ctx, body) {
} }
return h3Result; return h3Result;
} }
if (body.base_resp && body.base_resp.status_code !== 0) {
return { code: body.base_resp.status_code || 0, status: "FAILURE", progress: "100%", reason: body.base_resp.status_msg || "" };
}
const base = body.base_resp || {}; const base = body.base_resp || {};
const statuses = { Preparing: "IN_PROGRESS", Queueing: "IN_PROGRESS", Processing: "IN_PROGRESS", Success: "SUCCESS", Fail: "FAILURE" }; const statuses = { Preparing: "IN_PROGRESS", Queueing: "IN_PROGRESS", Processing: "IN_PROGRESS", Success: "SUCCESS", Fail: "FAILURE" };
const status = statuses[body.status] || "IN_PROGRESS"; const status = statuses[body.status];
if (!status) {
return { status: "UNKNOWN", reason: "unrecognized status: " + String(body.status || "") };
}
const progress = status === "SUCCESS" || status === "FAILURE" ? "100%" : body.status === "Processing" ? "50%" : "30%"; const progress = status === "SUCCESS" || status === "FAILURE" ? "100%" : body.status === "Processing" ? "50%" : "30%";
const reason = base.status_code !== 0 ? base.status_msg || "" : status === "FAILURE" ? "task failed" : ""; const reason = status === "FAILURE" ? "task failed" : "";
return { code: base.status_code || 0, status: status, progress: progress, reason: reason }; return { code: base.status_code || 0, status: status, progress: progress, reason: reason };
} }
......
plugins/tasks/jimeng/plugin.js
...@@ -7,7 +7,7 @@ export const meta = { ...@@ -7,7 +7,7 @@ export const meta = {
en: "Volcengine Jimeng video generation (text-to-video, image-to-video, and first-and-last-frame)", en: "Volcengine Jimeng video generation (text-to-video, image-to-video, and first-and-last-frame)",
zh: "火山引擎即梦视频生成（文生视频、图生视频、首尾帧）", zh: "火山引擎即梦视频生成（文生视频、图生视频、首尾帧）",
}, },
version: "1.0.0", version: "1.0.1",
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
channelTypes: [51], channelTypes: [51],
models: ["jimeng_vgfm_t2v_l20"], models: ["jimeng_vgfm_t2v_l20"],
...@@ -270,10 +270,8 @@ function filePlaceholder(image) { ...@@ -270,10 +270,8 @@ function filePlaceholder(image) {
} }
function queryReqKey(ctx) { function queryReqKey(ctx) {
const data = (ctx && ctx.data) || {}; const state = (ctx && ctx.state) || {};
if (typeof data.req_key === "string" && data.req_key.trim()) return data.req_key.trim(); if (typeof state.req_key === "string" && state.req_key.trim()) return state.req_key.trim();
const req = (ctx && ctx.requestBody) || {};
if (typeof req.req_key === "string" && req.req_key.trim()) return req.req_key.trim();
if (ctx && ctx.action === "image_to_video") return "jimeng_vgfm_i2v_l20"; if (ctx && ctx.action === "image_to_video") return "jimeng_vgfm_i2v_l20";
if (ctx && ctx.action === "first_tail_to_video") return "jimeng_i2v_first_tail_v30"; if (ctx && ctx.action === "first_tail_to_video") return "jimeng_i2v_first_tail_v30";
return "jimeng_vgfm_t2v_l20"; return "jimeng_vgfm_t2v_l20";
...@@ -349,7 +347,7 @@ export function parseSubmitResponse(ctx, resp) { ...@@ -349,7 +347,7 @@ export function parseSubmitResponse(ctx, resp) {
const body = resp.body || {}; const body = resp.body || {};
if (body.code !== 10000) throw new Error(body.message || "jimeng submit failed"); if (body.code !== 10000) throw new Error(body.message || "jimeng submit failed");
if (!body.data || !body.data.task_id) throw new Error("missing task_id"); if (!body.data || !body.data.task_id) throw new Error("missing task_id");
return { taskId: body.data.task_id, taskData: Object.assign({}, body, { req_key: submitReqKey(ctx) }) }; return { taskId: body.data.task_id, taskData: Object.assign({}, body, { req_key: submitReqKey(ctx) }), state: { req_key: submitReqKey(ctx) } };
} }
export function extractUsage(ctx) { export function extractUsage(ctx) {
...@@ -366,22 +364,20 @@ export function buildQueryRequest(ctx) { ...@@ -366,22 +364,20 @@ export function buildQueryRequest(ctx) {
export function parseTaskResult(ctx, body) { export function parseTaskResult(ctx, body) {
const data = body.data || {}; const data = body.data || {};
let status = "";
let progress = "";
if (body.code !== 10000) { if (body.code !== 10000) {
status = "FAILURE"; return { code: body.code || 0, status: "FAILURE", progress: "100%", reason: body.message || "" };
progress = "100%";
} }
if (data.status === "in_queue") { if (data.status === "in_queue") {
status = "QUEUED"; const result = { code: 0, status: "QUEUED", progress: "10%", reason: "" };
progress = "10%"; if (data.video_url) result.url = data.video_url;
} else if (data.status === "done") { return result;
status = "SUCCESS";
progress = "100%";
} }
const result = { code: body.code === 10000 ? 0 : body.code || 0, status: status, progress: progress, reason: body.code === 10000 ? "" : body.message || "" }; if (data.status === "done") {
const result = { code: 0, status: "SUCCESS", progress: "100%", reason: "" };
if (data.video_url) result.url = data.video_url; if (data.video_url) result.url = data.video_url;
return result; return result;
}
return { code: 0, status: "UNKNOWN", reason: "unrecognized status: " + String(data.status || "") };
} }
function artifactData(ctx) { function artifactData(ctx) {
......
plugins/tasks/kling/plugin.js
...@@ -7,7 +7,7 @@ export const meta = { ...@@ -7,7 +7,7 @@ export const meta = {
en: "Kuaishou Kling video generation (text-to-video and image-to-video)", en: "Kuaishou Kling video generation (text-to-video and image-to-video)",
zh: "快手可灵视频生成（文生视频、图生视频）", zh: "快手可灵视频生成（文生视频、图生视频）",
}, },
version: "1.0.0", version: "1.0.1",
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
channelTypes: [50], channelTypes: [50],
models: ["kling-v1", "kling-v1-6", "kling-v2-master"], models: ["kling-v1", "kling-v1-6", "kling-v2-master"],
...@@ -286,7 +286,7 @@ export function parseTaskResult(ctx, body) { ...@@ -286,7 +286,7 @@ export function parseTaskResult(ctx, body) {
const data = body.data || {}; const data = body.data || {};
const statuses = { submitted: "SUBMITTED", processing: "IN_PROGRESS", succeed: "SUCCESS", failed: "FAILURE" }; const statuses = { submitted: "SUBMITTED", processing: "IN_PROGRESS", succeed: "SUCCESS", failed: "FAILURE" };
const status = statuses[data.task_status]; const status = statuses[data.task_status];
if (!status) throw new Error("unknown task status: " + data.task_status); if (!status) return { status: "UNKNOWN", reason: "unknown task status: " + String(data.task_status || "") };
const videos = status === "SUCCESS" && data.task_result && data.task_result.videos ? data.task_result.videos : []; const videos = status === "SUCCESS" && data.task_result && data.task_result.videos ? data.task_result.videos : [];
const result = { code: body.code || 0, taskId: data.task_id, status: status, reason: data.task_status_msg || "" }; const result = { code: body.code || 0, taskId: data.task_id, status: status, reason: data.task_status_msg || "" };
if (videos.length && videos[0].url) result.url = videos[0].url; if (videos.length && videos[0].url) result.url = videos[0].url;
......
plugins/tasks/sora/plugin.js
...@@ -7,7 +7,7 @@ export const meta = { ...@@ -7,7 +7,7 @@ export const meta = {
en: "OpenAI Sora video generation (text-to-video, image-to-video, and remix)", en: "OpenAI Sora video generation (text-to-video, image-to-video, and remix)",
zh: "OpenAI Sora 视频生成（文生视频、图生视频、remix）", zh: "OpenAI Sora 视频生成（文生视频、图生视频、remix）",
}, },
version: "1.0.0", version: "1.0.1",
channelTypes: [55, 1], // OpenAI-type channels natively serve sora with the same wire format channelTypes: [55, 1], // OpenAI-type channels natively serve sora with the same wire format
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
models: ["sora-2", "sora-2-pro"], models: ["sora-2", "sora-2-pro"],
...@@ -145,7 +145,9 @@ export function parseTaskResult(ctx, body) { ...@@ -145,7 +145,9 @@ export function parseTaskResult(ctx, body) {
failed: "FAILURE", failed: "FAILURE",
cancelled: "FAILURE", cancelled: "FAILURE",
}; };
const result = { status: statuses[body.status] || "UNKNOWN" }; const mapped = statuses[body.status];
const result = { status: mapped || "UNKNOWN" };
if (!mapped) result.reason = "unrecognized status: " + String(body.status || "");
if (body.progress > 0 && body.progress < 100) result.progress = body.progress + "%"; if (body.progress > 0 && body.progress < 100) result.progress = body.progress + "%";
if (result.status === "FAILURE") result.reason = body.error && body.error.message ? body.error.message : "task failed"; if (result.status === "FAILURE") result.reason = body.error && body.error.message ? body.error.message : "task failed";
return result; return result;
......
plugins/tasks/sunoapi/plugin.js
...@@ -9,7 +9,7 @@ export const meta = { ...@@ -9,7 +9,7 @@ export const meta = {
en: "SunoAPI project music and lyrics generation", en: "SunoAPI project music and lyrics generation",
zh: "SunoAPI 项目 音乐与歌词生成", zh: "SunoAPI 项目 音乐与歌词生成",
}, },
version: "1.0.0", version: "1.0.1",
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
channelTypes: [36], channelTypes: [36],
models: ["suno_music", "suno_lyrics"], models: ["suno_music", "suno_lyrics"],
...@@ -133,19 +133,23 @@ export function extractUsage(ctx) { ...@@ -133,19 +133,23 @@ export function extractUsage(ctx) {
return { clips: action === "lyrics" ? 1 : 2, action: action }; return { clips: action === "lyrics" ? 1 : 2, action: action };
} }
export function buildBatchQueryRequest(ctx, taskIds) { export function buildBatchQueryRequest(ctx, tasks) {
return { return {
url: ctx.baseUrl + "/suno/fetch", url: ctx.baseUrl + "/suno/fetch",
method: "POST", method: "POST",
headers: { "Content-Type": "application/json", Authorization: "Bearer " + ctx.apiKey }, headers: { "Content-Type": "application/json", Authorization: "Bearer " + ctx.apiKey },
body: { ids: taskIds }, body: {
ids: (tasks || []).map(function (task) {
return task.taskId;
}),
},
}; };
} }
// Required v1 per-task hooks remain defined for contract compatibility. Suno's // Required v1 per-task hooks remain defined for contract compatibility. Suno's
// host polling path uses the batch hooks below. // host polling path uses the batch hooks below.
export function buildQueryRequest(ctx) { export function buildQueryRequest(ctx) {
return buildBatchQueryRequest(ctx, (ctx.requestBody || {}).ids || []); return buildBatchQueryRequest(ctx, [ctx]);
} }
export function parseBatchResult(ctx, body) { export function parseBatchResult(ctx, body) {
......
plugins/tasks/vertex-ai/plugin.js
...@@ -7,7 +7,7 @@ export const meta = { ...@@ -7,7 +7,7 @@ export const meta = {
en: "Google Veo video generation on Vertex AI (text-to-video and image-to-video)", en: "Google Veo video generation on Vertex AI (text-to-video and image-to-video)",
zh: "Google Veo 视频生成（文生视频、图生视频），Vertex AI 版本", zh: "Google Veo 视频生成（文生视频、图生视频），Vertex AI 版本",
}, },
version: "1.0.0", version: "1.0.1",
channelTypes: [41], channelTypes: [41],
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
models: ["veo-3.0-generate-001", "veo-3.0-fast-generate-001", "veo-3.1-generate-preview", "veo-3.1-fast-generate-preview"], models: ["veo-3.0-generate-001", "veo-3.0-fast-generate-001", "veo-3.1-generate-preview", "veo-3.1-fast-generate-preview"],
...@@ -225,7 +225,11 @@ export function buildQueryRequest(ctx) { ...@@ -225,7 +225,11 @@ export function buildQueryRequest(ctx) {
} }
export function parseTaskResult(ctx, body) { export function parseTaskResult(ctx, body) {
if (body.error && body.error.message) return { status: "FAILURE", progress: "100%", reason: body.error.message }; if (body.error && body.error.message) return { status: "FAILURE", progress: "100%", reason: body.error.message };
if (!body.done) return { status: "IN_PROGRESS", progress: "50%" }; // Google long-running operations omit `done` (proto3 default) while still
// running, so a missing key means in-progress; only a non-operation shape is
// unrecognized.
if (!body || typeof body !== "object" || !String(body.name || "").trim()) return { status: "UNKNOWN", reason: "unrecognized operation state" };
if (body.done !== true) return { status: "IN_PROGRESS", progress: "50%" };
const url = dataVideo(body.response || {}); const url = dataVideo(body.response || {});
return { status: "SUCCESS", progress: "100%", url: url, remoteUrl: url }; return { status: "SUCCESS", progress: "100%", url: url, remoteUrl: url };
} }
......
plugins/tasks/vidu/plugin.js
...@@ -7,7 +7,7 @@ export const meta = { ...@@ -7,7 +7,7 @@ export const meta = {
en: "Shengshu Vidu video generation (text-to-video, image-to-video, first-and-last-frame, and reference-to-video)", en: "Shengshu Vidu video generation (text-to-video, image-to-video, first-and-last-frame, and reference-to-video)",
zh: "生数 Vidu 视频生成（文生视频、图生视频、首尾帧、参考生视频）", zh: "生数 Vidu 视频生成（文生视频、图生视频、首尾帧、参考生视频）",
}, },
version: "1.0.0", version: "1.0.1",
author: { name: "QuantumNous" }, author: { name: "QuantumNous" },
channelTypes: [52], channelTypes: [52],
models: ["viduq2", "viduq1", "vidu2.0", "vidu1.5"], models: ["viduq2", "viduq1", "vidu2.0", "vidu1.5"],
...@@ -265,7 +265,7 @@ export function buildQueryRequest(ctx) { ...@@ -265,7 +265,7 @@ export function buildQueryRequest(ctx) {
export function parseTaskResult(ctx, body) { export function parseTaskResult(ctx, body) {
const statuses = { created: "SUBMITTED", queueing: "SUBMITTED", processing: "IN_PROGRESS", success: "SUCCESS", failed: "FAILURE" }; const statuses = { created: "SUBMITTED", queueing: "SUBMITTED", processing: "IN_PROGRESS", success: "SUCCESS", failed: "FAILURE" };
const status = statuses[body.state]; const status = statuses[body.state];
if (!status) throw new Error("unknown task state: " + body.state); if (!status) return { status: "UNKNOWN", reason: "unknown task state: " + String(body.state || "") };
const url = body.creations && body.creations.length ? body.creations[0].url || "" : ""; const url = body.creations && body.creations.length ? body.creations[0].url || "" : "";
const result = { status: status, reason: body.state === "failed" ? body.err_code || "" : "" }; const result = { status: status, reason: body.state === "failed" ? body.err_code || "" : "" };
if (url) result.url = url; if (url) result.url = url;
......
plugins/veo_poll_test.go 0 → 100644
package plugins_test
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/pkg/jsplugin"
builtinplugins "github.com/QuantumNous/new-api/plugins"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// Google long-running operations serialize proto3 defaults by omission: a
// still-running operation has no `done` key at all. Treating that as
// unrecognized would count every normal poll as a failure and fail the task
// at the poll-failure threshold while the video is still rendering.
func TestVeoParseTaskResultTreatsMissingDoneAsInProgress(t *testing.T) {
cases := []struct {
name string
body map[string]any
wantStatus string
}{
{
name: "running operation omits done",
body: map[string]any{"name": "operations/abc", "metadata": map[string]any{"@type": "x"}},
wantStatus: "IN_PROGRESS",
},
{
name: "explicit done false",
body: map[string]any{"name": "operations/abc", "done": false},
wantStatus: "IN_PROGRESS",
},
{
name: "body without operation name is unrecognized",
body: map[string]any{"foo": "bar"},
wantStatus: "UNKNOWN",
},
{
name: "operation error is failure",
body: map[string]any{"name": "operations/abc", "done": true, "error": map[string]any{"message": "quota exceeded"}},
wantStatus: "FAILURE",
},
}
for _, key := range []string{"google", "vertex-ai"} {
source, err := builtinplugins.Source(key)
require.NoError(t, err)
plugin, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: key})
require.NoError(t, err)
for _, tc := range cases {
t.Run(key+"/"+tc.name, func(t *testing.T) {
value, err := plugin.Engine.Call(t.Context(), "parseTaskResult", map[string]any{}, tc.body)
require.NoError(t, err)
encoded, err := common.Marshal(value)
require.NoError(t, err)
var result map[string]any
require.NoError(t, common.Unmarshal(encoded, &result))
assert.Equal(t, tc.wantStatus, result["status"])
})
}
}
}
relay/channel/adapter.go
...@@ -76,8 +76,8 @@ type TaskAdaptor interface { ...@@ -76,8 +76,8 @@ type TaskAdaptor interface {
// ── Polling ────────────────────────────────────────────────────── // ── Polling ──────────────────────────────────────────────────────
FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error) FetchTask(baseUrl, key string, task *model.Task, proxy string) (*http.Response, error)
ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) ParseTaskResult(task *model.Task, resp *http.Response, respBody []byte) (*relaycommon.TaskInfo, error)
} }
// TaskSubmitResponse is the transport-independent result of parsing an // TaskSubmitResponse is the transport-independent result of parsing an
...@@ -87,6 +87,7 @@ type TaskSubmitResponse struct { ...@@ -87,6 +87,7 @@ type TaskSubmitResponse struct {
TaskData []byte TaskData []byte
ClientResponse any ClientResponse any
Immediate *relaycommon.TaskInfo Immediate *relaycommon.TaskInfo
PluginState []byte
} }
type OpenAIVideoConverter interface { type OpenAIVideoConverter interface {
......
relay/channel/task/jsplugin/adaptor.go
...@@ -22,12 +22,12 @@ import ( ...@@ -22,12 +22,12 @@ import (
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
kitdto "github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay/channel" "github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
kitdto "github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
) )
...@@ -56,6 +56,7 @@ type submitResponse struct { ...@@ -56,6 +56,7 @@ type submitResponse struct {
TaskID string `json:"taskId"` TaskID string `json:"taskId"`
TaskData any `json:"taskData"` TaskData any `json:"taskData"`
Immediate *taskResult `json:"immediate"` Immediate *taskResult `json:"immediate"`
State any `json:"state"`
} }
type taskResult struct { type taskResult struct {
Code int `json:"code"` Code int `json:"code"`
...@@ -67,12 +68,16 @@ type taskResult struct { ...@@ -67,12 +68,16 @@ type taskResult struct {
RemoteURL string `json:"remoteUrl"` RemoteURL string `json:"remoteUrl"`
CompletionTokens float64 `json:"completionTokens"` CompletionTokens float64 `json:"completionTokens"`
TotalTokens float64 `json:"totalTokens"` TotalTokens float64 `json:"totalTokens"`
State any `json:"state"`
} }
var taskArtifactKeyPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._~-]{0,127}$`) var taskArtifactKeyPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._~-]{0,127}$`)
const maxTaskArtifacts = 64 const maxTaskArtifacts = 64
// maxTaskPluginPersistedJSONBytes is the shared ceiling for taskData and plugin state.
const maxTaskPluginPersistedJSONBytes = 1 << 20
type TaskAdaptor struct { type TaskAdaptor struct {
plugin *pluginruntime.LoadedPlugin plugin *pluginruntime.LoadedPlugin
info *relaycommon.RelayInfo info *relaycommon.RelayInfo
...@@ -497,10 +502,16 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r ...@@ -497,10 +502,16 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r
immediate != nil, immediate != nil,
time.Since(started).Milliseconds(), 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{ return &channel.TaskSubmitResponse{
UpstreamTaskID: parsed.TaskID, UpstreamTaskID: parsed.TaskID,
TaskData: taskData, TaskData: taskData,
Immediate: immediate, Immediate: immediate,
PluginState: pluginState,
}, nil }, nil
} }
...@@ -508,50 +519,32 @@ func (a *TaskAdaptor) GetModelList() []string { return append([]string(nil), a.p ...@@ -508,50 +519,32 @@ func (a *TaskAdaptor) GetModelList() []string { return append([]string(nil), a.p
func (a *TaskAdaptor) GetChannelName() string { return a.plugin.Meta.Name } func (a *TaskAdaptor) GetChannelName() string { return a.plugin.Meta.Name }
func (a *TaskAdaptor) FetchMode() string { return a.plugin.Meta.FetchMode } func (a *TaskAdaptor) FetchMode() string { return a.plugin.Meta.FetchMode }
func (a *TaskAdaptor) FetchBatchTasks(baseURL, key string, taskIDs []string, proxy string) (*http.Response, error) { func (a *TaskAdaptor) FetchBatchTasks(baseURL, key string, tasks []*model.Task, proxy string) (*http.Response, error) {
ctx := map[string]any{"baseUrl": baseURL} taskContexts := make([]map[string]any, 0, len(tasks))
auth, err := resolveAuth(a.plugin.Meta.Auth, key, proxy) for _, task := range tasks {
taskCtx, err := a.queryContext(task, key, baseURL, proxy)
if err != nil { if err != nil {
return nil, err return nil, err
} }
ctx["auth"] = auth taskContexts = append(taskContexts, taskCtx)
ctx["authHeader"] = auth["authHeader"] }
if a.plugin.Meta.Auth.Type == "" || a.plugin.Meta.Auth.Type == "none" || a.plugin.Meta.Auth.Type == "api_key" { ctx, err := a.batchQueryContext(key, baseURL, proxy, taskContexts)
ctx["apiKey"] = key if err != nil {
return nil, err
} }
value, err := a.plugin.Engine.Call(context.Background(), "buildBatchQueryRequest", ctx, taskIDs) value, err := a.plugin.Engine.Call(context.Background(), "buildBatchQueryRequest", ctx, taskContexts)
if err != nil { if err != nil {
return nil, err return nil, err
} }
return a.doFetchDescriptor(baseURL, proxy, value) return a.doFetchDescriptor(baseURL, proxy, value)
} }
func (a *TaskAdaptor) FetchTask(baseURL, key string, body map[string]any, proxy string) (*http.Response, error) { func (a *TaskAdaptor) FetchTask(baseURL, key string, task *model.Task, proxy string) (*http.Response, error) {
ctx := map[string]any{"taskId": body["task_id"], "action": body["action"], "requestBody": body, "baseUrl": baseURL} ctx, err := a.queryContext(task, key, baseURL, proxy)
// Query hooks are driver hooks and must see the same model identities as
// submit hooks. Polling has no relay info, so they arrive with the
// persisted task properties the caller puts in the fetch body.
originModel, _ := body["model"].(string)
upstreamModel, _ := body["upstream_model"].(string)
if upstreamModel == "" {
upstreamModel = originModel
}
ctx["model"] = originModel
ctx["upstreamModel"] = upstreamModel
auth, err := resolveAuth(a.plugin.Meta.Auth, key, proxy)
if err != nil { if err != nil {
return nil, err return nil, err
} }
ctx["auth"] = auth value, err := a.plugin.Engine.Call(context.Background(), "buildQueryRequest", ctx)
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
}
hook := "buildQueryRequest"
if a.plugin.Meta.FetchMode == "batch" && a.hasHook(context.Background(), "buildBatchQueryRequest") {
hook = "buildBatchQueryRequest"
}
value, err := a.plugin.Engine.Call(context.Background(), hook, ctx)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -616,14 +609,27 @@ func (a *TaskAdaptor) doFetchDescriptor(baseURL, proxy string, value any) (*http ...@@ -616,14 +609,27 @@ func (a *TaskAdaptor) doFetchDescriptor(baseURL, proxy string, value any) (*http
return resp, nil return resp, nil
} }
func (a *TaskAdaptor) ParseBatchResult(body []byte) (map[string]*service.BatchTaskResult, error) { func (a *TaskAdaptor) ParseBatchResult(tasks []*model.Task, resp *http.Response, body []byte) (map[string]*service.BatchTaskResult, error) {
started := time.Now() started := time.Now()
input := any(string(body)) input := any(string(body))
var decoded any var decoded any
if common.Unmarshal(body, &decoded) == nil { if common.Unmarshal(body, &decoded) == nil {
input = decoded input = decoded
} }
value, err := a.plugin.Engine.Call(context.Background(), "parseBatchResult", map[string]any{}, input) key, baseURL, proxy := a.queryCredentials()
taskContexts := make([]map[string]any, 0, len(tasks))
for _, task := range tasks {
taskCtx, err := a.queryContext(task, key, baseURL, proxy)
if err != nil {
return nil, err
}
taskContexts = append(taskContexts, taskCtx)
}
ctx, err := a.batchQueryContext(key, baseURL, proxy, taskContexts)
if err != nil {
return nil, err
}
value, err := a.plugin.Engine.Call(context.Background(), "parseBatchResult", ctx, input, hookHTTPResponse(resp))
if err != nil { 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()) 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 return nil, err
...@@ -639,6 +645,7 @@ func (a *TaskAdaptor) ParseBatchResult(body []byte) (map[string]*service.BatchTa ...@@ -639,6 +645,7 @@ func (a *TaskAdaptor) ParseBatchResult(body []byte) (map[string]*service.BatchTa
StartTime int64 `json:"startTime"` StartTime int64 `json:"startTime"`
FinishTime int64 `json:"finishTime"` FinishTime int64 `json:"finishTime"`
Data any `json:"data"` Data any `json:"data"`
State any `json:"state"`
} }
if err = convert(value, &parsed); err != nil { 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()) 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())
...@@ -651,12 +658,27 @@ func (a *TaskAdaptor) ParseBatchResult(body []byte) (map[string]*service.BatchTa ...@@ -651,12 +658,27 @@ func (a *TaskAdaptor) ParseBatchResult(body []byte) (map[string]*service.BatchTa
continue continue
} }
info := relaycommon.TaskInfo{TaskID: item.TaskID, Status: item.Status, Progress: item.Progress, Reason: item.Reason, Url: item.URL} info := relaycommon.TaskInfo{TaskID: item.TaskID, Status: item.Status, Progress: item.Progress, Reason: item.Reason, Url: item.URL}
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))
} else {
info.PluginState = pluginState
}
}
if hasCompletionUsage { if hasCompletionUsage {
usageBody := item.Data usageBody := item.Data
if usageBody == nil { if usageBody == nil {
usageBody = jsonValue(item) usageBody = jsonValue(item)
} }
facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", nil, jsonValue(&info), usageBody) itemCtx := ctx
for _, taskCtx := range taskContexts {
if fmt.Sprint(taskCtx["taskId"]) == item.TaskID {
itemCtx = taskCtx
break
}
}
facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", itemCtx, jsonValue(&info), usageBody)
if hookErr == nil { if hookErr == nil {
a.applyCompletionUsageFacts(&info, facts) a.applyCompletionUsageFacts(&info, facts)
} }
...@@ -675,14 +697,22 @@ func (a *TaskAdaptor) ParseBatchResult(body []byte) (map[string]*service.BatchTa ...@@ -675,14 +697,22 @@ func (a *TaskAdaptor) ParseBatchResult(body []byte) (map[string]*service.BatchTa
return results, nil return results, nil
} }
func (a *TaskAdaptor) ParseTaskResult(body []byte) (*relaycommon.TaskInfo, error) { func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, body []byte) (*relaycommon.TaskInfo, error) {
started := time.Now() started := time.Now()
input := any(string(body)) input := any(string(body))
var decoded any var decoded any
if common.Unmarshal(body, &decoded) == nil { if common.Unmarshal(body, &decoded) == nil {
input = decoded input = decoded
} }
value, err := a.plugin.Engine.Call(context.Background(), "parseTaskResult", map[string]any{}, input) key, baseURL, proxy := a.queryCredentials()
if task != nil && task.PrivateData.Key != "" {
key = task.PrivateData.Key
}
ctx, err := a.queryContext(task, key, baseURL, proxy)
if err != nil {
return nil, err
}
value, err := a.plugin.Engine.Call(context.Background(), "parseTaskResult", ctx, input, hookHTTPResponse(resp))
if err != nil { 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()) 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 return nil, err
...@@ -703,10 +733,17 @@ func (a *TaskAdaptor) ParseTaskResult(body []byte) (*relaycommon.TaskInfo, error ...@@ -703,10 +733,17 @@ func (a *TaskAdaptor) ParseTaskResult(body []byte) (*relaycommon.TaskInfo, error
CompletionTokens: positiveInt(parsed.CompletionTokens), CompletionTokens: positiveInt(parsed.CompletionTokens),
TotalTokens: positiveInt(parsed.TotalTokens), TotalTokens: positiveInt(parsed.TotalTokens),
} }
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)))
} else {
result.PluginState = pluginState
}
}
// The raw polling response only exists at this boundary. Capture upstream // The raw polling response only exists at this boundary. Capture upstream
// units here so the host settlement path can consume them from TaskInfo. // units here so the host settlement path can consume them from TaskInfo.
if a.hasHook(context.Background(), "extractUsageOnComplete") { if a.hasHook(context.Background(), "extractUsageOnComplete") {
facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", nil, jsonValue(result), input) facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", ctx, jsonValue(result), input)
if hookErr == nil { if hookErr == nil {
a.applyCompletionUsageFacts(result, facts) a.applyCompletionUsageFacts(result, facts)
} }
...@@ -901,15 +938,125 @@ func taskArtifactContext(task *model.Task) (map[string]any, error) { ...@@ -901,15 +938,125 @@ func taskArtifactContext(task *model.Task) (map[string]any, error) {
if task.PrivateData.Execution != nil && task.PrivateData.Execution.TaskPlugin != nil { if task.PrivateData.Execution != nil && task.PrivateData.Execution.TaskPlugin != nil {
producerVersion = task.PrivateData.Execution.TaskPlugin.Version producerVersion = task.PrivateData.Execution.TaskPlugin.Version
} }
var state any
if len(task.PrivateData.PluginState) > 0 {
if err := common.Unmarshal(task.PrivateData.PluginState, &state); err != nil {
return nil, fmt.Errorf("plugin state is invalid")
}
}
return map[string]any{ return map[string]any{
"taskId": task.TaskID, "taskId": task.TaskID,
"status": string(task.Status), "status": string(task.Status),
"action": task.Action, "action": task.Action,
"data": data, "data": data,
"state": state,
"producerVersion": producerVersion, "producerVersion": producerVersion,
}, nil }, nil
} }
func (a *TaskAdaptor) queryContext(task *model.Task, key, baseURL, proxy string) (map[string]any, error) {
ctx := map[string]any{
"taskId": "",
"publicTaskId": "",
"action": "",
"model": "",
"upstreamModel": "",
"baseUrl": baseURL,
"data": nil,
"state": nil,
}
if task != nil {
originModel := task.Properties.OriginModelName
upstreamModel := task.Properties.UpstreamModelName
if upstreamModel == "" {
upstreamModel = originModel
}
ctx["taskId"] = task.GetUpstreamTaskID()
ctx["publicTaskId"] = task.TaskID
ctx["action"] = constant.NormalizeTaskAction(task.Action)
ctx["model"] = originModel
ctx["upstreamModel"] = upstreamModel
if len(task.Data) > 0 {
var data any
if err := common.Unmarshal(task.Data, &data); err != nil {
return nil, fmt.Errorf("task data is invalid")
}
ctx["data"] = data
}
if len(task.PrivateData.PluginState) > 0 {
var state any
if err := common.Unmarshal(task.PrivateData.PluginState, &state); err != nil {
return nil, fmt.Errorf("plugin state is invalid")
}
ctx["state"] = state
}
if task.PrivateData.Key != "" {
key = task.PrivateData.Key
}
}
auth, err := resolveAuth(a.plugin.Meta.Auth, key, proxy)
if 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 {
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) queryCredentials() (key, baseURL, proxy string) {
if a.info == nil || !a.info.HasChannelMeta() {
return "", "", ""
}
return a.info.ApiKey, a.info.ChannelBaseUrl, a.info.ChannelSetting.Proxy
}
func hookHTTPResponse(resp *http.Response) map[string]any {
headers := map[string]string{}
status := 0
if resp != nil {
status = resp.StatusCode
for name, values := range resp.Header {
if len(values) > 0 {
headers[name] = values[0]
}
}
}
return map[string]any{"status": status, "headers": headers}
}
func encodeReturnedPluginState(value any) ([]byte, bool) {
object, ok := value.(map[string]any)
if !ok {
return nil, false
}
state, exists := object["state"]
if !exists || state == nil {
return nil, false
}
data, err := common.Marshal(state)
if err != nil {
return nil, false
}
return data, true
}
func validateTaskArtifacts(value any) ([]channel.TaskArtifact, error) { func validateTaskArtifacts(value any) ([]channel.TaskArtifact, error) {
encoded, err := common.Marshal(value) encoded, err := common.Marshal(value)
if err != nil { if err != nil {
......
relay/channel/task/jsplugin/adaptor_test.go
...@@ -503,12 +503,15 @@ func TestTaskAdaptorMapsJSContract(t *testing.T) { ...@@ -503,12 +503,15 @@ func TestTaskAdaptorMapsJSContract(t *testing.T) {
assert.Empty(t, recorder.Body.String(), "response parsing must not write before the durable task barrier") assert.Empty(t, recorder.Body.String(), "response parsing must not write before the durable task barrier")
assert.Equal(t, map[string]float64{"seconds": 7}, adaptor.AdjustBillingOnSubmit(info, []byte(`{"seconds":7}`))) assert.Equal(t, map[string]float64{"seconds": 7}, adaptor.AdjustBillingOnSubmit(info, []byte(`{"seconds":7}`)))
queryResp, err := adaptor.FetchTask(server.URL, "secret", map[string]any{"task_id": parsed.UpstreamTaskID, "action": info.Action}, "") queryResp, err := adaptor.FetchTask(server.URL, "secret", &model.Task{
Action: info.Action,
PrivateData: model.TaskPrivateData{UpstreamTaskID: parsed.UpstreamTaskID},
}, "")
require.NoError(t, err) require.NoError(t, err)
queryBody, err := io.ReadAll(queryResp.Body) queryBody, err := io.ReadAll(queryResp.Body)
require.NoError(t, err) require.NoError(t, err)
require.NoError(t, queryResp.Body.Close()) require.NoError(t, queryResp.Body.Close())
result, err := adaptor.ParseTaskResult(queryBody) result, err := adaptor.ParseTaskResult(&model.Task{}, queryResp, queryBody)
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, "SUCCESS", result.Status) assert.Equal(t, "SUCCESS", result.Status)
assert.Equal(t, "https://cdn.example/video.mp4", result.Url) assert.Equal(t, "https://cdn.example/video.mp4", result.Url)
...@@ -866,7 +869,7 @@ export function extractUsageOnComplete(task, result, body) { return (body || {}) ...@@ -866,7 +869,7 @@ export function extractUsageOnComplete(task, result, body) { return (body || {})
adaptor, _, _ := newRequest(t, map[string]any{}) adaptor, _, _ := newRequest(t, map[string]any{})
body, marshalErr := common.Marshal(map[string]any{"completionUsage": testCase.usage}) body, marshalErr := common.Marshal(map[string]any{"completionUsage": testCase.usage})
require.NoError(t, marshalErr) require.NoError(t, marshalErr)
result, parseErr := adaptor.ParseTaskResult(body) result, parseErr := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, body)
require.NoError(t, parseErr) require.NoError(t, parseErr)
assert.Nil(t, result.UsageFacts) assert.Nil(t, result.UsageFacts)
assert.Zero(t, result.TotalTokens) assert.Zero(t, result.TotalTokens)
...@@ -877,7 +880,7 @@ export function extractUsageOnComplete(task, result, body) { return (body || {}) ...@@ -877,7 +880,7 @@ export function extractUsageOnComplete(task, result, body) { return (body || {})
adaptor, _, _ := newRequest(t, map[string]any{}) adaptor, _, _ := newRequest(t, map[string]any{})
body, err := common.Marshal(map[string]any{"completionUsage": map[string]any{"tokens": 500000}}) body, err := common.Marshal(map[string]any{"completionUsage": map[string]any{"tokens": 500000}})
require.NoError(t, err) require.NoError(t, err)
result, err := adaptor.ParseTaskResult(body) result, err := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, body)
require.NoError(t, err) require.NoError(t, err)
assert.EqualValues(t, 500000, result.UsageFacts["tokens"]) assert.EqualValues(t, 500000, result.UsageFacts["tokens"])
}) })
...@@ -916,7 +919,7 @@ export function extractUsageOnComplete() { return {units: 3.5}; } ...@@ -916,7 +919,7 @@ export function extractUsageOnComplete() { return {units: 3.5}; }
body, err := common.Marshal(map[string]any{}) body, err := common.Marshal(map[string]any{})
require.NoError(t, err) require.NoError(t, err)
result, err := adaptor.ParseTaskResult(body) result, err := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, body)
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, 3.5, result.UsageFacts["units"]) assert.Equal(t, 3.5, result.UsageFacts["units"])
}) })
...@@ -925,7 +928,7 @@ export function extractUsageOnComplete() { return {units: 3.5}; } ...@@ -925,7 +928,7 @@ export function extractUsageOnComplete() { return {units: 3.5}; }
adaptor, _, _ := newRequest(t, map[string]any{}) adaptor, _, _ := newRequest(t, map[string]any{})
body, err := common.Marshal(map[string]any{"completionUsage": map[string]any{"upstreamUnits": 5000}}) body, err := common.Marshal(map[string]any{"completionUsage": map[string]any{"upstreamUnits": 5000}})
require.NoError(t, err) require.NoError(t, err)
result, err := adaptor.ParseTaskResult(body) result, err := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, body)
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, 5000, result.TotalTokens) assert.Equal(t, 5000, result.TotalTokens)
assert.EqualValues(t, 5000, result.UsageFacts["upstreamUnits"]) assert.EqualValues(t, 5000, result.UsageFacts["upstreamUnits"])
...@@ -988,7 +991,7 @@ export function parseTaskResult(ctx, body) { return {status: "SUCCESS", completi ...@@ -988,7 +991,7 @@ export function parseTaskResult(ctx, body) { return {status: "SUCCESS", completi
require.NoError(t, err) require.NoError(t, err)
adaptor := New(plugin) adaptor := New(plugin)
result, err := adaptor.ParseTaskResult([]byte(`{"completion":13,"total":17}`)) result, err := adaptor.ParseTaskResult(&model.Task{}, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}, []byte(`{"completion":13,"total":17}`))
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, 13, result.CompletionTokens) assert.Equal(t, 13, result.CompletionTokens)
assert.Equal(t, 17, result.TotalTokens) assert.Equal(t, 17, result.TotalTokens)
...@@ -1093,7 +1096,7 @@ export function buildSubmitRequest(ctx) { return { url: ctx.baseUrl + "/submit", ...@@ -1093,7 +1096,7 @@ export function buildSubmitRequest(ctx) { return { url: ctx.baseUrl + "/submit",
export function parseSubmitResponse(ctx, resp) { return { taskId: resp.body.id }; } export function parseSubmitResponse(ctx, resp) { return { taskId: resp.body.id }; }
export function buildQueryRequest(ctx) { return { url: ctx.baseUrl + "/tasks/" + ctx.taskId }; } export function buildQueryRequest(ctx) { return { url: ctx.baseUrl + "/tasks/" + ctx.taskId }; }
export function parseTaskResult(ctx, body) { return { taskId: body.id, status: "SUCCESS" }; } export function parseTaskResult(ctx, body) { return { taskId: body.id, status: "SUCCESS" }; }
export function buildBatchQueryRequest(ctx, taskIds) { return { url: ctx.baseUrl + "/batch", method: "POST", headers: { "X-Plugin": "batch" }, body: { ids: taskIds } }; } export function buildBatchQueryRequest(ctx, tasks) { return { url: ctx.baseUrl + "/batch", method: "POST", headers: { "X-Plugin": "batch" }, body: { ids: (tasks || []).map(function (task) { return task.taskId; }) } }; }
export function parseBatchResult(ctx, body) { export function parseBatchResult(ctx, body) {
return body.items.map(function (item) { return body.items.map(function (item) {
return { taskId: item.id, action: item.action, status: item.status, progress: item.progress, url: (item.urls || [])[0] || "", finishTime: item.finish || 0, data: item }; return { taskId: item.id, action: item.action, status: item.status, progress: item.progress, url: (item.urls || [])[0] || "", finishTime: item.finish || 0, data: item };
...@@ -1128,13 +1131,17 @@ func TestTaskAdaptorBatchBridge(t *testing.T) { ...@@ -1128,13 +1131,17 @@ func TestTaskAdaptorBatchBridge(t *testing.T) {
adaptor := New(plugin) adaptor := New(plugin)
require.Equal(t, "batch", adaptor.FetchMode()) require.Equal(t, "batch", adaptor.FetchMode())
resp, err := adaptor.FetchBatchTasks(server.URL, "secret", []string{"task-a", "task-b"}, "") tasks := []*model.Task{
{PrivateData: model.TaskPrivateData{UpstreamTaskID: "task-a"}},
{PrivateData: model.TaskPrivateData{UpstreamTaskID: "task-b"}},
}
resp, err := adaptor.FetchBatchTasks(server.URL, "secret", tasks, "")
require.NoError(t, err) require.NoError(t, err)
defer resp.Body.Close() defer resp.Body.Close()
payload, err := io.ReadAll(resp.Body) payload, err := io.ReadAll(resp.Body)
require.NoError(t, err) require.NoError(t, err)
results, err := adaptor.ParseBatchResult(payload) results, err := adaptor.ParseBatchResult(tasks, resp, payload)
require.NoError(t, err) require.NoError(t, err)
require.Len(t, results, 2, "entry without taskId must be skipped") require.Len(t, results, 2, "entry without taskId must be skipped")
...@@ -1233,26 +1240,176 @@ export function parseTaskResult(){return {status:"SUCCESS"}} ...@@ -1233,26 +1240,176 @@ export function parseTaskResult(){return {status:"SUCCESS"}}
testCases := []struct { testCases := []struct {
name string name string
body map[string]any task *model.Task
want string want string
}{ }{
{ {
name: "mapped model", name: "mapped model",
body: map[string]any{"task_id": "t1", "model": "alias", "upstream_model": "declared-model"}, task: &model.Task{
Properties: model.Properties{OriginModelName: "alias", UpstreamModelName: "declared-model"},
PrivateData: model.TaskPrivateData{UpstreamTaskID: "t1"},
},
want: "/tasks/alias/declared-model/t1", want: "/tasks/alias/declared-model/t1",
}, },
{ {
name: "unmapped model falls back to the origin name", name: "unmapped model falls back to the origin name",
body: map[string]any{"task_id": "t1", "model": "alias"}, task: &model.Task{
Properties: model.Properties{OriginModelName: "alias"},
PrivateData: model.TaskPrivateData{UpstreamTaskID: "t1"},
},
want: "/tasks/alias/alias/t1", want: "/tasks/alias/alias/t1",
}, },
} }
for _, testCase := range testCases { for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) { t.Run(testCase.name, func(t *testing.T) {
resp, fetchErr := adaptor.FetchTask(server.URL, "secret", testCase.body, "") resp, fetchErr := adaptor.FetchTask(server.URL, "secret", testCase.task, "")
require.NoError(t, fetchErr) require.NoError(t, fetchErr)
require.NoError(t, resp.Body.Close()) require.NoError(t, resp.Body.Close())
assert.Equal(t, testCase.want, requested) assert.Equal(t, testCase.want, requested)
}) })
} }
} }
func TestTaskAdaptorQueryContextOmitsRequestBody(t *testing.T) {
service.InitHttpClient()
var captured map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.NoError(t, common.DecodeJson(r.Body, &captured))
_, _ = w.Write([]byte(`{"ok":true}`))
}))
defer server.Close()
source := `
export const meta = {apiVersion:1,key:"query-ctx",name:"Query Ctx",version:"1.0.0",author:{name:"Test"},models:["alias"],fetchMode:"per_task"};
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit"}}
export function parseSubmitResponse(){return {taskId:"1",state:{req_key:"from-submit"}}}
export function buildQueryRequest(ctx){
return {url:ctx.baseUrl+"/query",method:"POST",body:{
keys: Object.keys(ctx).sort(),
taskId: ctx.taskId,
publicTaskId: ctx.publicTaskId,
action: ctx.action,
model: ctx.model,
upstreamModel: ctx.upstreamModel,
data: ctx.data,
state: ctx.state,
hasRequestBody: Object.prototype.hasOwnProperty.call(ctx, "requestBody")
}};
}
export function parseTaskResult(ctx, body, response){
return {status:"IN_PROGRESS",reason:String(response && response.status),url:ctx.taskId,state:{round:"poll"}};
}
`
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
require.NoError(t, err)
adaptor := New(plugin)
adaptor.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: server.URL, ApiKey: "secret"}})
task := &model.Task{
TaskID: "task_public",
Action: constant.TaskActionImageToVideo,
Properties: model.Properties{
OriginModelName: "alias",
UpstreamModelName: "declared",
},
Data: []byte(`{"snapshot":true}`),
PrivateData: model.TaskPrivateData{
UpstreamTaskID: "upstream-1",
PluginState: []byte(`{"req_key":"kept"}`),
},
}
resp, err := adaptor.FetchTask(server.URL, "secret", task, "")
require.NoError(t, err)
require.NoError(t, resp.Body.Close())
assert.Equal(t, "upstream-1", captured["taskId"])
assert.Equal(t, "task_public", captured["publicTaskId"])
assert.Equal(t, constant.NormalizeTaskAction(constant.TaskActionImageToVideo), captured["action"])
assert.Equal(t, "alias", captured["model"])
assert.Equal(t, "declared", captured["upstreamModel"])
assert.Equal(t, map[string]any{"snapshot": true}, captured["data"])
assert.Equal(t, map[string]any{"req_key": "kept"}, captured["state"])
assert.Equal(t, false, captured["hasRequestBody"])
keys, ok := captured["keys"].([]any)
require.True(t, ok)
assert.NotContains(t, keys, "requestBody")
result, err := adaptor.ParseTaskResult(task, &http.Response{StatusCode: http.StatusTeapot, Header: make(http.Header)}, []byte(`{"ok":true}`))
require.NoError(t, err)
assert.Equal(t, "IN_PROGRESS", result.Status)
assert.Equal(t, "418", result.Reason)
assert.Equal(t, "upstream-1", result.Url)
assert.JSONEq(t, `{"round":"poll"}`, string(result.PluginState))
}
func TestTaskAdaptorParseSubmitResponsePersistsState(t *testing.T) {
service.InitHttpClient()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(`{"id":"upstream-1"}`))
}))
defer server.Close()
source := `
export const meta = {apiVersion:1,key:"submit-state",name:"Submit State",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task"};
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit",method:"POST",body:{}}}
export function parseSubmitResponse(){return {taskId:"upstream-1",state:{req_key:"from-submit"}}}
export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"}}
export function parseTaskResult(){return {status:"SUCCESS"}}
`
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
require.NoError(t, err)
adaptor := New(plugin)
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: server.URL, ApiKey: "secret"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
c.Set("task_request", relaycommon.TaskSubmitReq{Prompt: "hello"})
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
body, err := adaptor.BuildRequestBody(c, info)
require.NoError(t, err)
resp, err := adaptor.DoRequest(c, info, body)
require.NoError(t, err)
parsed, taskErr := adaptor.ParseResponse(c, resp, info)
require.Nil(t, taskErr)
require.NotNil(t, parsed)
assert.JSONEq(t, `{"req_key":"from-submit"}`, string(parsed.PluginState))
}
func TestTaskAdaptorBatchQueryReceivesTaskObjects(t *testing.T) {
service.InitHttpClient()
var captured map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.NoError(t, common.DecodeJson(r.Body, &captured))
_, _ = w.Write([]byte(`{"items":[]}`))
}))
defer server.Close()
source := `
export const meta = {apiVersion:1,key:"batch-ctx",name:"Batch Ctx",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"batch"};
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit"}}
export function parseSubmitResponse(){return {taskId:"1"}}
export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/q"}}
export function parseTaskResult(){return {status:"SUCCESS"}}
export function buildBatchQueryRequest(ctx, tasks){
return {url:ctx.baseUrl+"/batch",method:"POST",body:{
ids: (tasks||[]).map(function(task){return task.taskId;}),
models: (tasks||[]).map(function(task){return task.model;}),
hasRequestBody: (tasks||[]).some(function(task){return Object.prototype.hasOwnProperty.call(task,"requestBody");})
}};
}
export function parseBatchResult(){return [];}
`
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
require.NoError(t, err)
adaptor := New(plugin)
tasks := []*model.Task{
{Properties: model.Properties{OriginModelName: "model-a"}, PrivateData: model.TaskPrivateData{UpstreamTaskID: "task-a"}},
{Properties: model.Properties{OriginModelName: "model-b"}, PrivateData: model.TaskPrivateData{UpstreamTaskID: "task-b"}},
}
resp, err := adaptor.FetchBatchTasks(server.URL, "secret", tasks, "")
require.NoError(t, err)
require.NoError(t, resp.Body.Close())
assert.Equal(t, []any{"task-a", "task-b"}, captured["ids"])
assert.Equal(t, []any{"model-a", "model-b"}, captured["models"])
assert.Equal(t, false, captured["hasRequestBody"])
}
relay/common/relay_info.go
...@@ -1028,6 +1028,7 @@ type TaskInfo struct { ...@@ -1028,6 +1028,7 @@ type TaskInfo struct {
CompletionTokens int `json:"completion_tokens,omitempty"` // 用于按倍率计费 CompletionTokens int `json:"completion_tokens,omitempty"` // 用于按倍率计费
TotalTokens int `json:"total_tokens,omitempty"` // 用于按倍率计费 TotalTokens int `json:"total_tokens,omitempty"` // 用于按倍率计费
UsageFacts map[string]any `json:"usage_facts,omitempty"` UsageFacts map[string]any `json:"usage_facts,omitempty"`
PluginState json.RawMessage `json:"plugin_state,omitempty"`
} }
func FailTaskInfo(reason string) *TaskInfo { func FailTaskInfo(reason string) *TaskInfo {
......
relay/relay_task.go
...@@ -32,6 +32,7 @@ type TaskSubmitResult struct { ...@@ -32,6 +32,7 @@ type TaskSubmitResult struct {
Platform constant.TaskPlatform Platform constant.TaskPlatform
Quota int Quota int
Immediate *relaycommon.TaskInfo Immediate *relaycommon.TaskInfo
PluginState []byte
//PerCallPrice types.PriceData //PerCallPrice types.PriceData
} }
...@@ -381,6 +382,7 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe ...@@ -381,6 +382,7 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
Platform: platform, Platform: platform,
Quota: finalQuota, Quota: finalQuota,
Immediate: parsed.Immediate, Immediate: parsed.Immediate,
PluginState: parsed.PluginState,
}, nil }, nil
} }
...@@ -517,12 +519,7 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte { ...@@ -517,12 +519,7 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte {
return nil return nil
} }
resp, err := adaptor.FetchTask(baseURL, channelModel.Key, map[string]any{ resp, err := adaptor.FetchTask(baseURL, channelModel.Key, task, proxy)
"task_id": task.GetUpstreamTaskID(),
"action": constant.NormalizeTaskAction(task.Action),
"model": task.Properties.OriginModelName,
"upstream_model": task.Properties.UpstreamModelName,
}, proxy)
if err != nil || resp == nil { if err != nil || resp == nil {
return nil return nil
} }
...@@ -532,7 +529,7 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte { ...@@ -532,7 +529,7 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte {
return nil return nil
} }
ti, err := adaptor.ParseTaskResult(body) ti, err := adaptor.ParseTaskResult(task, resp, body)
if err != nil || ti == nil { if err != nil || ti == nil {
return nil return nil
} }
......
service/task_billing_test.go
...@@ -1343,10 +1343,12 @@ type mockAdaptor struct { ...@@ -1343,10 +1343,12 @@ type mockAdaptor struct {
} }
func (m *mockAdaptor) Init(_ *relaycommon.RelayInfo) {} func (m *mockAdaptor) Init(_ *relaycommon.RelayInfo) {}
func (m *mockAdaptor) FetchTask(string, string, map[string]any, string) (*http.Response, error) { func (m *mockAdaptor) FetchTask(string, string, *model.Task, string) (*http.Response, error) {
return nil, nil
}
func (m *mockAdaptor) ParseTaskResult(*model.Task, *http.Response, []byte) (*relaycommon.TaskInfo, error) {
return nil, nil return nil, nil
} }
func (m *mockAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) { return nil, nil }
func (m *mockAdaptor) AdjustBillingOnComplete(_ *model.Task, _ *relaycommon.TaskInfo) int { func (m *mockAdaptor) AdjustBillingOnComplete(_ *model.Task, _ *relaycommon.TaskInfo) int {
return m.adjustReturn return m.adjustReturn
} }
......
service/task_plugin_view_test.go
...@@ -62,3 +62,24 @@ func TestBuildTaskPluginViewRewritesOnlyStructuredTaskIDFields(t *testing.T) { ...@@ -62,3 +62,24 @@ func TestBuildTaskPluginViewRewritesOnlyStructuredTaskIDFields(t *testing.T) {
assert.Equal(t, privateTaskID, nested[1]) assert.Equal(t, privateTaskID, nested[1])
} }
func TestBuildTaskPluginViewOmitsPrivatePollState(t *testing.T) {
task := &model.Task{
TaskID: "task_public_view",
Data: []byte(`{"ok":true}`),
PrivateData: model.TaskPrivateData{
PluginState: []byte(`{"req_key":"secret"}`),
PollFailures: 7,
},
}
view, err := BuildTaskPluginView(task)
require.NoError(t, err)
encoded, err := common.Marshal(view)
require.NoError(t, err)
var payload map[string]any
require.NoError(t, common.Unmarshal(encoded, &payload))
assert.NotContains(t, payload, "plugin_state")
assert.NotContains(t, payload, "poll_failures")
assert.NotContains(t, payload, "private_data")
}
service/task_polling.go
...@@ -18,7 +18,6 @@ import ( ...@@ -18,7 +18,6 @@ import (
"github.com/QuantumNous/new-api/pkg/billingexpr" "github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon" "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/bytedance/gopkg/util/gopool" "github.com/bytedance/gopkg/util/gopool"
"github.com/samber/lo" "github.com/samber/lo"
...@@ -27,8 +26,8 @@ import ( ...@@ -27,8 +26,8 @@ import (
// TaskPollingAdaptor 定义轮询所需的最小适配器接口，避免 service -> relay 的循环依赖 // TaskPollingAdaptor 定义轮询所需的最小适配器接口，避免 service -> relay 的循环依赖
type TaskPollingAdaptor interface { type TaskPollingAdaptor interface {
Init(info *relaycommon.RelayInfo) Init(info *relaycommon.RelayInfo)
FetchTask(baseURL string, key string, body map[string]any, proxy string) (*http.Response, error) FetchTask(baseURL string, key string, task *model.Task, proxy string) (*http.Response, error)
ParseTaskResult(body []byte) (*relaycommon.TaskInfo, error) ParseTaskResult(task *model.Task, resp *http.Response, body []byte) (*relaycommon.TaskInfo, error)
// AdjustBillingOnComplete 在任务到达终态（成功/失败）时由轮询循环调用。 // AdjustBillingOnComplete 在任务到达终态（成功/失败）时由轮询循环调用。
// 返回正数触发差额结算（补扣/退还），返回 0 保持预扣费金额不变。 // 返回正数触发差额结算（补扣/退还），返回 0 保持预扣费金额不变。
AdjustBillingOnComplete(task *model.Task, taskResult *relaycommon.TaskInfo) int AdjustBillingOnComplete(task *model.Task, taskResult *relaycommon.TaskInfo) int
...@@ -37,10 +36,21 @@ type TaskPollingAdaptor interface { ...@@ -37,10 +36,21 @@ type TaskPollingAdaptor interface {
type BatchTaskPollingAdaptor interface { type BatchTaskPollingAdaptor interface {
TaskPollingAdaptor TaskPollingAdaptor
FetchMode() string FetchMode() string
FetchBatchTasks(baseURL, key string, taskIDs []string, proxy string) (*http.Response, error) FetchBatchTasks(baseURL, key string, tasks []*model.Task, proxy string) (*http.Response, error)
ParseBatchResult(body []byte) (map[string]*BatchTaskResult, error) ParseBatchResult(tasks []*model.Task, resp *http.Response, body []byte) (map[string]*BatchTaskResult, error)
} }
const (
pollClassOK = "ok"
pollClassOtherClient = "other_client"
pollClassNotFound = "not_found"
pollClassAuth = "auth"
pollClassTransient = "transient"
pollClassUnrecognized = "unrecognized"
pollClassHookError = "hook_error"
pollClassTransport = "transport_error"
)
type BatchTaskResult struct { type BatchTaskResult struct {
TaskInfo relaycommon.TaskInfo TaskInfo relaycommon.TaskInfo
Action string Action string
...@@ -258,24 +268,39 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan ...@@ -258,24 +268,39 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
if baseURL == "" { if baseURL == "" {
baseURL = constant.GetChannelBaseURL(ch.Type) baseURL = constant.GetChannelBaseURL(ch.Type)
} }
resp, err := adaptor.FetchBatchTasks(baseURL, ch.Key, taskIds, proxy) tasks := make([]*model.Task, 0, len(taskIds))
for _, upstreamID := range taskIds {
if task := taskM[upstreamID]; task != nil {
tasks = append(tasks, task)
}
}
info := &relaycommon.RelayInfo{}
info.ChannelMeta = &relaycommon.ChannelMeta{ChannelBaseUrl: baseURL}
info.ApiKey = ch.Key
adaptor.Init(info)
resp, err := adaptor.FetchBatchTasks(baseURL, ch.Key, tasks, proxy)
if err != nil { if err != nil {
common.SysLog(fmt.Sprintf("Get Task Do req error: %v", err)) common.SysLog(fmt.Sprintf("Get Task Do req error: %v", err))
return err return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassTransport, 0, err.Error())
}
if resp.StatusCode != http.StatusOK {
logger.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode))
return fmt.Errorf("Get Task status code: %d", resp.StatusCode)
} }
defer resp.Body.Close() defer resp.Body.Close()
responseBody, err := io.ReadAll(resp.Body) responseBody, err := io.ReadAll(resp.Body)
if err != nil { if err != nil {
common.SysLog(fmt.Sprintf("Get Suno Task parse body error: %v", err)) common.SysLog(fmt.Sprintf("Get Suno Task parse body error: %v", err))
return err return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassTransport, resp.StatusCode, err.Error())
} }
responseItems, err := adaptor.ParseBatchResult(responseBody) switch classifyPollHTTP(resp.StatusCode) {
case pollClassNotFound:
return failTasksFromPoll(ctx, adaptor, tasks, fmt.Sprintf("upstream task not found (HTTP %d)", resp.StatusCode))
case pollClassAuth:
logger.LogWarn(ctx, fmt.Sprintf("task poll auth failure channel_id=%d http=%d", channelId, resp.StatusCode))
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassAuth, resp.StatusCode, "")
case pollClassTransient:
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassTransient, resp.StatusCode, "")
}
responseItems, err := adaptor.ParseBatchResult(tasks, resp, responseBody)
if err != nil { if err != nil {
return fmt.Errorf("parse batch result: %w", err) return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassHookError, resp.StatusCode, err.Error())
} }
for upstreamID, responseItem := range responseItems { for upstreamID, responseItem := range responseItems {
if ctx.Err() != nil { if ctx.Err() != nil {
...@@ -287,7 +312,27 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan ...@@ -287,7 +312,27 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
continue continue
} }
snap := task.Snapshot() snap := task.Snapshot()
task.Status = lo.If(model.TaskStatus(responseItem.TaskInfo.Status) != "", model.TaskStatus(responseItem.TaskInfo.Status)).Else(task.Status) httpClass := classifyPollHTTP(resp.StatusCode)
parsedStatus := model.TaskStatus(responseItem.TaskInfo.Status)
if parsedStatus == model.TaskStatusUnknown || parsedStatus == "" || !knownPollStatus(parsedStatus) {
if err := recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, responseItem.TaskInfo.Reason); err != nil {
common.SysLog("UpdateSunoTask task error: " + err.Error())
}
continue
}
if httpClass == pollClassOtherClient && isNonTerminalPollStatus(parsedStatus) {
if err := recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, responseItem.TaskInfo.Reason); err != nil {
common.SysLog("UpdateSunoTask task error: " + err.Error())
}
continue
}
if isNonTerminalPollStatus(parsedStatus) {
task.PrivateData.PollFailures = 0
}
if len(responseItem.TaskInfo.PluginState) > 0 {
task.PrivateData.PluginState = responseItem.TaskInfo.PluginState
}
task.Status = lo.If(parsedStatus != "", parsedStatus).Else(task.Status)
task.FailReason = lo.If(responseItem.TaskInfo.Reason != "", responseItem.TaskInfo.Reason).Else(task.FailReason) task.FailReason = lo.If(responseItem.TaskInfo.Reason != "", responseItem.TaskInfo.Reason).Else(task.FailReason)
task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime) task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime)
task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime) task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime)
...@@ -447,24 +492,28 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * ...@@ -447,24 +492,28 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
if privateData.Key != "" { if privateData.Key != "" {
key = privateData.Key key = privateData.Key
} }
resp, err := adaptor.FetchTask(baseURL, key, map[string]any{ snap := task.Snapshot()
"task_id": task.GetUpstreamTaskID(), resp, err := adaptor.FetchTask(baseURL, key, task, proxy)
"action": constant.NormalizeTaskAction(task.Action),
"model": task.Properties.OriginModelName,
"upstream_model": task.Properties.UpstreamModelName,
}, proxy)
if err != nil { if err != nil {
return fmt.Errorf("fetchTask failed for task %s: %w", taskId, err) return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassTransport, 0, err.Error())
} }
defer resp.Body.Close() defer resp.Body.Close()
responseBody, err := io.ReadAll(resp.Body) responseBody, err := io.ReadAll(resp.Body)
if err != nil { if err != nil {
return fmt.Errorf("readAll failed for task %s: %w", taskId, err) return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassTransport, resp.StatusCode, err.Error())
} }
logger.LogDebug(ctx, "updateVideoSingleTask response: %s", responseBody) logger.LogDebug(ctx, "updateVideoSingleTask response: %s", responseBody)
snap := task.Snapshot() switch classifyPollHTTP(resp.StatusCode) {
case pollClassNotFound:
return failTaskFromPoll(ctx, adaptor, task, snap.Status, fmt.Sprintf("upstream task not found (HTTP %d)", resp.StatusCode))
case pollClassAuth:
logger.LogWarn(ctx, fmt.Sprintf("task poll auth failure channel_id=%d task=%s http=%d", ch.Id, task.TaskID, resp.StatusCode))
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassAuth, resp.StatusCode, "")
case pollClassTransient:
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassTransient, resp.StatusCode, "")
}
taskResult := &relaycommon.TaskInfo{} taskResult := &relaycommon.TaskInfo{}
// try parse as New API response format // try parse as New API response format
...@@ -478,41 +527,33 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * ...@@ -478,41 +527,33 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
taskResult.Progress = t.Progress taskResult.Progress = t.Progress
taskResult.Reason = t.FailReason taskResult.Reason = t.FailReason
task.Data = t.Data task.Data = t.Data
} else if taskResult, err = adaptor.ParseTaskResult(responseBody); err != nil { } else if taskResult, err = adaptor.ParseTaskResult(task, resp, responseBody); err != nil {
return fmt.Errorf("parseTaskResult failed for task %s: %w", taskId, err) return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassHookError, resp.StatusCode, err.Error())
} }
task.Data = redactVideoResponseBody(responseBody)
logger.LogDebug(ctx, "updateVideoSingleTask taskResult: %+v", taskResult) logger.LogDebug(ctx, "updateVideoSingleTask taskResult: %+v", taskResult)
now := time.Now().Unix() parsedStatus := model.TaskStatus(taskResult.Status)
if taskResult.Status == "" { if parsedStatus == model.TaskStatusUnknown || parsedStatus == "" || !knownPollStatus(parsedStatus) {
//taskResult = relaycommon.FailTaskInfo("upstream returned empty status") return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, unrecognizedPollDetail(taskResult.Reason, responseBody))
errorResult := &dto.GeneralErrorResponse{}
if err = common.Unmarshal(responseBody, &errorResult); err == nil {
openaiError := errorResult.TryToOpenAIError()
if openaiError != nil {
// 返回规范的 OpenAI 错误格式，提取错误信息，判断错误是否为任务失败
if openaiError.Code == "429" {
// 429 错误通常表示请求过多或速率限制，暂时不认为是任务失败，保持原状态等待下一轮轮询
return nil
} }
if classifyPollHTTP(resp.StatusCode) == pollClassOtherClient && isNonTerminalPollStatus(parsedStatus) {
// 其他错误认为是任务失败，记录错误信息并更新任务状态 return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, unrecognizedPollDetail(taskResult.Reason, responseBody))
taskResult = relaycommon.FailTaskInfo("upstream returned error")
} else {
// unknown error format, log original response
logger.LogError(ctx, fmt.Sprintf("Task %s returned empty status with unrecognized error format, response: %s", taskId, string(responseBody)))
taskResult = relaycommon.FailTaskInfo("upstream returned unrecognized message")
} }
task.Data = redactVideoResponseBody(responseBody)
if len(taskResult.PluginState) > 0 {
task.PrivateData.PluginState = taskResult.PluginState
} }
if isNonTerminalPollStatus(parsedStatus) {
task.PrivateData.PollFailures = 0
} }
now := time.Now().Unix()
shouldFinalizeBilling := false shouldFinalizeBilling := false
task.Status = model.TaskStatus(taskResult.Status) task.Status = parsedStatus
switch taskResult.Status { switch parsedStatus {
case model.TaskStatusSubmitted: case model.TaskStatusSubmitted:
task.Progress = taskcommon.ProgressSubmitted task.Progress = taskcommon.ProgressSubmitted
case model.TaskStatusQueued: case model.TaskStatusQueued:
...@@ -549,8 +590,6 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * ...@@ -549,8 +590,6 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
logger.LogInfo(ctx, fmt.Sprintf("Task %s failed: %s", task.TaskID, task.FailReason)) logger.LogInfo(ctx, fmt.Sprintf("Task %s failed: %s", task.TaskID, task.FailReason))
taskResult.Progress = taskcommon.ProgressComplete taskResult.Progress = taskcommon.ProgressComplete
shouldFinalizeBilling = true shouldFinalizeBilling = true
default:
return fmt.Errorf("unknown task status %s for task %s", taskResult.Status, task.TaskID)
} }
if taskResult.Progress != "" { if taskResult.Progress != "" {
task.Progress = taskResult.Progress task.Progress = taskResult.Progress
...@@ -670,3 +709,131 @@ func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor ...@@ -670,3 +709,131 @@ func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor
} }
return false return false
} }
func classifyPollHTTP(statusCode int) string {
switch {
case statusCode >= 200 && statusCode < 300:
return pollClassOK
case statusCode == http.StatusNotFound || statusCode == http.StatusGone:
return pollClassNotFound
case statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden:
return pollClassAuth
case statusCode == http.StatusTooManyRequests || statusCode >= 500:
return pollClassTransient
case statusCode >= 400 && statusCode < 500:
return pollClassOtherClient
default:
return pollClassTransient
}
}
func knownPollStatus(status model.TaskStatus) bool {
switch status {
case model.TaskStatusNotStart, model.TaskStatusSubmitted, model.TaskStatusQueued, model.TaskStatusInProgress, model.TaskStatusSuccess, model.TaskStatusFailure:
return true
default:
return false
}
}
func isNonTerminalPollStatus(status model.TaskStatus) bool {
switch status {
case model.TaskStatusNotStart, model.TaskStatusSubmitted, model.TaskStatusQueued, model.TaskStatusInProgress:
return true
default:
return false
}
}
func pollFailureReason(class string, statusCode int, detail string) string {
reason := fmt.Sprintf("poll failed: %s", class)
if statusCode > 0 {
reason = fmt.Sprintf("poll failed: %s (HTTP %d)", class, statusCode)
}
if detail != "" {
reason = reason + ": " + detail
}
return reason
}
// unrecognizedPollDetail pairs the plugin's reason with a bounded copy of the
// upstream body so the WARN line is enough to diagnose a parser gap.
func unrecognizedPollDetail(reason string, body []byte) string {
const maxBodyChars = 512
redacted := string(redactVideoResponseBody(body))
if len(redacted) > maxBodyChars {
redacted = redacted[:maxBodyChars] + "…"
}
if strings.TrimSpace(reason) == "" {
return "body=" + redacted
}
return reason + "; body=" + redacted
}
func recordPollFailure(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, fromStatus model.TaskStatus, class string, statusCode int, detail string) error {
task.PrivateData.PollFailures++
if class == pollClassUnrecognized || class == pollClassHookError {
// The redacted body is intentionally not persisted to Task.Data on these
// paths, so the WARN line is the only operator-visible copy of what the
// plugin could not interpret.
logger.LogWarn(ctx, fmt.Sprintf("task %s poll %s (failures=%d, http=%d): %s", task.TaskID, class, task.PrivateData.PollFailures, statusCode, detail))
}
// TASK_POLL_MAX_FAILURES <= 0 disables the consecutive-failure cutoff, matching
// TASK_TIMEOUT_MINUTES semantics; the 24h sweep remains the only backstop.
if constant.TaskPollMaxFailures > 0 && task.PrivateData.PollFailures >= constant.TaskPollMaxFailures {
return failTaskFromPoll(ctx, adaptor, task, fromStatus, pollFailureReason(class, statusCode, detail))
}
if _, err := task.UpdateWithStatus(fromStatus); err != nil {
return err
}
return nil
}
func recordPollFailureForTasks(ctx context.Context, adaptor TaskPollingAdaptor, tasks []*model.Task, class string, statusCode int, detail string) error {
var firstErr error
for _, task := range tasks {
if task == nil {
continue
}
if err := recordPollFailure(ctx, adaptor, task, task.Status, class, statusCode, detail); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}
func failTaskFromPoll(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, fromStatus model.TaskStatus, reason string) error {
now := time.Now().Unix()
task.Status = model.TaskStatusFailure
task.Progress = taskcommon.ProgressComplete
if task.FinishTime == 0 {
task.FinishTime = now
}
task.FailReason = reason
won, err := task.UpdateWithStatus(fromStatus)
if err != nil {
return err
}
if !won {
return nil
}
taskResult := relaycommon.FailTaskInfo(reason)
billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
if !billingSettled && task.Quota != 0 {
RefundTaskQuota(ctx, task, reason)
}
return nil
}
func failTasksFromPoll(ctx context.Context, adaptor TaskPollingAdaptor, tasks []*model.Task, reason string) error {
var firstErr error
for _, task := range tasks {
if task == nil {
continue
}
if err := failTaskFromPoll(ctx, adaptor, task, task.Status, reason); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}
service/task_polling_test.go
...@@ -40,12 +40,15 @@ type batchPollingAdaptor struct { ...@@ -40,12 +40,15 @@ type batchPollingAdaptor struct {
} }
func (a *batchPollingAdaptor) FetchMode() string { return "batch" } func (a *batchPollingAdaptor) FetchMode() string { return "batch" }
func (a *batchPollingAdaptor) FetchBatchTasks(_ string, _ string, taskIDs []string, _ string) (*http.Response, error) { func (a *batchPollingAdaptor) FetchBatchTasks(_ string, _ string, tasks []*model.Task, _ string) (*http.Response, error) {
a.batchCalls++ a.batchCalls++
a.batchIDs = append([]string(nil), taskIDs...) a.batchIDs = a.batchIDs[:0]
for _, task := range tasks {
a.batchIDs = append(a.batchIDs, task.GetUpstreamTaskID())
}
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader([]byte(`{}`)))}, nil return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader([]byte(`{}`)))}, nil
} }
func (a *batchPollingAdaptor) ParseBatchResult([]byte) (map[string]*BatchTaskResult, error) { func (a *batchPollingAdaptor) ParseBatchResult(_ []*model.Task, _ *http.Response, _ []byte) (map[string]*BatchTaskResult, error) {
if a.results != nil { if a.results != nil {
return a.results, nil return a.results, nil
} }
...@@ -58,8 +61,11 @@ func (a *batchPollingAdaptor) ParseBatchResult([]byte) (map[string]*BatchTaskRes ...@@ -58,8 +61,11 @@ func (a *batchPollingAdaptor) ParseBatchResult([]byte) (map[string]*BatchTaskRes
func (a *taskPollingFetchAdaptor) Init(_ *relaycommon.RelayInfo) {} func (a *taskPollingFetchAdaptor) Init(_ *relaycommon.RelayInfo) {}
func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, body map[string]any, _ string) (*http.Response, error) { func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, task *model.Task, _ string) (*http.Response, error) {
taskID, _ := body["task_id"].(string) taskID := ""
if task != nil {
taskID = task.GetUpstreamTaskID()
}
if taskID == a.blockTaskID && a.releaseBlock != nil { if taskID == a.blockTaskID && a.releaseBlock != nil {
a.blockOnce.Do(func() { a.blockOnce.Do(func() {
if a.blockStarted != nil { if a.blockStarted != nil {
...@@ -97,7 +103,7 @@ func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, body map[string] ...@@ -97,7 +103,7 @@ func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, body map[string]
}, nil }, nil
} }
func (a *taskPollingFetchAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) { func (a *taskPollingFetchAdaptor) ParseTaskResult(*model.Task, *http.Response, []byte) (*relaycommon.TaskInfo, error) {
return &relaycommon.TaskInfo{Status: model.TaskStatusInProgress}, nil return &relaycommon.TaskInfo{Status: model.TaskStatusInProgress}, nil
} }
...@@ -721,3 +727,273 @@ func TestSweepTimedOutTasksHonorsRefundRolloutBoundary(t *testing.T) { ...@@ -721,3 +727,273 @@ func TestSweepTimedOutTasksHonorsRefundRolloutBoundary(t *testing.T) {
assert.Equal(t, initialQuota+modernTaskQuota, getUserQuota(t, userID)) assert.Equal(t, initialQuota+modernTaskQuota, getUserQuota(t, userID))
assert.Equal(t, int64(1), countLogs(t)) assert.Equal(t, int64(1), countLogs(t))
} }
type scriptedPollingAdaptor struct {
statusCode int
body []byte
fetchErr error
parse *relaycommon.TaskInfo
parseErr error
}
func (a *scriptedPollingAdaptor) Init(*relaycommon.RelayInfo) {}
func (a *scriptedPollingAdaptor) FetchTask(string, string, *model.Task, string) (*http.Response, error) {
if a.fetchErr != nil {
return nil, a.fetchErr
}
code := a.statusCode
if code == 0 {
code = http.StatusOK
}
body := a.body
if body == nil {
body = []byte(`{}`)
}
return &http.Response{StatusCode: code, Body: io.NopCloser(bytes.NewReader(body))}, nil
}
func (a *scriptedPollingAdaptor) ParseTaskResult(*model.Task, *http.Response, []byte) (*relaycommon.TaskInfo, error) {
if a.parseErr != nil {
return nil, a.parseErr
}
if a.parse != nil {
return a.parse, nil
}
return &relaycommon.TaskInfo{Status: model.TaskStatusInProgress}, nil
}
func (a *scriptedPollingAdaptor) AdjustBillingOnComplete(*model.Task, *relaycommon.TaskInfo) int {
return 0
}
type scriptedBatchPollingAdaptor struct {
scriptedPollingAdaptor
results map[string]*BatchTaskResult
}
func (a *scriptedBatchPollingAdaptor) FetchMode() string { return "batch" }
func (a *scriptedBatchPollingAdaptor) FetchBatchTasks(string, string, []*model.Task, string) (*http.Response, error) {
return a.FetchTask("", "", nil, "")
}
func (a *scriptedBatchPollingAdaptor) ParseBatchResult([]*model.Task, *http.Response, []byte) (map[string]*BatchTaskResult, error) {
if a.parseErr != nil {
return nil, a.parseErr
}
return a.results, nil
}
func TestUpdateVideoSingleTaskPollClassification(t *testing.T) {
testCases := []struct {
name string
statusCode int
fetchErr error
parse *relaycommon.TaskInfo
parseErr error
priorFailures int
priorState string
maxFailures int
wantStatus model.TaskStatus
wantFailures int
wantRefund bool
wantReason string
wantState string
wantUnchanged bool
}{
{
name: "404 fails immediately and refunds",
statusCode: http.StatusNotFound,
wantStatus: model.TaskStatusFailure,
wantRefund: true,
wantReason: "upstream task not found (HTTP 404)",
wantUnchanged: false,
},
{
name: "401 increments without changing status",
statusCode: http.StatusUnauthorized,
wantStatus: model.TaskStatusInProgress,
wantFailures: 1,
wantUnchanged: true,
},
{
name: "429 reaches threshold and refunds",
statusCode: http.StatusTooManyRequests,
priorFailures: 2,
maxFailures: 3,
wantStatus: model.TaskStatusFailure,
wantFailures: 3,
wantRefund: true,
wantReason: "poll failed: transient (HTTP 429)",
},
{
name: "UNKNOWN increments",
statusCode: http.StatusOK,
parse: &relaycommon.TaskInfo{Status: model.TaskStatusUnknown, Reason: "weird"},
wantStatus: model.TaskStatusInProgress,
wantFailures: 1,
wantUnchanged: true,
},
{
name: "valid 2xx resets the failure counter",
statusCode: http.StatusOK,
parse: &relaycommon.TaskInfo{Status: model.TaskStatusInProgress},
priorFailures: 5,
wantStatus: model.TaskStatusInProgress,
wantFailures: 0,
},
{
name: "omit state preserves previous plugin state",
statusCode: http.StatusOK,
parse: &relaycommon.TaskInfo{Status: model.TaskStatusInProgress},
priorState: `{"req_key":"keep"}`,
wantStatus: model.TaskStatusInProgress,
wantState: `{"req_key":"keep"}`,
},
{
name: "returned state replaces plugin state",
statusCode: http.StatusOK,
parse: &relaycommon.TaskInfo{
Status: model.TaskStatusInProgress,
PluginState: []byte(`{"req_key":"new"}`),
},
priorState: `{"req_key":"old"}`,
wantStatus: model.TaskStatusInProgress,
wantState: `{"req_key":"new"}`,
},
{
name: "other 4xx non-terminal is unrecognized",
statusCode: http.StatusBadRequest,
parse: &relaycommon.TaskInfo{Status: model.TaskStatusInProgress},
wantStatus: model.TaskStatusInProgress,
wantFailures: 1,
wantUnchanged: true,
},
}
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
truncate(t)
const userID, tokenID, channelID = 510, 510, 510
const initialQuota, preConsumed, tokenRemain = 10_000, 4_000, 7_000
seedUser(t, userID, initialQuota)
seedToken(t, tokenID, userID, "sk-poll-class", tokenRemain)
ch := &model.Channel{Id: channelID, Type: constant.ChannelTypeKling, Name: "poll", Key: "sk-test", Status: common.ChannelStatusEnabled}
if testCase.maxFailures > 0 {
previous := constant.TaskPollMaxFailures
constant.TaskPollMaxFailures = testCase.maxFailures
t.Cleanup(func() { constant.TaskPollMaxFailures = previous })
}
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
task.TaskID = "task_poll_class"
task.PrivateData.UpstreamTaskID = "upstream_poll_class"
task.PrivateData.PollFailures = testCase.priorFailures
if testCase.priorState != "" {
task.PrivateData.PluginState = []byte(testCase.priorState)
}
require.NoError(t, model.DB.Create(task).Error)
adaptor := &scriptedPollingAdaptor{statusCode: testCase.statusCode, fetchErr: testCase.fetchErr, parse: testCase.parse, parseErr: testCase.parseErr}
require.NoError(t, updateVideoSingleTask(context.Background(), adaptor, ch, task.GetUpstreamTaskID(), map[string]*model.Task{
task.GetUpstreamTaskID(): task,
}))
var persisted model.Task
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
assert.EqualValues(t, testCase.wantStatus, persisted.Status)
assert.Equal(t, testCase.wantFailures, persisted.PrivateData.PollFailures)
if testCase.wantUnchanged {
assert.Empty(t, persisted.FailReason)
}
if testCase.wantReason != "" {
assert.Contains(t, persisted.FailReason, testCase.wantReason)
}
if testCase.wantState != "" {
assert.JSONEq(t, testCase.wantState, string(persisted.PrivateData.PluginState))
}
if testCase.wantRefund {
assert.Equal(t, initialQuota+preConsumed, getUserQuota(t, userID))
assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
assert.Zero(t, persisted.Quota)
log := getLastLog(t)
require.NotNil(t, log)
assert.Equal(t, model.LogTypeRefund, log.Type)
} else {
assert.Equal(t, initialQuota, getUserQuota(t, userID))
assert.Equal(t, tokenRemain, getTokenRemainQuota(t, tokenID))
}
})
}
}
func TestUpdateBatchTasksPollClassification(t *testing.T) {
testCases := []struct {
name string
statusCode int
resultStatus model.TaskStatus
wantStatus model.TaskStatus
wantFailures int
wantRefund bool
wantReason string
}{
{
name: "404 fails the batch and refunds",
statusCode: http.StatusNotFound,
wantStatus: model.TaskStatusFailure,
wantRefund: true,
wantReason: "upstream task not found (HTTP 404)",
},
{
name: "401 increments every task",
statusCode: http.StatusUnauthorized,
wantStatus: model.TaskStatusInProgress,
wantFailures: 1,
},
{
name: "UNKNOWN increments",
statusCode: http.StatusOK,
resultStatus: model.TaskStatusUnknown,
wantStatus: model.TaskStatusInProgress,
wantFailures: 1,
},
}
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
truncate(t)
const userID, tokenID, channelID = 610, 610, 610
const initialQuota, preConsumed, tokenRemain = 10_000, 4_000, 7_000
seedUser(t, userID, initialQuota)
seedToken(t, tokenID, userID, "sk-batch-class", tokenRemain)
seedTaskPollingChannel(t, channelID, true)
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
task.TaskID = "task_batch_class"
task.PrivateData.UpstreamTaskID = "upstream_batch_class"
require.NoError(t, model.DB.Create(task).Error)
upstreamID := task.GetUpstreamTaskID()
adaptor := &scriptedBatchPollingAdaptor{
scriptedPollingAdaptor: scriptedPollingAdaptor{statusCode: testCase.statusCode},
}
if testCase.resultStatus != "" {
adaptor.results = map[string]*BatchTaskResult{
upstreamID: {TaskInfo: relaycommon.TaskInfo{TaskID: upstreamID, Status: string(testCase.resultStatus), Reason: "weird"}},
}
}
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: {upstreamID}}, map[string]*model.Task{upstreamID: task}))
var persisted model.Task
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
assert.EqualValues(t, testCase.wantStatus, persisted.Status)
assert.Equal(t, testCase.wantFailures, persisted.PrivateData.PollFailures)
if testCase.wantReason != "" {
assert.Contains(t, persisted.FailReason, testCase.wantReason)
}
if testCase.wantRefund {
assert.Equal(t, initialQuota+preConsumed, getUserQuota(t, userID))
assert.Zero(t, persisted.Quota)
} else {
assert.Equal(t, initialQuota, getUserQuota(t, userID))
}
})
}
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment