cneicy/new-api-custom
1
0 Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

feat: sync deployed v1.0.0-rc.38-custom.42

Browse Source 
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
cneicy committed 2026-09-26 04:49:11 +08:00
1 parent aa6c3f50f4
commit ceddfd3d6b
828 files changed
+122253 -13702
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

controller/relay_task_plugin_test.go
+193 -5
View File
@@ -3,8 +3,10 @@ package controller
import (
"context"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
@@ -22,14 +24,18 @@ import (
"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
onSettle func()
refunds int
events *[]string
settleErr error
reserveErr error
onSettle func()
refunds int
}
func (b *taskSubmissionTestBilling) Settle(int) error {
@@ -49,7 +55,7 @@ func (b *taskSubmissionTestBilling) NeedsRefund() bool { return b.refunds
func (b *taskSubmissionTestBilling) GetPreConsumedQuota() int { return 0 }
func (b *taskSubmissionTestBilling) Reserve(int) error {
*b.events = append(*b.events, "reserve")
return nil
return b.reserveErr
}
func TestPresentTaskSubmissionUsesNativePresenterAfterPersistence(t *testing.T) {
@@ -387,3 +393,185 @@ func taskSubmissionRelayInfo(billing relaycommon.BillingSettler) *relaycommon.Re
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.
func TestImmediateTaskSettlementDatabase(t *testing.T) {
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() })
models := []any{&model.User{}, &model.Channel{}, &model.Task{}, &model.Log{}}
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)
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())
}