mirror of
https://gh-proxy.org/https://github.com/klinsmaya/new-api.git
synced 2026-11-08 01:17:09 +08:00
CaIon
55e2ae54d1
A task retry that moved to another plugin or channel kept the previous attempt's usage expression. A retry from an expression-priced plugin to one priced per request charged, and later settled, at the first plugin's price. Each attempt now prices itself from scratch. A retry to a more expensive channel was topped up only after the upstream had accepted the task. A user who could not cover the difference went into arrears, or the task was dropped while the upstream was already running it. The retry now raises the reservation before the request is sent, with the same balance check as the first pre-consume, and a channel the balance cannot cover is never called. A wallet top-up taken before submission replaced the amount held by the first pre-consume instead of adding to it, so a request that then failed refunded only the top-up. Image requests whose quantity was raised before submission were affected too. The full amount is now refunded.
663 lines
28 KiB
Go
663 lines
28 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/model"
|
|
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
|
"github.com/QuantumNous/new-api/relay"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/QuantumNous/new-api/types"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
)
|
|
|
|
type taskSubmissionTestBilling struct {
|
|
events *[]string
|
|
settleErr error
|
|
reserveErr error
|
|
onSettle func()
|
|
refunds int
|
|
}
|
|
|
|
func (b *taskSubmissionTestBilling) Settle(int) error {
|
|
*b.events = append(*b.events, "settle")
|
|
if b.onSettle != nil {
|
|
b.onSettle()
|
|
}
|
|
return b.settleErr
|
|
}
|
|
|
|
func (b *taskSubmissionTestBilling) Refund(*gin.Context) {
|
|
*b.events = append(*b.events, "refund")
|
|
b.refunds++
|
|
}
|
|
|
|
func (b *taskSubmissionTestBilling) NeedsRefund() bool { return b.refunds == 0 }
|
|
func (b *taskSubmissionTestBilling) GetPreConsumedQuota() int { return 0 }
|
|
func (b *taskSubmissionTestBilling) Reserve(int) error {
|
|
*b.events = append(*b.events, "reserve")
|
|
return b.reserveErr
|
|
}
|
|
|
|
func (b *taskSubmissionTestBilling) ReserveBeforeSubmit(quota int) error { return b.Reserve(quota) }
|
|
|
|
func TestPresentTaskSubmissionUsesNativePresenterAfterPersistence(t *testing.T) {
|
|
plugin, err := pluginruntime.CompilePlugin(`
|
|
export const meta = {apiVersion:1,key:"presenter-test",name:"Presenter",version:"1.0.0",author:{name:"Test"},models:["model"],fetchMode:"per_task",routes:[{method:"POST",path:"/vendor/jobs",type:"submit",decode:"decode",render:"created"}]};
|
|
export const native = {decode:function(ctx){return {kind:"submit",model:"model",requestBody:ctx.body.value};},created:function(ctx,task){return {data:{task_id:task.task_id},upstream:task.data};}};
|
|
export function buildSubmitRequest(){return {}} export function parseSubmitResponse(){return {taskId:"upstream"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}}
|
|
`, pluginruntime.Options{})
|
|
require.NoError(t, err)
|
|
priceData := types.PriceData{}
|
|
priceData.AddOtherRatio("seconds", 5)
|
|
task := &model.Task{TaskID: "task_public", SubmitTime: 123}
|
|
task.SetData(map[string]any{"task_id": "upstream_private"})
|
|
outcome := &taskSubmissionOutcome{
|
|
Result: &relay.TaskSubmitResult{},
|
|
Task: task,
|
|
RelayInfo: &relaycommon.RelayInfo{PriceData: priceData},
|
|
}
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"model"}`))
|
|
c.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{Plugin: plugin, Route: plugin.Meta.Routes[0]})
|
|
c.Set(pluginruntime.ContextKeyRouteRequest, pluginruntime.RouteRequestContext{Path: "/vendor/jobs", Method: http.MethodPost, Body: map[string]any{"kind": "json", "value": map[string]any{"model": "model"}}})
|
|
|
|
presentTaskSubmission(c, outcome)
|
|
|
|
assert.JSONEq(t, `{
|
|
"data":{"task_id":"task_public"},
|
|
"upstream":{"task_id":"upstream_private"}
|
|
}`, recorder.Body.String())
|
|
assert.JSONEq(t, `{"seconds":5}`, recorder.Header().Get("X-New-Api-Other-Ratios"))
|
|
}
|
|
|
|
func TestPresentTaskSubmissionFallbackUsesPersistedPublicID(t *testing.T) {
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
outcome := &taskSubmissionOutcome{
|
|
Result: &relay.TaskSubmitResult{},
|
|
Task: &model.Task{TaskID: "task_persisted", SubmitTime: 456},
|
|
RelayInfo: &relaycommon.RelayInfo{OriginModelName: "video-model"},
|
|
}
|
|
|
|
presentTaskSubmission(c, outcome)
|
|
|
|
assert.JSONEq(t, `{
|
|
"id":"task_persisted",
|
|
"task_id":"task_persisted",
|
|
"status":"queued",
|
|
"model":"video-model",
|
|
"created_at":456
|
|
}`, recorder.Body.String())
|
|
}
|
|
|
|
func TestPresentTaskSubmissionUsesHostOpenAIVideoCreateReceipt(t *testing.T) {
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{
|
|
Protocol: "openai_video",
|
|
Operation: pluginruntime.HostProtocolOperation{Name: "create"},
|
|
})
|
|
task := &model.Task{
|
|
TaskID: "task_public",
|
|
Status: model.TaskStatusSubmitted,
|
|
Progress: "0%",
|
|
CreatedAt: 456,
|
|
Properties: model.Properties{OriginModelName: "video-model"},
|
|
}
|
|
outcome := &taskSubmissionOutcome{Result: &relay.TaskSubmitResult{}, Task: task, RelayInfo: &relaycommon.RelayInfo{}}
|
|
|
|
presentTaskSubmission(c, outcome)
|
|
|
|
assert.JSONEq(t, `{"id":"task_public","object":"video","model":"video-model","status":"queued","progress":0,"created_at":456}`, recorder.Body.String())
|
|
assert.NotContains(t, recorder.Body.String(), "task_id")
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionRefundsWhenInsertFails(t *testing.T) {
|
|
events := make([]string, 0, 3)
|
|
database := setupTaskSubmissionDatabase(t, false, &events)
|
|
_ = database
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
c := taskSubmissionTestContext()
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream_private",
|
|
Platform: constant.TaskPlatform("plugin"),
|
|
}, nil
|
|
})
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "task_insert_failed", taskErr.Code)
|
|
assert.Equal(t, []string{"reserve", "insert", "refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionSettlementFailureStaysDurableAndWritesNothing(t *testing.T) {
|
|
events := make([]string, 0, 3)
|
|
database := setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events, settleErr: errors.New("settlement failed")}
|
|
c := taskSubmissionTestContext()
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream_private",
|
|
Platform: constant.TaskPlatform("plugin"),
|
|
}, nil
|
|
})
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "task_billing_settlement_failed", taskErr.Code)
|
|
assert.Equal(t, []string{"reserve", "insert", "settle"}, events)
|
|
assert.Zero(t, billing.refunds)
|
|
var count int64
|
|
require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error)
|
|
assert.Equal(t, int64(1), count)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionPersistsPinnedPluginProvenance(t *testing.T) {
|
|
events := make([]string, 0, 3)
|
|
database := setupTaskSubmissionDatabase(t, true, &events)
|
|
previousLogConsumeEnabled := common.LogConsumeEnabled
|
|
common.LogConsumeEnabled = false
|
|
t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled })
|
|
|
|
c := taskSubmissionTestContext()
|
|
c.Set(common.RequestIdKey, "request-public")
|
|
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{
|
|
Generation: &pluginruntime.RoutingGeneration{Number: 42},
|
|
Plugin: &pluginruntime.LoadedPlugin{Meta: pluginruntime.Meta{
|
|
Key: "document-parser",
|
|
Name: "Document Parser",
|
|
Version: "1.2.3",
|
|
APIVersion: 1,
|
|
Author: pluginruntime.AuthorMeta{
|
|
Name: "Community Author",
|
|
URL: "https://plugins.example/author",
|
|
},
|
|
}},
|
|
})
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream-private",
|
|
Platform: constant.TaskPlatform("document-parser"),
|
|
}, nil
|
|
})
|
|
|
|
require.Nil(t, taskErr)
|
|
require.NotNil(t, outcome)
|
|
require.NotNil(t, outcome.Task.PrivateData.Execution)
|
|
require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin)
|
|
assert.Equal(t, "request-public", outcome.Task.PrivateData.Execution.RequestID)
|
|
assert.Equal(t, "/plugin/submit", outcome.Task.PrivateData.Execution.RequestPath)
|
|
assert.Equal(t, "1.2.3", outcome.Task.PrivateData.Execution.TaskPlugin.Version)
|
|
assert.Equal(t, uint64(42), outcome.Task.PrivateData.Execution.TaskPlugin.Generation)
|
|
require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin.Author)
|
|
assert.Equal(t, "Community Author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.Name)
|
|
assert.Equal(t, "https://plugins.example/author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.URL)
|
|
|
|
var stored model.Task
|
|
require.NoError(t, database.Where("task_id = ?", "task_public").First(&stored).Error)
|
|
require.NotNil(t, stored.PrivateData.Execution)
|
|
require.NotNil(t, stored.PrivateData.Execution.TaskPlugin)
|
|
assert.Equal(t, "document-parser", stored.PrivateData.Execution.TaskPlugin.Key)
|
|
require.NotNil(t, stored.PrivateData.Execution.TaskPlugin.Author)
|
|
assert.Equal(t, "Community Author", stored.PrivateData.Execution.TaskPlugin.Author.Name)
|
|
assert.Equal(t, "upstream-private", stored.PrivateData.UpstreamTaskID)
|
|
}
|
|
|
|
// A submit route declaring retainResult: false persists the task row for
|
|
// billing but never writes the upstream snapshot for an immediate terminal
|
|
// result, while the in-memory task still carries it for the presenter. An
|
|
// asynchronous result on the same route is retained because polling and
|
|
// retrieval depend on it.
|
|
func TestExecuteTaskSubmissionHonorsRouteRetainResult(t *testing.T) {
|
|
retainFalse := false
|
|
for _, tc := range []struct {
|
|
name string
|
|
immediate *relaycommon.TaskInfo
|
|
wantDiscarded bool
|
|
}{
|
|
{"immediate success is discarded", &relaycommon.TaskInfo{Status: model.TaskStatusSuccess, Progress: "100%"}, true},
|
|
{"immediate failure is discarded", &relaycommon.TaskInfo{Status: model.TaskStatusFailure, Reason: "rejected"}, true},
|
|
{"asynchronous result is retained", nil, false},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
events := make([]string, 0, 3)
|
|
database, _ := openTaskDialectDatabase(t, &model.Task{}, &model.User{}, &model.Channel{})
|
|
previousDB := model.DB
|
|
model.DB = database
|
|
t.Cleanup(func() { model.DB = previousDB })
|
|
previousLogConsumeEnabled := common.LogConsumeEnabled
|
|
common.LogConsumeEnabled = false
|
|
t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled })
|
|
|
|
c := taskSubmissionTestContext()
|
|
c.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{
|
|
Plugin: &pluginruntime.LoadedPlugin{Meta: pluginruntime.Meta{Key: "sync-images"}},
|
|
Route: pluginruntime.Route{Method: http.MethodPost, Path: "/sync/images", Type: pluginruntime.RouteTypeSubmit, RetainResult: &retainFalse},
|
|
})
|
|
info := taskSubmissionRelayInfo(&taskSubmissionTestBilling{events: &events})
|
|
upstream := []byte(`{"data":[{"url":"https://cdn.example/a.png"}]}`)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "task_public",
|
|
Platform: constant.TaskPlatform("sync-images"),
|
|
TaskData: upstream,
|
|
Immediate: tc.immediate,
|
|
}, nil
|
|
})
|
|
require.Nil(t, taskErr)
|
|
require.NotNil(t, outcome)
|
|
assert.JSONEq(t, string(upstream), string(outcome.Task.Data), "presenter keeps the in-memory snapshot")
|
|
assert.Equal(t, tc.wantDiscarded, outcome.Task.PrivateData.ResultDiscarded)
|
|
assert.Equal(t, !tc.wantDiscarded, outcome.Task.ResultRetrievable())
|
|
|
|
var stored model.Task
|
|
require.NoError(t, database.Where("task_id = ?", "task_public").First(&stored).Error)
|
|
assert.Equal(t, tc.wantDiscarded, stored.PrivateData.ResultDiscarded)
|
|
var nullCount int64
|
|
require.NoError(t, database.Model(&model.Task{}).Where("task_id = ? AND data IS NULL", "task_public").Count(&nullCount).Error)
|
|
if tc.wantDiscarded {
|
|
assert.Equal(t, int64(1), nullCount, "snapshot column must stay NULL")
|
|
artifacts, err := projectTaskArtifacts(&stored)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, artifacts)
|
|
} else {
|
|
assert.Equal(t, int64(0), nullCount)
|
|
assert.JSONEq(t, string(upstream), string(stored.Data))
|
|
}
|
|
listed := model.TaskGetAllUserTask(1, 0, 10, model.SyncTaskQueryParams{})
|
|
require.Len(t, listed, 1)
|
|
assert.Empty(t, listed[0].Data, "task lists never select the snapshot column")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionRefundsCancellationBeforeDurableBarrier(t *testing.T) {
|
|
events := make([]string, 0, 2)
|
|
setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
c := taskSubmissionTestContext()
|
|
requestContext, cancel := context.WithCancel(c.Request.Context())
|
|
c.Request = c.Request.WithContext(requestContext)
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
cancel()
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream_private",
|
|
Platform: constant.TaskPlatform("plugin"),
|
|
}, nil
|
|
})
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "request_cancelled", taskErr.Code)
|
|
assert.Equal(t, []string{"refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionDisconnectBeforeUpstreamAcceptanceSkipsSubmitAndRefunds(t *testing.T) {
|
|
events := make([]string, 0, 1)
|
|
setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
c := taskSubmissionTestContext()
|
|
requestContext, cancel := context.WithCancel(c.Request.Context())
|
|
cancel()
|
|
c.Request = c.Request.WithContext(requestContext)
|
|
info := taskSubmissionRelayInfo(billing)
|
|
submitted := false
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
submitted = true
|
|
return nil, nil
|
|
})
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "request_cancelled", taskErr.Code)
|
|
assert.False(t, submitted)
|
|
assert.Equal(t, []string{"refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionCallerCancellationDuringSubmitRefundsBeforeDurableBarrier(t *testing.T) {
|
|
events := make([]string, 0, 1)
|
|
setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
c := taskSubmissionTestContext()
|
|
requestContext, cancel := context.WithCancel(c.Request.Context())
|
|
c.Request = c.Request.WithContext(requestContext)
|
|
info := taskSubmissionRelayInfo(billing)
|
|
submitStarted := make(chan struct{})
|
|
done := make(chan struct{})
|
|
var outcome *taskSubmissionOutcome
|
|
var taskErr *dto.TaskError
|
|
|
|
go func() {
|
|
defer close(done)
|
|
outcome, taskErr = executeTaskSubmissionWith(c, info, func(c *gin.Context, _ *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
close(submitStarted)
|
|
<-c.Request.Context().Done()
|
|
return nil, service.TaskErrorWrapperLocal(c.Request.Context().Err(), "do_request_failed", http.StatusInternalServerError)
|
|
})
|
|
}()
|
|
select {
|
|
case <-submitStarted:
|
|
case <-time.After(2 * time.Second):
|
|
require.FailNow(t, "submission did not start")
|
|
}
|
|
cancel()
|
|
select {
|
|
case <-done:
|
|
case <-time.After(2 * time.Second):
|
|
require.FailNow(t, "submission did not stop after disconnect")
|
|
}
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "request_cancelled", taskErr.Code)
|
|
assert.Equal(t, []string{"refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionDisconnectAfterDurableInsertDoesNotRefund(t *testing.T) {
|
|
events := make([]string, 0, 3)
|
|
database := setupTaskSubmissionDatabase(t, true, &events)
|
|
previousLogConsumeEnabled := common.LogConsumeEnabled
|
|
common.LogConsumeEnabled = false
|
|
t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled })
|
|
c := taskSubmissionTestContext()
|
|
requestContext, cancel := context.WithCancel(c.Request.Context())
|
|
c.Request = c.Request.WithContext(requestContext)
|
|
billing := &taskSubmissionTestBilling{
|
|
events: &events,
|
|
onSettle: cancel,
|
|
}
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream_private",
|
|
Platform: constant.TaskPlatform("plugin"),
|
|
}, nil
|
|
})
|
|
|
|
require.Nil(t, taskErr)
|
|
require.NotNil(t, outcome)
|
|
assert.Equal(t, "task_public", outcome.Task.TaskID)
|
|
assert.Equal(t, []string{"reserve", "insert", "settle"}, events)
|
|
assert.Zero(t, billing.refunds)
|
|
var count int64
|
|
require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error)
|
|
assert.Equal(t, int64(1), count)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
// setupTaskSubmissionDatabase opens the dialect selected by
|
|
// TEST_TASK_DB_DIALECT (SQLite in memory by default) and records every task
|
|
// INSERT in events so tests can assert the reserve → insert → settle order.
|
|
// Without migrate the task table does not exist and inserts fail.
|
|
func setupTaskSubmissionDatabase(t *testing.T, migrate bool, events *[]string) *gorm.DB {
|
|
t.Helper()
|
|
previousDB := model.DB
|
|
var models []any
|
|
if migrate {
|
|
models = append(models, &model.Task{})
|
|
}
|
|
database, _ := openTaskDialectDatabase(t, models...)
|
|
require.NoError(t, database.Callback().Create().Before("gorm:create").Register("test:task-submit-order", func(*gorm.DB) {
|
|
*events = append(*events, "insert")
|
|
}))
|
|
model.DB = database
|
|
t.Cleanup(func() { model.DB = previousDB })
|
|
return database
|
|
}
|
|
|
|
func taskSubmissionTestContext() *gin.Context {
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/plugin/submit", strings.NewReader(`{}`))
|
|
return c
|
|
}
|
|
|
|
func taskSubmissionRelayInfo(billing relaycommon.BillingSettler) *relaycommon.RelayInfo {
|
|
return &relaycommon.RelayInfo{
|
|
UserId: 1,
|
|
UsingGroup: "default",
|
|
OriginModelName: "plugin-model",
|
|
Billing: billing,
|
|
TaskRelayInfo: &relaycommon.TaskRelayInfo{
|
|
PublicTaskID: "task_public",
|
|
LockedChannel: &model.Channel{Id: 1, Type: constant.ChannelTypeTaskPlugin, Name: "plugin"},
|
|
},
|
|
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1, ChannelType: constant.ChannelTypeTaskPlugin},
|
|
}
|
|
}
|
|
|
|
// Uses the real submit adaptor, expression evaluator, BillingSession, task row
|
|
// and consume log. Set TEST_TASK_DB_DIALECT plus TEST_MYSQL_DSN or
|
|
// TEST_POSTGRES_DSN to exercise the same contract on an external test database.
|
|
// Unique table prefixes keep the fixture isolated from all existing tables.
|
|
// openTaskDialectDatabase opens the engine selected by TEST_TASK_DB_DIALECT
|
|
// (default SQLite in memory; MySQL and PostgreSQL through TEST_MYSQL_DSN and
|
|
// TEST_POSTGRES_DSN) with a unique table prefix, migrates the given models and
|
|
// drops them on cleanup. It logs the engine version so database verification
|
|
// runs leave a record.
|
|
func openTaskDialectDatabase(t *testing.T, models ...any) (*gorm.DB, common.DatabaseType) {
|
|
t.Helper()
|
|
dialect := common.DatabaseType(os.Getenv("TEST_TASK_DB_DIALECT"))
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "", common.DatabaseTypeSQLite:
|
|
dialect = common.DatabaseTypeSQLite
|
|
driver = sqlite.Open(":memory:")
|
|
case common.DatabaseTypeMySQL:
|
|
require.NotEmpty(t, os.Getenv("TEST_MYSQL_DSN"))
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
case common.DatabaseTypePostgreSQL:
|
|
require.NotEmpty(t, os.Getenv("TEST_POSTGRES_DSN"))
|
|
driver = postgres.New(postgres.Config{DSN: os.Getenv("TEST_POSTGRES_DSN"), PreferSimpleProtocol: true})
|
|
default:
|
|
t.Fatalf("unsupported test dialect %q", dialect)
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("tsubmit_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
sqlDB, err := db.DB()
|
|
require.NoError(t, err)
|
|
sqlDB.SetMaxOpenConns(1)
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
|
require.NoError(t, db.AutoMigrate(models...))
|
|
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(models...)) })
|
|
var version string
|
|
if dialect == common.DatabaseTypeSQLite {
|
|
require.NoError(t, db.Raw("SELECT sqlite_version()").Scan(&version).Error)
|
|
} else {
|
|
require.NoError(t, db.Raw("SELECT version()").Scan(&version).Error)
|
|
}
|
|
t.Logf("database: %s %s", dialect, version)
|
|
return db, dialect
|
|
}
|
|
|
|
func TestImmediateTaskSettlementDatabase(t *testing.T) {
|
|
db, dialect := openTaskDialectDatabase(t, &model.User{}, &model.Channel{}, &model.Task{}, &model.Log{})
|
|
oldDB, oldLogDB := model.DB, model.LOG_DB
|
|
oldMain, oldLog := common.MainDatabaseType(), common.LogDatabaseType()
|
|
oldRedis, oldMemory, oldBatch, oldConsume, oldExport := common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled
|
|
model.DB, model.LOG_DB = db, db
|
|
common.SetDatabaseTypes(dialect, dialect)
|
|
common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled = false, false, false, true, false
|
|
t.Cleanup(func() {
|
|
model.DB, model.LOG_DB = oldDB, oldLogDB
|
|
common.SetDatabaseTypes(oldMain, oldLog)
|
|
common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled = oldRedis, oldMemory, oldBatch, oldConsume, oldExport
|
|
})
|
|
|
|
const expression = `u("units") == 7 ? tier("missing", u("missing") * 1.0) : u("units") == 8 ? tier("invalid", -1.0) : u("units") > 3 ? tier("bulk", u("units") * 0.01) : tier("small", u("units") * 0.01)`
|
|
withTieredBillingConfig(t, map[string]string{"document-model": "tiered_expr"}, map[string]string{"document-model": expression})
|
|
const source = `
|
|
export const meta={apiVersion:1,key:"generic-settlement",name:"Generic settlement",version:"1.0.0",author:{name:"Test"},models:["document-model"],fetchMode:"per_task",usageSchema:{units:{type:"number",unit:"count"}}};
|
|
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/compile",body:ctx.requestBody};}
|
|
export function parseSubmitResponse(ctx,resp){return {taskId:"vendor-job",taskData:resp.body,immediate:{status:resp.body.status,reason:"provider rejected job"}};}
|
|
export function extractUsage(){return {units:4};}
|
|
export function extractUsageOnComplete(ctx,result,body){return body.usage;}
|
|
export function parseTaskResult(){throw new Error("completed submissions must not poll");}
|
|
export function buildQueryRequest(){throw new Error("completed submissions must not poll");}
|
|
`
|
|
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
|
|
require.NoError(t, err)
|
|
for index, tc := range []struct {
|
|
name, status string
|
|
actual any
|
|
count float64
|
|
}{
|
|
{"partial", "SUCCESS", 2, 2}, {"zero", "SUCCESS", 0, 0}, {"larger", "SUCCESS", 6, 6},
|
|
{"invalid usage", "SUCCESS", -1, 4}, {"expression failure", "SUCCESS", 7, 4}, {"negative result", "SUCCESS", 8, 4}, {"missing usage", "SUCCESS", nil, 4}, {"failed", "FAILURE", 9, 0},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
body := map[string]any{"status": tc.status}
|
|
if tc.actual != nil {
|
|
body["usage"] = map[string]any{"units": tc.actual}
|
|
}
|
|
encoded, err := common.Marshal(body)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
_, _ = w.Write(encoded)
|
|
}))
|
|
defer server.Close()
|
|
initial := int(20 * common.QuotaPerUnit)
|
|
user := model.User{Username: fmt.Sprintf("task_user_%d", index), AffCode: fmt.Sprintf("task_aff_%d", index), Quota: initial}
|
|
require.NoError(t, db.Create(&user).Error)
|
|
ch := model.Channel{Name: "test provider", Type: constant.ChannelTypeTaskPlugin}
|
|
require.NoError(t, db.Create(&ch).Error)
|
|
c := taskSubmissionTestContext()
|
|
c.Set("group", "default")
|
|
c.Set("username", user.Username)
|
|
c.Set("task_request", map[string]any{"model": "document-model"})
|
|
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Plugin: plugin})
|
|
common.SetContextKey(c, constant.ContextKeyOriginalModel, "document-model")
|
|
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, server.URL)
|
|
common.SetContextKey(c, constant.ContextKeyChannelId, ch.Id)
|
|
common.SetContextKey(c, constant.ContextKeyChannelType, ch.Type)
|
|
info := taskSubmissionRelayInfo(nil)
|
|
info.UserId = user.Id
|
|
info.OriginModelName = "document-model"
|
|
info.UserGroup = "default"
|
|
info.UsingGroup = "default"
|
|
info.IsPlayground = true
|
|
info.UserSetting.BillingPreference = "wallet_only"
|
|
info.PublicTaskID = model.GenerateTaskID()
|
|
info.LockedChannel = &ch
|
|
outcome, taskErr := executeTaskSubmission(c, info)
|
|
require.Nil(t, taskErr)
|
|
require.NotNil(t, outcome)
|
|
want := common.QuotaRound(tc.count * 0.01 * common.QuotaPerUnit)
|
|
assert.Equal(t, want, outcome.Result.Quota)
|
|
assert.Equal(t, want, info.PriceData.Quota)
|
|
var stored model.Task
|
|
require.NoError(t, db.Where("task_id = ?", info.PublicTaskID).First(&stored).Error)
|
|
assert.Equal(t, want, stored.Quota)
|
|
assert.Equal(t, model.TaskStatus(tc.status), stored.Status)
|
|
assert.Positive(t, stored.FinishTime)
|
|
assert.Equal(t, float64(4), info.TieredBillingSnapshot.EstimatedQuotaBeforeGroup/(0.01*common.QuotaPerUnit))
|
|
if tc.status == "SUCCESS" {
|
|
assert.Equal(t, tc.count, stored.PrivateData.BillingContext.TieredSnapshot.UsageFacts["units"])
|
|
}
|
|
var updated model.User
|
|
require.NoError(t, db.First(&updated, user.Id).Error)
|
|
assert.Equal(t, initial-want, updated.Quota)
|
|
assert.Equal(t, want, updated.UsedQuota)
|
|
var logs []model.Log
|
|
require.NoError(t, db.Where("user_id = ?", user.Id).Find(&logs).Error)
|
|
require.Len(t, logs, 1)
|
|
assert.Equal(t, want, logs[0].Quota)
|
|
var other map[string]any
|
|
require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other))
|
|
if tc.status == "SUCCESS" {
|
|
assert.Equal(t, tc.count, other["usage_facts"].(map[string]any)["units"])
|
|
}
|
|
assert.False(t, c.Writer.Written(), "presentation must follow persistence and settlement")
|
|
require.NoError(t, info.Billing.Settle(want))
|
|
info.Billing.Refund(c)
|
|
require.NoError(t, db.First(&updated, user.Id).Error)
|
|
assert.Equal(t, initial-want, updated.Quota, "terminal settlement is idempotent")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAcceptedSubmitStreamNeverRetries(t *testing.T) {
|
|
c := taskSubmissionTestContext()
|
|
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "task_accepted", Source: "system"}, decideTaskRetry(c, &dto.TaskError{StatusCode: 502, LocalError: true, NoRetry: true}, 3))
|
|
}
|
|
|
|
// Local task rejections carry a message but no cause; the response and the
|
|
// decision record must still be produced.
|
|
func TestRespondTaskSubmissionErrorWithoutCause(t *testing.T) {
|
|
previousErrorLog := constant.ErrorLogEnabled
|
|
constant.ErrorLogEnabled = false
|
|
t.Cleanup(func() { constant.ErrorLogEnabled = previousErrorLog })
|
|
c := taskSubmissionTestContext()
|
|
taskErr := &dto.TaskError{Code: "get_channel_failed", Message: "no channel", StatusCode: http.StatusServiceUnavailable, LocalError: true}
|
|
require.NotPanics(t, func() { respondTaskSubmissionError(c, taskErr) })
|
|
assert.Equal(t, http.StatusServiceUnavailable, c.Writer.Status())
|
|
events := service.RequestPolicy(c).Events()
|
|
require.Len(t, events, 1)
|
|
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "request_failed", Source: "system"}, events[0].Decision)
|
|
assert.Equal(t, http.StatusServiceUnavailable, events[0].Status)
|
|
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "local_rejection", Source: "system"}, decideTaskRetry(c, &dto.TaskError{StatusCode: http.StatusForbidden, LocalError: true, Message: "billing"}, 2), "local errors stop once the status rules do not force a retry")
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionRefundsWhenFinalReserveFails(t *testing.T) {
|
|
events := []string{}
|
|
setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events, reserveErr: errors.New("insufficient funds")}
|
|
c := taskSubmissionTestContext()
|
|
info := taskSubmissionRelayInfo(billing)
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{Platform: "plugin", Quota: 600, Immediate: &relaycommon.TaskInfo{Status: "SUCCESS"}}, nil
|
|
})
|
|
require.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, http.StatusForbidden, taskErr.StatusCode)
|
|
assert.Equal(t, []string{"reserve", "refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|