From 3fbad6a72f7745b1d9c3da4b6281eb25b77f93c4 Mon Sep 17 00:00:00 2001
From: CaIon <i@caion.me>
Date: Tue, 7 Jul 2026 15:41:47 +0800
Subject: [PATCH] fix(price): add default token estimate for tiered expression
 pre-consume

---
 relay/helper/price.go      | 11 ++++--
 relay/helper/price_test.go | 78 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+), 3 deletions(-)

diff --git a/relay/helper/price.go b/relay/helper/price.go
index 313424fb..e2bff583 100644
--- a/relay/helper/price.go
+++ b/relay/helper/price.go
@@ -35,6 +35,11 @@ func modelPriceNotConfiguredError(modelName string, userId int) error {
 // https://docs.claude.com/en/docs/build-with-claude/prompt-caching#1-hour-cache-duration
 const claudeCacheCreation1hMultiplier = 6 / 3.75
 
+// defaultTieredPreConsumeMaxTokens is the fallback completion-token estimate
+// used for tiered expression pre-consume when the client omits max_tokens, so
+// the pre-consumed quota still reflects a plausible output cost in paid groups.
+const defaultTieredPreConsumeMaxTokens = 8192
+
 // HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
 func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) types.GroupRatioInfo {
 	groupRatioInfo := types.GroupRatioInfo{
@@ -244,9 +249,9 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT
 		return types.PriceData{}, fmt.Errorf("model %s is configured as tiered_expr but has no billing expression", info.OriginModelName)
 	}
 
-	estimatedCompletionTokens := 0
-	if meta.MaxTokens != 0 {
-		estimatedCompletionTokens = meta.MaxTokens
+	estimatedCompletionTokens := meta.MaxTokens
+	if estimatedCompletionTokens == 0 && groupRatioInfo.GroupRatio != 0 {
+		estimatedCompletionTokens = defaultTieredPreConsumeMaxTokens
 	}
 
 	requestInput, err := ResolveIncomingBillingExprRequestInput(c, info)
diff --git a/relay/helper/price_test.go b/relay/helper/price_test.go
index afa64c4b..a8410e81 100644
--- a/relay/helper/price_test.go
+++ b/relay/helper/price_test.go
@@ -60,3 +60,81 @@ func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) {
 	require.Equal(t, billing_setting.BillingModeTieredExpr, info.TieredBillingSnapshot.BillingMode)
 	require.Equal(t, common.QuotaPerUnit, info.TieredBillingSnapshot.QuotaPerUnit)
 }
+
+func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	saved := map[string]string{}
+	require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
+		saved[key] = value
+		return nil
+	}))
+	t.Cleanup(func() {
+		require.NoError(t, config.GlobalConfig.LoadFromDB(saved))
+	})
+
+	require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
+		"billing_setting.billing_mode":    `{"tiered-fallback-model":"tiered_expr"}`,
+		"billing_setting.billing_expr":    `{"tiered-fallback-model":"tier(\"base\", p * 3 + c * 15)"}`,
+		"group_ratio_setting.group_ratio": `{"default":1,"free":0}`,
+	}))
+
+	const promptTokens = 1000
+
+	cases := []struct {
+		name      string
+		group     string
+		maxTokens int
+		expected  int
+	}{
+		{
+			// max_tokens omitted in a paid group -> fall back to 8192 completion tokens.
+			// p*3 + c*15 = 1000*3 + 8192*15 = 125880 -> /1e6 * 500000 = 62940
+			name:      "non-free group falls back to 8192 completion tokens",
+			group:     "default",
+			maxTokens: 0,
+			expected:  62940,
+		},
+		{
+			// explicit max_tokens is used verbatim, no fallback.
+			// 1000*3 + 100*15 = 4500 -> /1e6 * 500000 = 2250
+			name:      "explicit max_tokens is used verbatim",
+			group:     "default",
+			maxTokens: 100,
+			expected:  2250,
+		},
+		{
+			// free group (ratio 0) stays zero; fallback is gated on non-zero group ratio.
+			name:      "free group stays zero without fallback",
+			group:     "free",
+			maxTokens: 0,
+			expected:  0,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			recorder := httptest.NewRecorder()
+			ctx, _ := gin.CreateTestContext(recorder)
+			req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
+			req.Header.Set("Content-Type", "application/json")
+			ctx.Request = req
+			ctx.Set("group", tc.group)
+
+			info := &relaycommon.RelayInfo{
+				OriginModelName: "tiered-fallback-model",
+				UserGroup:       tc.group,
+				UsingGroup:      tc.group,
+				RequestHeaders:  map[string]string{"Content-Type": "application/json"},
+				BillingRequestInput: &billingexpr.RequestInput{
+					Headers: map[string]string{"Content-Type": "application/json"},
+					Body:    []byte(`{}`),
+				},
+			}
+
+			priceData, err := ModelPriceHelper(ctx, info, promptTokens, &types.TokenCountMeta{MaxTokens: tc.maxTokens})
+			require.NoError(t, err)
+			require.Equal(t, tc.expected, priceData.QuotaToPreConsume)
+		})
+	}
+}