fix(billing): 封住发布前计费竞态
ci / verify (pull_request) Failing after 8s
ci / verify (pull_request) Failing after 8s
阻止上游提交状态不明的任务被租约接管后重复执行,并为人工复核保留可操作的结算记录。 将生产提交绑定到当前估价签名,统一复用预处理快照,并补强规则形状、定点溢出与历史规则兼容校验。 已通过 PostgreSQL 16 集成测试、Go 全量测试与静态检查、前端测试与构建、OpenAPI、依赖审计、镜像、迁移、流水线和 SemVer 门禁。
This commit is contained in:
@@ -931,6 +931,10 @@ func (s *Server) estimatePricing(w http.ResponseWriter, r *http.Request) {
|
||||
writeErrorWithDetails(w, statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), store.ModelCandidateErrorCode(err))
|
||||
return
|
||||
}
|
||||
if code := clients.ErrorCode(err); code == "bad_request" || code == "invalid_parameter" {
|
||||
writeErrorWithDetails(w, http.StatusBadRequest, runErrorMessage(err), runErrorDetails(err), code)
|
||||
return
|
||||
}
|
||||
s.logger.Error("estimate pricing failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "estimate pricing failed")
|
||||
return
|
||||
|
||||
@@ -165,6 +165,13 @@ func validPricingRuleSetInput(input store.PricingRuleSetInput) bool {
|
||||
default:
|
||||
return false
|
||||
}
|
||||
calculator := strings.TrimSpace(rule.CalculatorType)
|
||||
if calculator == "" {
|
||||
calculator = store.DefaultEffectivePricingCalculator(rule.ResourceType)
|
||||
}
|
||||
if store.ValidateEffectivePricingRuleShape(strings.TrimSpace(rule.ResourceType), store.NormalizeEffectivePricingRuleUnit(rule.ResourceType, rule.Unit), calculator) != nil {
|
||||
return false
|
||||
}
|
||||
effectiveFrom, fromOK := pricingEffectiveTime(rule.EffectiveFrom)
|
||||
effectiveTo, toOK := pricingEffectiveTime(rule.EffectiveTo)
|
||||
if !fromOK || !toOK || (!effectiveFrom.IsZero() && !effectiveTo.IsZero() && !effectiveFrom.Before(effectiveTo)) {
|
||||
|
||||
@@ -2,7 +2,10 @@ package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -19,6 +22,9 @@ func (s *Service) renewTaskExecutionLease(ctx context.Context, cancel context.Ca
|
||||
return
|
||||
case <-ticker.C:
|
||||
if err := s.store.RenewTaskExecutionLease(ctx, taskID, executionToken, taskExecutionLeaseTTL); err != nil {
|
||||
if errors.Is(err, store.ErrTaskExecutionFinished) {
|
||||
return
|
||||
}
|
||||
s.logger.Warn("task execution lease lost", "taskID", taskID, "error_category", "task_execution_lease_lost")
|
||||
cancel()
|
||||
return
|
||||
|
||||
@@ -38,8 +38,11 @@ func (s *Service) Estimate(ctx context.Context, kind string, model string, body
|
||||
}
|
||||
estimates := make([]candidateEstimate, 0, len(candidates))
|
||||
for _, candidate := range candidates {
|
||||
candidateBody := preprocessRequest(kind, cloneMap(body), candidate)
|
||||
estimate, candidateErr := s.estimateCandidateV2(ctx, user, kind, candidateBody, candidate)
|
||||
preprocessing := s.preprocessRequestWithScripts(ctx, kind, body, candidate)
|
||||
if preprocessing.Err != nil {
|
||||
return EstimateResult{}, parameterPreprocessClientError(preprocessing.Err)
|
||||
}
|
||||
estimate, candidateErr := s.estimateCandidateV2(ctx, user, kind, preprocessing.Body, candidate)
|
||||
if candidateErr != nil {
|
||||
return EstimateResult{}, candidateErr
|
||||
}
|
||||
|
||||
@@ -20,9 +20,11 @@ import (
|
||||
const (
|
||||
pricingVersionV2 = "effective-pricing-v2"
|
||||
fixedScale = int64(1_000_000_000)
|
||||
maxFixedAmount = fixedAmount(math.MaxInt64)
|
||||
)
|
||||
|
||||
var ErrPricingUnavailable = errors.New("pricing unavailable")
|
||||
var errFixedAmountOverflow = errors.New("fixed amount overflow")
|
||||
|
||||
type PricingUnavailableError struct {
|
||||
Reason string
|
||||
@@ -129,11 +131,11 @@ func fixedAmountFromAny(value any) (fixedAmount, error) {
|
||||
case json.Number:
|
||||
return parseFixedAmount(typed.String())
|
||||
case int:
|
||||
return fixedAmount(int64(typed) * fixedScale), nil
|
||||
return parseFixedAmount(strconv.FormatInt(int64(typed), 10))
|
||||
case int64:
|
||||
return fixedAmount(typed * fixedScale), nil
|
||||
return parseFixedAmount(strconv.FormatInt(typed, 10))
|
||||
case int32:
|
||||
return fixedAmount(int64(typed) * fixedScale), nil
|
||||
return parseFixedAmount(strconv.FormatInt(int64(typed), 10))
|
||||
case float64:
|
||||
if math.IsNaN(typed) || math.IsInf(typed, 0) {
|
||||
return 0, fmt.Errorf("non-finite decimal")
|
||||
@@ -159,26 +161,37 @@ func (a fixedAmount) String() string {
|
||||
func (a fixedAmount) Float64() float64 { return float64(a) / float64(fixedScale) }
|
||||
func (a fixedAmount) IsZero() bool { return a == 0 }
|
||||
|
||||
func (a fixedAmount) Add(other fixedAmount) fixedAmount { return a + other }
|
||||
|
||||
func (a fixedAmount) MulInt(multiplier int) fixedAmount {
|
||||
return fixedAmount(int64(a) * int64(multiplier))
|
||||
func addFixedAmounts(left fixedAmount, right fixedAmount) (fixedAmount, error) {
|
||||
value := new(big.Int).Add(big.NewInt(int64(left)), big.NewInt(int64(right)))
|
||||
return fixedAmountFromBigInt(value)
|
||||
}
|
||||
|
||||
func (a fixedAmount) Mul(other fixedAmount) fixedAmount {
|
||||
product := new(big.Int).Mul(big.NewInt(int64(a)), big.NewInt(int64(other)))
|
||||
return fixedAmount(roundBigIntRatio(product, big.NewInt(fixedScale)))
|
||||
func multiplyFixedAmountByInt(amount fixedAmount, multiplier int) (fixedAmount, error) {
|
||||
value := new(big.Int).Mul(big.NewInt(int64(amount)), big.NewInt(int64(multiplier)))
|
||||
return fixedAmountFromBigInt(value)
|
||||
}
|
||||
|
||||
func (a fixedAmount) MulRatio(numerator int, denominator int) fixedAmount {
|
||||
func multiplyFixedAmounts(left fixedAmount, right fixedAmount) (fixedAmount, error) {
|
||||
product := new(big.Int).Mul(big.NewInt(int64(left)), big.NewInt(int64(right)))
|
||||
return fixedAmountFromBigInt(roundBigIntRatio(product, big.NewInt(fixedScale)))
|
||||
}
|
||||
|
||||
func multiplyFixedAmountRatio(amount fixedAmount, numerator int, denominator int) (fixedAmount, error) {
|
||||
if denominator == 0 {
|
||||
return 0
|
||||
return 0, fmt.Errorf("division by zero")
|
||||
}
|
||||
product := new(big.Int).Mul(big.NewInt(int64(a)), big.NewInt(int64(numerator)))
|
||||
return fixedAmount(roundBigIntRatio(product, big.NewInt(int64(denominator))))
|
||||
product := new(big.Int).Mul(big.NewInt(int64(amount)), big.NewInt(int64(numerator)))
|
||||
return fixedAmountFromBigInt(roundBigIntRatio(product, big.NewInt(int64(denominator))))
|
||||
}
|
||||
|
||||
func roundBigIntRatio(numerator *big.Int, denominator *big.Int) int64 {
|
||||
func fixedAmountFromBigInt(value *big.Int) (fixedAmount, error) {
|
||||
if value == nil || !value.IsInt64() {
|
||||
return 0, errFixedAmountOverflow
|
||||
}
|
||||
return fixedAmount(value.Int64()), nil
|
||||
}
|
||||
|
||||
func roundBigIntRatio(numerator *big.Int, denominator *big.Int) *big.Int {
|
||||
negative := numerator.Sign() < 0
|
||||
absolute := new(big.Int).Abs(new(big.Int).Set(numerator))
|
||||
quotient, remainder := new(big.Int), new(big.Int)
|
||||
@@ -189,7 +202,7 @@ func roundBigIntRatio(numerator *big.Int, denominator *big.Int) int64 {
|
||||
if negative {
|
||||
quotient.Neg(quotient)
|
||||
}
|
||||
return quotient.Int64()
|
||||
return quotient
|
||||
}
|
||||
|
||||
type resolvedPricing struct {
|
||||
@@ -248,6 +261,105 @@ func (pricing resolvedPricing) price(resource string, keys ...string) (fixedAmou
|
||||
return 0, false, nil
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) weight(resource string, key string, name string) (fixedAmount, error) {
|
||||
if strings.TrimSpace(name) == "" {
|
||||
return fixedAmount(fixedScale), nil
|
||||
}
|
||||
keys := weightKeyAliases(key)
|
||||
names := weightValueAliases(key, name)
|
||||
value, found := pricingWeightValue(pricing.Config, resource, keys, names)
|
||||
if !found {
|
||||
return fixedAmount(fixedScale), nil
|
||||
}
|
||||
weight, err := fixedAmountFromAny(value)
|
||||
if err != nil || weight <= 0 {
|
||||
reason := fmt.Sprintf("invalid %s weight for %s", key, name)
|
||||
if err != nil {
|
||||
reason += ": " + err.Error()
|
||||
}
|
||||
return 0, &PricingUnavailableError{Reason: reason, ResourceType: resource, RuleSetID: pricing.RuleSetID}
|
||||
}
|
||||
return weight, nil
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) calculate(resource string, base fixedAmount, integerFactors []int, fixedFactors ...fixedAmount) (fixedAmount, error) {
|
||||
amount := base
|
||||
var err error
|
||||
for _, factor := range integerFactors {
|
||||
amount, err = multiplyFixedAmountByInt(amount, factor)
|
||||
if err != nil {
|
||||
return 0, pricing.calculationError(resource, err)
|
||||
}
|
||||
}
|
||||
for _, factor := range fixedFactors {
|
||||
amount, err = multiplyFixedAmounts(amount, factor)
|
||||
if err != nil {
|
||||
return 0, pricing.calculationError(resource, err)
|
||||
}
|
||||
}
|
||||
return amount, nil
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) calculateRatio(resource string, base fixedAmount, numerator int, denominator int, fixedFactors ...fixedAmount) (fixedAmount, error) {
|
||||
amount, err := multiplyFixedAmountRatio(base, numerator, denominator)
|
||||
if err != nil {
|
||||
return 0, pricing.calculationError(resource, err)
|
||||
}
|
||||
return pricing.calculate(resource, amount, nil, fixedFactors...)
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) add(resource string, left fixedAmount, right fixedAmount) (fixedAmount, error) {
|
||||
amount, err := addFixedAmounts(left, right)
|
||||
if err != nil {
|
||||
return 0, pricing.calculationError(resource, err)
|
||||
}
|
||||
return amount, nil
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) calculationError(resource string, err error) error {
|
||||
return &PricingUnavailableError{Reason: "pricing calculation failed: " + err.Error(), ResourceType: resource, RuleSetID: pricing.RuleSetID}
|
||||
}
|
||||
|
||||
func pricingWeightValue(config map[string]any, resource string, keys []string, names []string) (any, bool) {
|
||||
if value, ok := pricingWeightValueFromConfig(config, keys, names); ok {
|
||||
return value, true
|
||||
}
|
||||
resourceConfig, _ := config[resource].(map[string]any)
|
||||
if len(resourceConfig) == 0 && resource == "image_edit" {
|
||||
resourceConfig, _ = config["image"].(map[string]any)
|
||||
}
|
||||
return pricingWeightValueFromConfig(resourceConfig, keys, names)
|
||||
}
|
||||
|
||||
func pricingWeightValueFromConfig(config map[string]any, keys []string, names []string) (any, bool) {
|
||||
if len(config) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
for _, key := range keys {
|
||||
weights, _ := config[key].(map[string]any)
|
||||
for _, name := range names {
|
||||
if value, ok := weights[name]; ok {
|
||||
return value, true
|
||||
}
|
||||
}
|
||||
}
|
||||
dynamic, _ := config["dynamicWeight"].(map[string]any)
|
||||
for _, name := range names {
|
||||
if value, ok := dynamic[name]; ok {
|
||||
return value, true
|
||||
}
|
||||
}
|
||||
for _, key := range keys {
|
||||
weights, _ := dynamic[key].(map[string]any)
|
||||
for _, name := range names {
|
||||
if value, ok := weights[name]; ok {
|
||||
return value, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
type candidateEstimate struct {
|
||||
Items []any
|
||||
Amount fixedAmount
|
||||
@@ -474,8 +586,14 @@ func (s *Service) billingsWithResolvedPricingV2(
|
||||
items := make([]any, 0, 3)
|
||||
total := fixedAmount(0)
|
||||
if uncachedInputTokens > 0 || cachedInputTokens == 0 {
|
||||
amount := inputPrice.MulRatio(uncachedInputTokens, 1000).Mul(discount)
|
||||
total = total.Add(amount)
|
||||
amount, calculationErr := pricing.calculateRatio("text_input", inputPrice, uncachedInputTokens, 1000, discount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
total, calculationErr = pricing.add("text", total, amount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
items = append(items, buildLine("text_input", "1k_tokens", uncachedInputTokens, amount, map[string]any{
|
||||
"inputTokens": inputTokens, "uncachedInputTokens": uncachedInputTokens,
|
||||
"cachedInputTokens": cachedInputTokens, "pricePer1k": inputPrice.Float64(),
|
||||
@@ -486,8 +604,14 @@ func (s *Service) billingsWithResolvedPricingV2(
|
||||
if priceErr != nil {
|
||||
return nil, 0, resolvedPricing{}, priceErr
|
||||
}
|
||||
amount := cachedPrice.MulRatio(cachedInputTokens, 1000).Mul(discount)
|
||||
total = total.Add(amount)
|
||||
amount, calculationErr := pricing.calculateRatio("text_cached_input", cachedPrice, cachedInputTokens, 1000, discount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
total, calculationErr = pricing.add("text", total, amount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
items = append(items, buildLine("text_cached_input", "1k_tokens", cachedInputTokens, amount, map[string]any{
|
||||
"inputTokens": inputTokens, "uncachedInputTokens": uncachedInputTokens,
|
||||
"cachedInputTokens": cachedInputTokens, "pricePer1k": cachedPrice.Float64(),
|
||||
@@ -498,8 +622,14 @@ func (s *Service) billingsWithResolvedPricingV2(
|
||||
if priceErr != nil {
|
||||
return nil, 0, resolvedPricing{}, priceErr
|
||||
}
|
||||
amount := outputPrice.MulRatio(outputTokens, 1000).Mul(discount)
|
||||
total = total.Add(amount)
|
||||
amount, calculationErr := pricing.calculateRatio("text_output", outputPrice, outputTokens, 1000, discount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
total, calculationErr = pricing.add("text", total, amount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
items = append(items, buildLine("text_output", "1k_tokens", outputTokens, amount, map[string]any{"pricePer1k": outputPrice.Float64()}))
|
||||
}
|
||||
return items, total, pricing, nil
|
||||
@@ -524,12 +654,26 @@ func (s *Service) billingsWithResolvedPricingV2(
|
||||
if priceErr != nil {
|
||||
return nil, 0, resolvedPricing{}, priceErr
|
||||
}
|
||||
amount := price.MulInt(count).MulInt(durationUnits).
|
||||
Mul(fixedWeight(resourceWeight(pricing.Config, resource, "resolutionWeights", firstNonEmptyString(stringFromMap(body, "resolution"), stringFromMap(body, "size"))))).
|
||||
Mul(fixedWeight(resourceWeight(pricing.Config, resource, "audioWeights", boolWeightKey(audioEnabled)))).
|
||||
Mul(fixedWeight(resourceWeight(pricing.Config, resource, "referenceVideoWeights", boolWeightKey(requestHasReferenceVideo(body))))).
|
||||
Mul(fixedWeight(resourceWeight(pricing.Config, resource, "voiceSpecifiedWeights", boolWeightKey(requestHasVoiceID(body, audioEnabled))))).
|
||||
Mul(discount)
|
||||
resolutionWeight, weightErr := pricing.weight(resource, "resolutionWeights", firstNonEmptyString(stringFromMap(body, "resolution"), stringFromMap(body, "size")))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
audioWeight, weightErr := pricing.weight(resource, "audioWeights", boolWeightKey(audioEnabled))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
referenceVideoWeight, weightErr := pricing.weight(resource, "referenceVideoWeights", boolWeightKey(requestHasReferenceVideo(body)))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
voiceWeight, weightErr := pricing.weight(resource, "voiceSpecifiedWeights", boolWeightKey(requestHasVoiceID(body, audioEnabled)))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
amount, calculationErr := pricing.calculate(resource, price, []int{count, durationUnits}, resolutionWeight, audioWeight, referenceVideoWeight, voiceWeight, discount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
item := buildLine(resource, unit, count*durationUnits, amount, map[string]any{
|
||||
"count": count, "audio": audioEnabled, "audioSource": audioSource,
|
||||
"durationSeconds": duration, "durationSource": durationSource,
|
||||
@@ -555,14 +699,32 @@ func (s *Service) billingsWithResolvedPricingV2(
|
||||
if err != nil {
|
||||
return nil, 0, resolvedPricing{}, err
|
||||
}
|
||||
amount := price.MulInt(count)
|
||||
amount, calculationErr := pricing.calculate(resource, price, []int{count})
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
if resource == "image" || resource == "image_edit" {
|
||||
amount = amount.
|
||||
Mul(fixedWeight(resourceWeight(pricing.Config, resource, "qualityWeights", stringFromMap(body, "quality")))).
|
||||
Mul(fixedWeight(resourceWeight(pricing.Config, resource, "sizeWeights", stringFromMap(body, "size")))).
|
||||
Mul(fixedWeight(resourceWeight(pricing.Config, resource, "resolutionWeights", firstNonEmptyString(stringFromMap(body, "resolution"), stringFromMap(body, "size")))))
|
||||
qualityWeight, weightErr := pricing.weight(resource, "qualityWeights", stringFromMap(body, "quality"))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
sizeWeight, weightErr := pricing.weight(resource, "sizeWeights", stringFromMap(body, "size"))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
resolutionWeight, weightErr := pricing.weight(resource, "resolutionWeights", firstNonEmptyString(stringFromMap(body, "resolution"), stringFromMap(body, "size")))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
amount, calculationErr = pricing.calculate(resource, amount, nil, qualityWeight, sizeWeight, resolutionWeight)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
}
|
||||
amount, calculationErr = pricing.calculate(resource, amount, nil, discount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
amount = amount.Mul(discount)
|
||||
return []any{buildLine(resource, unit, count, amount, nil)}, amount, pricing, nil
|
||||
}
|
||||
|
||||
@@ -580,14 +742,6 @@ func (pricing resolvedPricing) requiredTextPrice(resource string, keys ...string
|
||||
return 0, &PricingUnavailableError{Reason: "missing, invalid, or not explicitly free", ResourceType: resource, RuleSetID: pricing.RuleSetID}
|
||||
}
|
||||
|
||||
func fixedWeight(value float64) fixedAmount {
|
||||
weight, err := fixedAmountFromAny(value)
|
||||
if err != nil || weight <= 0 {
|
||||
return fixedAmount(fixedScale)
|
||||
}
|
||||
return weight
|
||||
}
|
||||
|
||||
func maximumCandidateEstimate(estimates []candidateEstimate) (candidateEstimate, error) {
|
||||
if len(estimates) == 0 {
|
||||
return candidateEstimate{}, &PricingUnavailableError{Reason: "no candidate has effective pricing"}
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
@@ -16,11 +18,29 @@ func TestFixedAmountPreservesNineDecimalPlaces(t *testing.T) {
|
||||
}
|
||||
|
||||
price := mustFixedAmount(t, "0.000000001")
|
||||
if got, want := price.MulInt(3).String(), "0.000000003"; got != want {
|
||||
product, err := multiplyFixedAmountByInt(price, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("multiply fixed amount: %v", err)
|
||||
}
|
||||
if got, want := product.String(), "0.000000003"; got != want {
|
||||
t.Fatalf("nano amount multiplication=%q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixedAmountOperationsRejectOverflow(t *testing.T) {
|
||||
maximum := fixedAmount(math.MaxInt64)
|
||||
if _, err := multiplyFixedAmountByInt(maximum, 2); !errors.Is(err, errFixedAmountOverflow) {
|
||||
t.Fatalf("multiply overflow error=%v", err)
|
||||
}
|
||||
if _, err := addFixedAmounts(maximum, 1); !errors.Is(err, errFixedAmountOverflow) {
|
||||
t.Fatalf("add overflow error=%v", err)
|
||||
}
|
||||
pricing := resolvedPricing{RuleSetID: "overflow-rule"}
|
||||
if _, err := pricing.calculate("image", maximum, []int{2}); !isPricingUnavailable(err) {
|
||||
t.Fatalf("pricing overflow should be unavailable: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimatedOutputTokensUsesAliasesAndCapabilityFallback(t *testing.T) {
|
||||
candidate := store.RuntimeModelCandidate{
|
||||
ModelType: "text_generate",
|
||||
@@ -97,6 +117,27 @@ func TestPricingAvailabilityRequiresExplicitFree(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPricingWeightsUseFixedPrecisionAndRejectInvalidValues(t *testing.T) {
|
||||
pricing := resolvedPricing{Config: map[string]any{
|
||||
"image": map[string]any{
|
||||
"dynamicWeight": map[string]any{
|
||||
"qualityFactors": map[string]any{"high": "1.123456789", "broken": 0},
|
||||
},
|
||||
},
|
||||
}}
|
||||
weight, err := pricing.weight("image", "qualityWeights", "high")
|
||||
if err != nil || weight.String() != "1.123456789" {
|
||||
t.Fatalf("exact weight=%s err=%v", weight.String(), err)
|
||||
}
|
||||
if _, err := pricing.weight("image", "qualityWeights", "broken"); !isPricingUnavailable(err) {
|
||||
t.Fatalf("invalid configured weight should make pricing unavailable: %v", err)
|
||||
}
|
||||
defaultWeight, err := pricing.weight("image", "qualityWeights", "unconfigured")
|
||||
if err != nil || defaultWeight.String() != "1.000000000" {
|
||||
t.Fatalf("default weight=%s err=%v", defaultWeight.String(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func mustFixedAmount(t *testing.T, value string) fixedAmount {
|
||||
t.Helper()
|
||||
amount, err := parseFixedAmount(value)
|
||||
|
||||
@@ -49,6 +49,10 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs
|
||||
w.service.logger.Debug("river async task execution lease already held", "taskID", task.ID, "riverJobID", job.ID)
|
||||
return nil
|
||||
}
|
||||
if errors.Is(runErr, store.ErrTaskExecutionManualReview) {
|
||||
w.service.logger.Warn("river async task moved to manual review after ambiguous upstream submission", "taskID", task.ID, "riverJobID", job.ID)
|
||||
return nil
|
||||
}
|
||||
var queuedErr *TaskQueuedError
|
||||
if errors.As(runErr, &queuedErr) {
|
||||
return river.JobSnooze(queuedErr.Delay)
|
||||
|
||||
@@ -302,6 +302,7 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
|
||||
}
|
||||
}
|
||||
pricingByCandidate := map[string]resolvedPricing{}
|
||||
preprocessingByCandidate := map[string]parameterPreprocessResult{}
|
||||
reservationBillings := []any(nil)
|
||||
reservationPricingSnapshot := map[string]any(nil)
|
||||
if task.RunMode == "production" {
|
||||
@@ -316,9 +317,17 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
|
||||
}
|
||||
estimates := make([]candidateEstimate, 0, len(candidates))
|
||||
var pricingErr error
|
||||
var preprocessingErr error
|
||||
var preprocessingCandidate store.RuntimeModelCandidate
|
||||
for _, candidate := range candidates {
|
||||
pricingBody := preprocessRequest(task.Kind, cloneMap(body), candidate)
|
||||
estimate, estimateErr := s.estimateCandidateV2(ctx, user, task.Kind, pricingBody, candidate)
|
||||
preprocessing := s.preprocessRequestWithScripts(ctx, task.Kind, body, candidate)
|
||||
preprocessingByCandidate[pricingCandidateKey(candidate)] = preprocessing
|
||||
if preprocessing.Err != nil {
|
||||
preprocessingErr = parameterPreprocessClientError(preprocessing.Err)
|
||||
preprocessingCandidate = candidate
|
||||
break
|
||||
}
|
||||
estimate, estimateErr := s.estimateCandidateV2(ctx, user, task.Kind, preprocessing.Body, candidate)
|
||||
if estimateErr != nil {
|
||||
pricingErr = estimateErr
|
||||
if billingMode == "enforce" {
|
||||
@@ -329,8 +338,23 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
|
||||
estimates = append(estimates, estimate)
|
||||
pricingByCandidate[pricingCandidateKey(candidate)] = estimate.Pricing
|
||||
}
|
||||
if preprocessingErr != nil {
|
||||
preprocessing := preprocessingByCandidate[pricingCandidateKey(preprocessingCandidate)]
|
||||
s.recordFailedAttempt(ctx, failedAttemptRecord{
|
||||
Task: task, Body: preprocessing.Body, Candidate: &preprocessingCandidate,
|
||||
AttemptNo: task.AttemptCount + 1, Code: clients.ErrorCode(preprocessingErr), Cause: preprocessingErr,
|
||||
Simulated: false, Scope: "parameter_preprocessing", Reason: "parameter_preprocessing_failed",
|
||||
ExtraMetrics: []map[string]any{parameterPreprocessingMetrics(preprocessing.Log)}, Preprocessing: &preprocessing.Log,
|
||||
ModelType: preprocessingCandidate.ModelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, clients.ErrorCode(preprocessingErr), preprocessingErr.Error(), false, preprocessingErr, parameterPreprocessingMetrics(preprocessing.Log))
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, preprocessingErr
|
||||
}
|
||||
if billingMode == "observe" {
|
||||
legacyItems, legacyAmount := s.maximumLegacyCandidateEstimate(ctx, user, task.Kind, body, candidates)
|
||||
legacyItems, legacyAmount := s.maximumLegacyCandidateEstimate(ctx, user, task.Kind, body, candidates, preprocessingByCandidate)
|
||||
reservationBillings = legacyItems
|
||||
candidateSnapshots := make([]any, 0, len(estimates))
|
||||
for _, estimate := range estimates {
|
||||
@@ -396,7 +420,10 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
|
||||
}
|
||||
}()
|
||||
if len(candidates) > 0 {
|
||||
preprocessing := s.preprocessRequestWithScripts(ctx, task.Kind, body, candidates[0])
|
||||
preprocessing, ok := preprocessingByCandidate[pricingCandidateKey(candidates[0])]
|
||||
if !ok {
|
||||
preprocessing = s.preprocessRequestWithScripts(ctx, task.Kind, body, candidates[0])
|
||||
}
|
||||
firstCandidateBody = preprocessing.Body
|
||||
firstPreprocessing = preprocessing.Log
|
||||
normalizedModelType = candidates[0].ModelType
|
||||
@@ -472,7 +499,10 @@ candidatesLoop:
|
||||
var candidateErr error
|
||||
for clientAttempt := 1; clientAttempt <= clientAttempts; clientAttempt++ {
|
||||
nextAttemptNo := attemptNo + 1
|
||||
preprocessing := s.preprocessRequestWithScripts(ctx, task.Kind, body, candidate)
|
||||
preprocessing, ok := preprocessingByCandidate[pricingCandidateKey(candidate)]
|
||||
if !ok {
|
||||
preprocessing = s.preprocessRequestWithScripts(ctx, task.Kind, body, candidate)
|
||||
}
|
||||
preprocessingLog := preprocessing.Log
|
||||
lastPreprocessing = &preprocessingLog
|
||||
if preprocessing.Err != nil {
|
||||
@@ -768,11 +798,14 @@ func normalizedBillingEngineMode(value string) string {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) maximumLegacyCandidateEstimate(ctx context.Context, user *auth.User, kind string, body map[string]any, candidates []store.RuntimeModelCandidate) ([]any, fixedAmount) {
|
||||
func (s *Service) maximumLegacyCandidateEstimate(ctx context.Context, user *auth.User, kind string, body map[string]any, candidates []store.RuntimeModelCandidate, preprocessingByCandidate map[string]parameterPreprocessResult) ([]any, fixedAmount) {
|
||||
var maximumItems []any
|
||||
maximumAmount := fixedAmount(0)
|
||||
for index, candidate := range candidates {
|
||||
candidateBody := preprocessRequest(kind, cloneMap(body), candidate)
|
||||
if preprocessing, ok := preprocessingByCandidate[pricingCandidateKey(candidate)]; ok && preprocessing.Err == nil {
|
||||
candidateBody = preprocessing.Body
|
||||
}
|
||||
items := s.estimatedBillings(ctx, user, kind, candidateBody, candidate)
|
||||
amount := billingItemsFixedTotal(items)
|
||||
if index == 0 || amount > maximumAmount {
|
||||
@@ -792,7 +825,11 @@ func billingItemsFixedTotal(items []any) fixedAmount {
|
||||
}
|
||||
amount, err := fixedAmountFromAny(line["amount"])
|
||||
if err == nil && amount > 0 {
|
||||
total = total.Add(amount)
|
||||
next, addErr := addFixedAmounts(total, amount)
|
||||
if addErr != nil {
|
||||
return maxFixedAmount
|
||||
}
|
||||
total = next
|
||||
}
|
||||
}
|
||||
return total
|
||||
@@ -925,6 +962,11 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
|
||||
PreviousResponseTurns: responseExecution.PreviousTurns,
|
||||
})
|
||||
callFinishedAt := time.Now()
|
||||
if err == nil {
|
||||
if markErr := s.store.SetAttemptUpstreamSubmissionStatus(context.WithoutCancel(ctx), attemptID, "response_received"); markErr != nil {
|
||||
return clients.Response{}, &upstreamSubmissionUnknownError{AttemptID: attemptID, Cause: markErr}
|
||||
}
|
||||
}
|
||||
if response.ResponseStartedAt.IsZero() {
|
||||
response.ResponseStartedAt = callStartedAt
|
||||
}
|
||||
@@ -939,7 +981,9 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
|
||||
}
|
||||
if err != nil {
|
||||
if clients.ErrorResponseMetadata(err).StatusCode > 0 {
|
||||
_ = s.store.SetAttemptUpstreamSubmissionStatus(context.WithoutCancel(ctx), attemptID, "response_received")
|
||||
if markErr := s.store.SetAttemptUpstreamSubmissionStatus(context.WithoutCancel(ctx), attemptID, "response_received"); markErr != nil {
|
||||
return clients.Response{}, &upstreamSubmissionUnknownError{AttemptID: attemptID, Cause: markErr}
|
||||
}
|
||||
}
|
||||
retryable := clients.IsRetryable(err)
|
||||
requestID, metrics, responseStartedAt, responseFinishedAt, responseDurationMS := failureMetrics(err, simulated)
|
||||
|
||||
@@ -62,6 +62,149 @@ func TestTaskIdempotencyAndExecutionLease(t *testing.T) {
|
||||
if err != nil || finished.ErrorCode != "new_worker" {
|
||||
t.Fatalf("new worker terminal task=%+v err=%v", finished, err)
|
||||
}
|
||||
if err := db.RenewTaskExecutionLease(ctx, created.Task.ID, secondToken, 5*time.Minute); !errors.Is(err, ErrTaskExecutionFinished) {
|
||||
t.Fatalf("terminal task renewal error=%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpiredExecutionLeaseDoesNotReplayAmbiguousUpstreamSubmission(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-review-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
|
||||
Request: map[string]any{"model": "billing-v2-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
})
|
||||
|
||||
firstToken := uuid.NewString()
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, firstToken, 5*time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
|
||||
TaskID: created.ID, AttemptNo: 1, Status: "running", RequestSnapshot: map[string]any{"model": created.Model},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, "submitting"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE gateway_tasks SET execution_lease_expires_at=now()-interval '1 second' WHERE id=$1::uuid`, created.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, uuid.NewString(), 5*time.Minute); !errors.Is(err, ErrTaskExecutionManualReview) {
|
||||
t.Fatalf("ambiguous submission takeover error=%v", err)
|
||||
}
|
||||
|
||||
review, err := db.GetTask(ctx, created.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_submission_unknown" {
|
||||
t.Fatalf("manual review task=%+v", review)
|
||||
}
|
||||
var outboxStatus string
|
||||
var outboxAction string
|
||||
var reviewReason string
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT status, action, COALESCE(manual_review_reason, '')
|
||||
FROM settlement_outbox
|
||||
WHERE task_id=$1::uuid AND event_type='task.billing.review'`, created.ID).Scan(&outboxStatus, &outboxAction, &reviewReason); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if outboxStatus != "manual_review" || outboxAction != "release" || reviewReason != "upstream_submission_unknown" {
|
||||
t.Fatalf("review outbox status=%s action=%s reason=%s", outboxStatus, outboxAction, reviewReason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpiredExecutionLeaseCanResumeAfterKnownRejectedResponse(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
_, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-known-response-" + uuid.NewString(), GatewayUserID: gatewayUserID}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
|
||||
Request: map[string]any{"model": "billing-v2-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
})
|
||||
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, uuid.NewString(), 5*time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{TaskID: created.ID, AttemptNo: 1, Status: "running"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, "response_received"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.FinishTaskAttempt(ctx, FinishTaskAttemptInput{AttemptID: attemptID, Status: "failed", ErrorCode: "upstream_rejected"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE gateway_tasks SET execution_lease_expires_at=now()-interval '1 second' WHERE id=$1::uuid`, created.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, uuid.NewString(), 5*time.Minute); err != nil {
|
||||
t.Fatalf("known rejected response should remain retryable: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFinishTaskManualReviewCreatesVisibleBillingRecord(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
_, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-direct-review-" + uuid.NewString(), GatewayUserID: gatewayUserID}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
|
||||
Request: map[string]any{"model": "billing-v2-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
})
|
||||
token := uuid.NewString()
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, token, 5*time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{TaskID: created.ID, AttemptNo: 1, Status: "running"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, "submitting"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.FinishTaskManualReview(ctx, FinishTaskManualReviewInput{
|
||||
TaskID: created.ID, ExecutionToken: token, AttemptID: attemptID, TaskStatus: "failed",
|
||||
Code: "upstream_submission_unknown", Message: "upstream submission result is unknown",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var visible bool
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM settlement_outbox
|
||||
WHERE task_id=$1::uuid AND status='manual_review'
|
||||
AND action='release' AND manual_review_reason='upstream_submission_unknown'
|
||||
)`, created.ID).Scan(&visible); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !visible {
|
||||
t.Fatal("manual review billing record is not visible in settlement outbox")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingSettlementStaleTakeoverDebitsExactlyOnce(t *testing.T) {
|
||||
|
||||
@@ -59,6 +59,8 @@ var (
|
||||
ErrIdempotencyKeyReused = errors.New("idempotency key was reused for a different request")
|
||||
ErrTaskExecutionLeaseUnavailable = errors.New("task execution lease is unavailable")
|
||||
ErrTaskExecutionLeaseLost = errors.New("task execution lease was lost")
|
||||
ErrTaskExecutionFinished = errors.New("task execution already finished")
|
||||
ErrTaskExecutionManualReview = errors.New("task execution requires manual review")
|
||||
ErrProtectedDefault = errors.New("protected default resource cannot be deleted")
|
||||
ErrUserAlreadyExists = errors.New("user already exists")
|
||||
ErrWeakPassword = errors.New("password must be at least 8 characters")
|
||||
|
||||
@@ -335,6 +335,9 @@ ORDER BY priority ASC, resource_type ASC, rule_key ASC`, id)
|
||||
default:
|
||||
return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s uses unsupported calculator %s", ruleKey, calculatorType)
|
||||
}
|
||||
if err := ValidateEffectivePricingRuleShape(resourceType, unit, calculatorType); err != nil {
|
||||
return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s: %w", ruleKey, err)
|
||||
}
|
||||
if strings.HasPrefix(strings.TrimSpace(basePrice), "-") {
|
||||
return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s has a negative base price", ruleKey)
|
||||
}
|
||||
@@ -372,6 +375,51 @@ ORDER BY priority ASC, resource_type ASC, rule_key ASC`, id)
|
||||
}, nil
|
||||
}
|
||||
|
||||
func ValidateEffectivePricingRuleShape(resourceType string, unit string, calculatorType string) error {
|
||||
unit = NormalizeEffectivePricingRuleUnit(resourceType, unit)
|
||||
if strings.TrimSpace(calculatorType) == "" {
|
||||
calculatorType = DefaultEffectivePricingCalculator(resourceType)
|
||||
}
|
||||
allowed := false
|
||||
switch resourceType {
|
||||
case "text_input", "text_cached_input", "text_output", "text_total":
|
||||
allowed = unit == "1k_tokens" && calculatorType == "token_usage"
|
||||
case "image", "image_edit":
|
||||
allowed = unit == "image" && calculatorType == "unit_weight"
|
||||
case "video":
|
||||
allowed = unit == "5s" && calculatorType == "duration_weight"
|
||||
case "music":
|
||||
allowed = (unit == "song" || unit == "item") && calculatorType == "unit_weight"
|
||||
case "audio":
|
||||
allowed = unit == "character" && calculatorType == "unit_weight"
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
if !allowed {
|
||||
return fmt.Errorf("unit %q and calculator %q do not match resource %q", unit, calculatorType, resourceType)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func NormalizeEffectivePricingRuleUnit(resourceType string, unit string) string {
|
||||
unit = strings.TrimSpace(unit)
|
||||
if resourceType == "video" && unit == "video" {
|
||||
return "5s"
|
||||
}
|
||||
return unit
|
||||
}
|
||||
|
||||
func DefaultEffectivePricingCalculator(resourceType string) string {
|
||||
switch strings.TrimSpace(resourceType) {
|
||||
case "text_input", "text_cached_input", "text_output", "text_total":
|
||||
return "token_usage"
|
||||
case "video":
|
||||
return "duration_weight"
|
||||
default:
|
||||
return "unit_weight"
|
||||
}
|
||||
}
|
||||
|
||||
func addPricingRuleToConfig(config map[string]any, resourceType string, basePrice string, dynamicWeight map[string]any, formulaConfig map[string]any) {
|
||||
resourceConfig := map[string]any{"basePrice": basePrice}
|
||||
if len(dynamicWeight) > 0 {
|
||||
@@ -591,7 +639,7 @@ func normalizePricingRule(input PricingRuleInput, index int, defaultCurrency str
|
||||
input.RuleKey = strings.TrimSpace(input.RuleKey)
|
||||
input.DisplayName = strings.TrimSpace(input.DisplayName)
|
||||
input.ResourceType = strings.TrimSpace(input.ResourceType)
|
||||
input.Unit = strings.TrimSpace(input.Unit)
|
||||
input.Unit = NormalizeEffectivePricingRuleUnit(input.ResourceType, input.Unit)
|
||||
input.Currency = strings.TrimSpace(input.Currency)
|
||||
input.CalculatorType = strings.TrimSpace(input.CalculatorType)
|
||||
input.Status = strings.TrimSpace(input.Status)
|
||||
@@ -608,7 +656,7 @@ func normalizePricingRule(input PricingRuleInput, index int, defaultCurrency str
|
||||
input.Currency = defaultCurrency
|
||||
}
|
||||
if input.CalculatorType == "" {
|
||||
input.CalculatorType = "unit_weight"
|
||||
input.CalculatorType = DefaultEffectivePricingCalculator(input.ResourceType)
|
||||
}
|
||||
if input.Priority == 0 {
|
||||
input.Priority = (index + 1) * 10
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
package store
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestValidateEffectivePricingRuleShape(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
resource string
|
||||
unit string
|
||||
calculator string
|
||||
valid bool
|
||||
}{
|
||||
{name: "text", resource: "text_input", unit: "1k_tokens", calculator: "token_usage", valid: true},
|
||||
{name: "image", resource: "image", unit: "image", calculator: "unit_weight", valid: true},
|
||||
{name: "video", resource: "video", unit: "5s", calculator: "duration_weight", valid: true},
|
||||
{name: "wrong video unit", resource: "video", unit: "second", calculator: "duration_weight"},
|
||||
{name: "wrong image calculator", resource: "image", unit: "image", calculator: "token_usage"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
err := ValidateEffectivePricingRuleShape(test.resource, test.unit, test.calculator)
|
||||
if (err == nil) != test.valid {
|
||||
t.Fatalf("valid=%v err=%v", test.valid, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -157,7 +157,83 @@ func (s *Store) ClaimTaskExecution(ctx context.Context, taskID string, execution
|
||||
if leaseTTL <= 0 {
|
||||
leaseTTL = 5 * time.Minute
|
||||
}
|
||||
task, err := scanGatewayTask(s.pool.QueryRow(ctx, `
|
||||
var task GatewayTask
|
||||
manualReview := false
|
||||
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
var queuedReady bool
|
||||
var runningExpired bool
|
||||
var production bool
|
||||
var hasGatewayUser bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT status = 'queued' AND next_run_at <= now(),
|
||||
status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now()),
|
||||
run_mode = 'production',
|
||||
gateway_user_id IS NOT NULL
|
||||
FROM gateway_tasks
|
||||
WHERE id = $1::uuid
|
||||
FOR UPDATE`, taskID).Scan(&queuedReady, &runningExpired, &production, &hasGatewayUser); err != nil {
|
||||
return err
|
||||
}
|
||||
if !queuedReady && !runningExpired {
|
||||
return ErrTaskExecutionLeaseUnavailable
|
||||
}
|
||||
if runningExpired && production {
|
||||
var submissionAmbiguous bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE((
|
||||
SELECT (
|
||||
(status = 'running' AND upstream_submission_status IN ('submitting', 'response_received'))
|
||||
OR (status = 'failed' AND upstream_submission_status = 'submitting')
|
||||
)
|
||||
FROM gateway_task_attempts
|
||||
WHERE task_id = $1::uuid
|
||||
ORDER BY attempt_no DESC, started_at DESC
|
||||
LIMIT 1
|
||||
), false)`, taskID).Scan(&submissionAmbiguous); err != nil {
|
||||
return err
|
||||
}
|
||||
if submissionAmbiguous {
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'failed',
|
||||
billing_status = CASE WHEN $2 THEN 'manual_review' ELSE 'not_required' END,
|
||||
billing_updated_at = now(),
|
||||
error = 'upstream submission result is unknown',
|
||||
error_code = 'upstream_submission_unknown',
|
||||
error_message = 'upstream submission result is unknown',
|
||||
locked_by = NULL,
|
||||
locked_at = NULL,
|
||||
heartbeat_at = NULL,
|
||||
execution_token = NULL,
|
||||
execution_lease_expires_at = NULL,
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil {
|
||||
return err
|
||||
}
|
||||
if hasGatewayUser {
|
||||
payloadJSON, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID, "classification": "upstream_submission_unknown",
|
||||
})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO settlement_outbox (
|
||||
task_id, event_type, action, amount, currency, pricing_snapshot, payload,
|
||||
status, next_attempt_at, manual_review_reason
|
||||
)
|
||||
SELECT id, 'task.billing.review', 'release', reservation_amount, billing_currency,
|
||||
pricing_snapshot, $2::jsonb, 'manual_review', now(), 'upstream_submission_unknown'
|
||||
FROM gateway_tasks
|
||||
WHERE id = $1::uuid
|
||||
ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
manualReview = true
|
||||
return nil
|
||||
}
|
||||
}
|
||||
var err error
|
||||
task, err = scanGatewayTask(tx.QueryRow(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'running',
|
||||
execution_token = $2::uuid,
|
||||
@@ -171,10 +247,18 @@ WHERE id = $1::uuid
|
||||
OR (status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now()))
|
||||
)
|
||||
RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second)))
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return GatewayTask{}, ErrTaskExecutionLeaseUnavailable
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return GatewayTask{}, ErrTaskExecutionLeaseUnavailable
|
||||
}
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
return task, err
|
||||
if manualReview {
|
||||
return GatewayTask{}, ErrTaskExecutionManualReview
|
||||
}
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func (s *Store) RenewTaskExecutionLease(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) error {
|
||||
@@ -193,6 +277,14 @@ WHERE id = $1::uuid
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
var stillRunning bool
|
||||
if err := s.pool.QueryRow(ctx, `
|
||||
SELECT COALESCE((SELECT status = 'running' FROM gateway_tasks WHERE id = $1::uuid), false)`, taskID).Scan(&stillRunning); err != nil {
|
||||
return err
|
||||
}
|
||||
if !stillRunning {
|
||||
return ErrTaskExecutionFinished
|
||||
}
|
||||
return ErrTaskExecutionLeaseLost
|
||||
}
|
||||
return nil
|
||||
@@ -902,8 +994,7 @@ SET status = $2,
|
||||
response_started_at = $6::timestamptz,
|
||||
response_finished_at = $7::timestamptz,
|
||||
response_duration_ms = $8,
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
finished_at = now()
|
||||
WHERE id = $1::uuid`, input.AttemptID, attemptStatus, input.RequestID, input.Code, input.Message,
|
||||
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS); err != nil {
|
||||
return err
|
||||
@@ -942,7 +1033,22 @@ WHERE id = $1::uuid
|
||||
if tag.RowsAffected() != 1 {
|
||||
return ErrTaskExecutionLeaseLost
|
||||
}
|
||||
return nil
|
||||
payloadJSON, _ := json.Marshal(map[string]any{
|
||||
"taskId": input.TaskID, "classification": input.Code,
|
||||
})
|
||||
_, err = tx.Exec(ctx, `
|
||||
INSERT INTO settlement_outbox (
|
||||
task_id, event_type, action, amount, currency, pricing_snapshot, payload,
|
||||
status, next_attempt_at, manual_review_reason
|
||||
)
|
||||
SELECT id, 'task.billing.review', 'release', reservation_amount, billing_currency,
|
||||
pricing_snapshot, $2::jsonb, 'manual_review', now(), $3
|
||||
FROM gateway_tasks
|
||||
WHERE id = $1::uuid
|
||||
AND run_mode = 'production'
|
||||
AND gateway_user_id IS NOT NULL
|
||||
ON CONFLICT (task_id, event_type) DO NOTHING`, input.TaskID, string(payloadJSON), input.Code)
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
|
||||
@@ -95,6 +95,20 @@ ALTER TABLE model_pricing_rules
|
||||
|
||||
UPDATE model_pricing_rules SET is_free = false WHERE is_free IS NULL;
|
||||
|
||||
UPDATE model_pricing_rules
|
||||
SET calculator_type = CASE
|
||||
WHEN resource_type IN ('text_input', 'text_cached_input', 'text_output', 'text_total') THEN 'token_usage'
|
||||
WHEN resource_type = 'video' THEN 'duration_weight'
|
||||
ELSE calculator_type
|
||||
END,
|
||||
unit = CASE
|
||||
WHEN resource_type = 'video' AND unit = 'video' THEN '5s'
|
||||
ELSE unit
|
||||
END,
|
||||
updated_at = now()
|
||||
WHERE (resource_type IN ('text_input', 'text_cached_input', 'text_output', 'text_total') AND calculator_type = 'unit_weight')
|
||||
OR (resource_type = 'video' AND (calculator_type = 'unit_weight' OR unit = 'video'));
|
||||
|
||||
ALTER TABLE model_pricing_rules
|
||||
ADD CONSTRAINT model_pricing_rules_explicit_free_v2_check CHECK (
|
||||
base_price > 0 OR (base_price = 0 AND is_free)
|
||||
@@ -152,6 +166,30 @@ WHERE task.status = 'succeeded'
|
||||
AND transaction.transaction_type = 'task_billing'
|
||||
);
|
||||
|
||||
INSERT INTO settlement_outbox (
|
||||
task_id, event_type, action, amount, currency, pricing_snapshot, payload,
|
||||
status, next_attempt_at, manual_review_reason
|
||||
)
|
||||
SELECT task.id,
|
||||
'task.billing.review',
|
||||
'settle',
|
||||
COALESCE(task.final_charge_amount, 0),
|
||||
COALESCE(NULLIF(task.billing_summary->>'currency', ''), 'resource'),
|
||||
COALESCE(task.pricing_snapshot, '{}'::jsonb),
|
||||
jsonb_build_object(
|
||||
'taskId', task.id,
|
||||
'classification', 'historical_success_without_charge'
|
||||
),
|
||||
'manual_review',
|
||||
now(),
|
||||
'historical_success_without_charge'
|
||||
FROM gateway_tasks task
|
||||
WHERE task.billing_status = 'manual_review'
|
||||
AND task.status = 'succeeded'
|
||||
AND task.run_mode = 'production'
|
||||
AND task.gateway_user_id IS NOT NULL
|
||||
ON CONFLICT (task_id, event_type) DO NOTHING;
|
||||
|
||||
INSERT INTO settlement_outbox (
|
||||
task_id, event_type, action, amount, currency, payload, status, next_attempt_at
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user