diff --git a/controller/plugin_native_e2e_test.go b/controller/plugin_native_e2e_test.go index f0eb4eef0..f0d1c2df4 100644 --- a/controller/plugin_native_e2e_test.go +++ b/controller/plugin_native_e2e_test.go @@ -63,6 +63,10 @@ func (b *nativeRouteBilling) Reserve(quota int) error { return nil } +// ReserveBeforeSubmit holds nothing: this fixture charges in Reserve once the +// upstream has accepted. +func (*nativeRouteBilling) ReserveBeforeSubmit(int) error { return nil } + func TestKlingNativeRouteSubmitPollSettleAndQuery(t *testing.T) { gin.SetMode(gin.TestMode) service.InitHttpClient() diff --git a/controller/relay_task_plugin_test.go b/controller/relay_task_plugin_test.go index 3e0139a43..aadf68c2e 100644 --- a/controller/relay_task_plugin_test.go +++ b/controller/relay_task_plugin_test.go @@ -58,6 +58,8 @@ func (b *taskSubmissionTestBilling) Reserve(int) error { 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"}]}; diff --git a/relay/alpha_search_handler_test.go b/relay/alpha_search_handler_test.go index 412b00974..585dd7621 100644 --- a/relay/alpha_search_handler_test.go +++ b/relay/alpha_search_handler_test.go @@ -71,6 +71,8 @@ func (*recordedSettlement) NeedsRefund() bool { return false } func (*recordedSettlement) GetPreConsumedQuota() int { return 0 } func (*recordedSettlement) Reserve(int) error { return nil } +func (*recordedSettlement) ReserveBeforeSubmit(int) error { return nil } + // /v1/tools/search and /v1/tools/fetch follow /v1/alpha/search: the raw body // reaches /v1/tools/* unchanged and the response returns unchanged. With only // the built-in prices, a response DeepSeek bills (results, page content) diff --git a/relay/common/billing.go b/relay/common/billing.go index 3971426b6..c404f3ad5 100644 --- a/relay/common/billing.go +++ b/relay/common/billing.go @@ -21,4 +21,8 @@ type BillingSettler interface { // Reserve 将预扣额度补到目标值；若目标值不高于当前预扣额度则不做任何事。 Reserve(targetQuota int) error + + // ReserveBeforeSubmit 与 Reserve 相同，用于发起上游请求前价格已经确定的补扣 + // （例如任务重试换到更贵的渠道）：钱包余额不足时返回额度不足，不记为欠费。 + ReserveBeforeSubmit(targetQuota int) error } diff --git a/relay/convert_request_error_test.go b/relay/convert_request_error_test.go index a0c40e858..8002a02c5 100644 --- a/relay/convert_request_error_test.go +++ b/relay/convert_request_error_test.go @@ -42,6 +42,8 @@ func (*imageReservation) Settle(int) error { return nil } func (*imageReservation) Refund(*gin.Context) {} func (*imageReservation) NeedsRefund() bool { return false } +func (s *imageReservation) ReserveBeforeSubmit(quota int) error { return s.Reserve(quota) } + func TestImageRequestReservesFinalQuantityBeforeUpstream(t *testing.T) { service.InitHttpClient() for _, tc := range []struct { diff --git a/relay/helper/price_test.go b/relay/helper/price_test.go index 6a1fefa0f..1731ca188 100644 --- a/relay/helper/price_test.go +++ b/relay/helper/price_test.go @@ -800,3 +800,5 @@ func (s *priceTestReservation) Refund(*gin.Context) {} func (s *priceTestReservation) NeedsRefund() bool { return false } func (s *priceTestReservation) GetPreConsumedQuota() int { return s.held } func (s *priceTestReservation) Reserve(quota int) error { s.held = max(s.held, quota); return nil } + +func (s *priceTestReservation) ReserveBeforeSubmit(quota int) error { return s.Reserve(quota) } diff --git a/relay/relay_task.go b/relay/relay_task.go index fc293cc7a..0ca23b633 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -194,7 +194,7 @@ func ApplyOriginTaskAffinity(c *gin.Context, info *relaycommon.RelayInfo) *dto.T // RelayTaskSubmit 完成 task 提交的全部流程（每次尝试调用一次）： // 刷新渠道元数据 → 确定 platform/adaptor → 验证请求 → -// 估算计费(EstimateBilling) → 计算价格 → 预扣费（仅首次）→ +// 估算计费(EstimateBilling) → 计算价格 → 预扣费（重试时补足差额）→ // 构建/发送/解析上游请求 → 提交后计费调整(AdjustBillingOnSubmit)。 // 共享控制器编排负责未落库退款、最终额度预留、落库和结算。 func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitResult, *dto.TaskError) { @@ -247,8 +247,10 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } } - // 4. 价格计算：基础模型价格 + // 4. 价格计算：基础模型价格。重试可能换了插件或渠道，每次尝试都重新定价， + // 不沿用上一次尝试的表达式快照。 info.OriginModelName = modelName + info.TieredBillingSnapshot = nil var priceData types.PriceData var err error pluginKey := c.GetString("task_plugin_key") @@ -331,8 +333,9 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe noteTaskQuotaClamp(info, clamp) } - // 7. 预扣费（仅首次 — 重试时 info.Billing 已存在，跳过） - if info.Billing == nil && !info.PriceData.FreeModel { + // 7. 预扣费：首次尝试创建计费会话；重试时会话已存在，换到更贵的渠道要在 + // 发起上游请求前补足预扣，余额不足就不发请求。 + if info.Billing != nil || !info.PriceData.FreeModel { info.ForcePreConsume = true if apiErr := service.PreConsumeBilling(c, info.PriceData.Quota, info); apiErr != nil { return nil, service.TaskErrorFromAPIError(apiErr) diff --git a/relay/relay_task_test.go b/relay/relay_task_test.go index 0427a5b49..cab325c41 100644 --- a/relay/relay_task_test.go +++ b/relay/relay_task_test.go @@ -19,6 +19,7 @@ import ( "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/billing_setting" "github.com/QuantumNous/new-api/setting/config" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" "github.com/stretchr/testify/assert" @@ -531,7 +532,7 @@ export function parseTaskResult() { return {status:"SUCCESS"}; } c.Set("group", "default") info.UserGroup, info.UsingGroup = "default", "default" info.OriginModelName = "declared-model" - // An existing reservation skips pre-consume so the fixture reaches the upstream call. + // An existing reservation stands in for the wallet so the fixture reaches the upstream call. info.Billing = &imageReservation{limit: 1 << 30} pinMappingOrderPlugin(t, c, source) @@ -551,3 +552,102 @@ export function parseTaskResult() { return {status:"SUCCESS"}; } }) } } + +// A retry prices the plugin it lands on from scratch and holds that price +// before the upstream sees the request. Plugin A fails upstream on the first +// attempt; the retry runs on plugin B, pinned the way +// SetupContextForSelectedChannel pins another candidate plugin's channel. +func TestRelayTaskSubmitRetryRepricesAndReservesBeforeUpstream(t *testing.T) { + service.InitHttpClient() + const source = ` +export const meta = {apiVersion:1,key:"KEY",name:"KEY",version:"1.0.0",author:{name:"Test"},models:["declared-model"],fetchMode:"per_task",usageSchema:{seconds:{type:"number",unit:"second",description:"Video generation unit price"}}}; +export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl+"/submit", method:"POST", body:{model: ctx.model}, action:"text_to_video"}; } +export function parseSubmitResponse(ctx, response) { return {taskId: response.body.id}; } +export function buildQueryRequest(ctx) { return {url: ctx.baseUrl+"/query"}; } +export function parseTaskResult() { return {status:"SUCCESS"}; } +export function extractUsage() { return {seconds:2}; } +` + pricier := map[string]string{"retry-a::declared-model": `tier("a", u("seconds") * 0.05)`, "retry-b::declared-model": `tier("b", u("seconds") * 0.4)`} + for _, tc := range []struct { + name, perCallPrice, wantCode string + expressions map[string]string + wallet, wantQuota int + wantSubmitted []string + }{ + {name: "per-call plugin drops the previous expression", expressions: map[string]string{"retry-a::declared-model": `tier("a", u("seconds") * 0.02)`}, perCallPrice: `{"declared-model":0.04}`, wallet: 1 << 30, wantQuota: 40000, wantSubmitted: []string{"a", "b"}}, + {name: "pricier plugin is held before submission", expressions: pricier, perCallPrice: `{}`, wallet: 1 << 30, wantQuota: 400000, wantSubmitted: []string{"a", "b"}}, + {name: "unaffordable plugin is never submitted", expressions: pricier, perCallPrice: `{}`, wallet: 100000, wantCode: "insufficient_user_quota", wantSubmitted: []string{"a"}}, + } { + t.Run(tc.name, func(t *testing.T) { + saveBillingConfig(t) + savedPrices := ratio_setting.ModelPrice2JSONString() + t.Cleanup(func() { require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(savedPrices)) }) + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(tc.perCallPrice)) + expressions, err := common.Marshal(tc.expressions) + require.NoError(t, err) + require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{ + billing_setting.PluginBillingExprOption: string(expressions), + "billing_setting.billing_mode": `{}`, + "billing_setting.billing_expr": `{}`, + "group_ratio_setting.group_ratio": `{"default":1}`, + })) + registry := pluginruntime.NewRegistry() + for _, key := range []string{"retry-a", "retry-b"} { + _, err := registry.Register(strings.ReplaceAll(source, "KEY", key), pluginruntime.Options{}) + require.NoError(t, err) + } + // The fake stands in for the session the first attempt created; its + // limit is the most the wallet can hold. + billing := &imageReservation{limit: tc.wallet} + var submitted []string + var heldAtSubmit []int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempt, _, _ := strings.Cut(strings.TrimPrefix(r.URL.Path, "/"), "/") + submitted = append(submitted, attempt) + heldAtSubmit = append(heldAtSubmit, billing.held) + w.Header().Set("Content-Type", "application/json") + if attempt == "a" { + w.WriteHeader(http.StatusInternalServerError) + } + _, _ = w.Write([]byte(`{"id":"job-` + attempt + `"}`)) + })) + defer server.Close() + + c, info := newTaskSubmitContext(t, "declared-model", "") + c.Set("group", "default") + info.UserGroup, info.UsingGroup, info.OriginModelName = "default", "default", "declared-model" + info.Billing = billing + var result *TaskSubmitResult + var taskErr *dto.TaskError + for _, attempt := range []string{"a", "b"} { + plugin, ok := registry.Generation().Get("retry-" + attempt) + require.True(t, ok) + c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Generation: registry.Generation(), Plugin: plugin}) + c.Set("task_plugin_key", plugin.Meta.Key) + common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, server.URL+"/"+attempt) + result, taskErr = RelayTaskSubmit(c, info) + if attempt == "a" { + require.NotNil(t, taskErr) + require.Equal(t, http.StatusInternalServerError, taskErr.StatusCode, "the first attempt fails upstream and is retried") + } + } + + require.Equal(t, tc.wantSubmitted, submitted) + if tc.wantCode != "" { + require.NotNil(t, taskErr) + assert.Equal(t, tc.wantCode, taskErr.Code) + assert.Equal(t, http.StatusForbidden, taskErr.StatusCode) + return + } + require.Nil(t, taskErr, "submission error: %+v", taskErr) + assert.Equal(t, tc.wantQuota, result.Quota) + assert.Equal(t, tc.wantQuota, heldAtSubmit[1], "plugin B must not see the request before its price is held") + if tc.perCallPrice != `{}` { + assert.Nil(t, info.TieredBillingSnapshot, "a per-call plugin must not settle with the previous plugin's expression") + return + } + require.NotNil(t, info.TieredBillingSnapshot) + assert.Equal(t, tc.expressions["retry-b::declared-model"], info.TieredBillingSnapshot.ExprString) + }) + } +} diff --git a/service/billing.go b/service/billing.go index e5e0360eb..091fbf309 100644 --- a/service/billing.go +++ b/service/billing.go @@ -1,6 +1,7 @@ package service import ( + "errors" "fmt" "net/http" @@ -18,6 +19,8 @@ const ( // PreConsumeBilling 根据用户计费偏好创建 BillingSession 并执行预扣费。 // 会话存储在 relayInfo.Billing 上，供后续 Settle / Refund 使用。 +// 会话已存在时（任务重试换到另一个渠道）把预扣补到 preConsumedQuota， +// 和首次预扣一样必须由余额覆盖。 func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.NewAPIError { if relayInfo != nil && relayInfo.QuotaClamp != nil { return types.NewErrorWithStatusCode( @@ -35,6 +38,16 @@ func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycom types.ErrOptionWithSkipRetry(), ) } + if relayInfo != nil && relayInfo.Billing != nil { + if err := relayInfo.Billing.ReserveBeforeSubmit(preConsumedQuota); err != nil { + var apiErr *types.NewAPIError + if errors.As(err, &apiErr) { + return apiErr + } + return types.NewErrorWithStatusCode(err, types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + return nil + } session, apiErr := NewBillingSession(c, relayInfo, preConsumedQuota) if apiErr != nil { return apiErr diff --git a/service/billing_session.go b/service/billing_session.go index e2a24b974..61c6f9c7f 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -156,15 +156,23 @@ func (s *BillingSession) GetPreConsumedQuota() int { } func (s *BillingSession) Reserve(targetQuota int) error { + return s.reserve(targetQuota, false) +} + +// ReserveBeforeSubmit 在发起上游请求前补扣，钱包余额必须覆盖差额。 +func (s *BillingSession) ReserveBeforeSubmit(targetQuota int) error { + return s.reserve(targetQuota, true) +} + +func (s *BillingSession) reserve(targetQuota int, requireAvailableQuota bool) error { s.mu.Lock() defer s.mu.Unlock() - imageRequest := false if s.relayInfo != nil { - _, imageRequest = s.relayInfo.Request.(*dto.ImageRequest) - imageRequest = imageRequest || s.relayInfo.ImageRequestCount > 0 + _, imageRequest := s.relayInfo.Request.(*dto.ImageRequest) + requireAvailableQuota = requireAvailableQuota || imageRequest || s.relayInfo.ImageRequestCount > 0 } - if s.settled || s.refunded || s.trusted && !imageRequest || targetQuota <= s.preConsumedQuota { + if s.settled || s.refunded || s.trusted && !requireAvailableQuota || targetQuota <= s.preConsumedQuota { return nil } @@ -173,7 +181,7 @@ func (s *BillingSession) Reserve(targetQuota int) error { return nil } - if err := s.reserveFunding(delta, imageRequest); err != nil { + if err := s.reserveFunding(delta, requireAvailableQuota); err != nil { return err } if err := s.reserveToken(delta); err != nil { @@ -184,7 +192,7 @@ func (s *BillingSession) Reserve(targetQuota int) error { s.preConsumedQuota += delta s.tokenConsumed += delta s.extraReserved += delta - if imageRequest { + if requireAvailableQuota { s.trusted = false } s.syncRelayInfo() @@ -257,8 +265,8 @@ func (s *BillingSession) reserveFunding(delta int, requireAvailableQuota bool) e switch funding := s.funding.(type) { case *WalletFunding: if requireAvailableQuota { - // Image quantity is known before submission, including retries and - // overrides. Reserve atomically instead of admitting wallet debt. + // Image quantities and task retry prices are known before + // submission. Reserve atomically instead of admitting wallet debt. if err := funding.PreConsume(delta); err != nil { if errors.Is(err, ErrInsufficientWalletQuota) { return types.NewErrorWithStatusCode(err, types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) diff --git a/service/funding_source.go b/service/funding_source.go index e4a356b5e..c22210b27 100644 --- a/service/funding_source.go +++ b/service/funding_source.go @@ -50,7 +50,8 @@ func (w *WalletFunding) PreConsume(amount int) error { if !reserved { return ErrInsufficientWalletQuota } - w.consumed = amount + // 发送前补扣会再次预扣，累加后退款才能退回全部预扣。 + w.consumed += amount return nil } diff --git a/service/text_quota_test.go b/service/text_quota_test.go index 21840b420..d8cfa569d 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -110,7 +110,7 @@ func runFixedPriceAccountingCases(t *testing.T, db, logDB *gorm.DB) { usage *dto.Usage audio, stream, refund, insufficient, tool bool realtime, reserveInsufficient bool - wallet, outboundImages int + wallet, outboundImages, retryReservation int groupRatio float64 want int unit billingexpr.BillingUnit @@ -140,6 +140,8 @@ func runFixedPriceAccountingCases(t *testing.T, db, logDB *gorm.DB) { {name: "image quantity exceeds one-image wallet before submission", expression: `tier("image", fixed(0.04)) * image_count`, requestedImages: 4, wallet: 20000, insufficient: true}, {name: "image override reserves extra quantity", expression: `tier("image", fixed(0.04)) * image_count`, requestedImages: 1, outboundImages: 4, want: 80000, unit: billingexpr.BillingUnitRequest}, {name: "image override cannot exceed remaining wallet", expression: `tier("image", fixed(0.04)) * image_count`, requestedImages: 1, outboundImages: 4, wallet: 40000, reserveInsufficient: true, refund: true}, + {name: "task retry holds a pricier channel and refunds all of it", expression: flat, retryReservation: 40000, refund: true}, + {name: "task retry cannot exceed remaining wallet", expression: flat, retryReservation: 40000, wallet: 20000, reserveInsufficient: true, refund: true}, } { t.Run(tc.name, func(t *testing.T) { quota := startingQuota @@ -203,6 +205,17 @@ func runFixedPriceAccountingCases(t *testing.T, db, logDB *gorm.DB) { assert.Equal(t, tc.want, info.Billing.GetPreConsumedQuota()) } } + if tc.retryReservation > 0 { + retryErr := PreConsumeBilling(ctx, tc.retryReservation, info) + if tc.reserveInsufficient { + require.NotNil(t, retryErr) + assert.Equal(t, types.ErrorCodeInsufficientUserQuota, retryErr.GetErrorCode()) + assert.Equal(t, reservation, info.Billing.GetPreConsumedQuota()) + } else { + require.Nil(t, retryErr) + assert.Equal(t, tc.retryReservation, info.Billing.GetPreConsumedQuota()) + } + } if tc.refund { refunded := make(chan struct{}, 1) const callback = "fixed_billing_refund_observed" diff --git a/service/tiered_settle_test.go b/service/tiered_settle_test.go index 6302e9118..7d97e739a 100644 --- a/service/tiered_settle_test.go +++ b/service/tiered_settle_test.go @@ -340,6 +340,10 @@ func (s *recordingBillingSettler) Reserve(targetQuota int) error { return nil } +func (s *recordingBillingSettler) ReserveBeforeSubmit(targetQuota int) error { + return s.Reserve(targetQuota) +} + func TestPrepareTieredBillingForSelectedGroupUpdatesReservation(t *testing.T) { const expr = `tier("base", p)` billing := &recordingBillingSettler{preConsumedQuota: 50_000}