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 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()) }