feat: add model and group retry time budgets

config.example.yaml

@@ -141,6 +141,8 @@ modelconfigs:
     rpm: 3500
     tpm: 80000
     retry_times: 3
+    # Omit to inherit the global budget; 0 disables it. Range: 0-180 seconds.
+    retry_budget: 60
     timeout_config:
       request_timeout: 300
       stream_request_timeout: 600
@@ -240,6 +242,7 @@ options:
 
   # Retry settings
   RetryTimes: "3"
+  RetryBudget: "0"  # Seconds; RETRY_BUDGET overrides this value, maximum 180.
 
   # Group settings
   GroupMaxTokenNum: "0"  # 0 means unlimited

config.md

@@ -264,6 +264,7 @@ options:
 
   # Retry settings
   RetryTimes: "3"
+  RetryBudget: "0"  # Seconds, maximum 180; 0 disables the time limit
 
   # Error rate alerts
   DefaultWarnNotifyErrorRate: "0.5"
@@ -285,12 +286,28 @@ options:
 - `LogDetailResponseBodyMaxSize`: Max size of response body to log
 - `DisableServe`: Disable API serving (for maintenance)
 - `RetryTimes`: Number of retry attempts
+- `RetryBudget`: Time budget in seconds, from 0 to 180; 0 disables the time limit. The `RETRY_BUDGET` environment variable overrides this global default and is capped at 180 seconds.
 - `DefaultChannelModels`: Default models for new channels (JSON array)
 - `GroupMaxTokenNum`: Max tokens per group
 - `DefaultWarnNotifyErrorRate`: Default error rate warning threshold
 - `UsageAlertThreshold`: Usage alert threshold
 - `FuzzyTokenThreshold`: Fuzzy token matching threshold
 
+### Retry Limits
+
+The budget starts with the first upstream attempt and includes upstream calls and retry backoff. After it expires, no further retry starts. An in-flight call keeps its existing request or stream timeout.
+
+Models can set `retry_budget` to override the global budget, set it to `0` to disable the budget, or omit it to inherit. Group model configs use `override_retry_budget` and `retry_budget` to override the model value, including zero. Retry counts continue to use `retry_times` and `override_retry_times`.
+
+| Effective configuration | Behavior |
+| --- | --- |
+| Count only | Stop after the configured number of retries |
+| Budget only | Retry until the budget expires, with no count limit |
+| Count and budget | Stop when either limit is reached |
+| Neither | No retries |
+
+A model that explicitly sets a positive budget and has no positive retry count uses only that budget, without inheriting the global count. With no local budget, global count and budget defaults apply together. A group can clear an inherited model count with `override_retry_times: true` and `retry_times: 0`; when a budget is active, this enables retries for the remaining budget. Without a budget, a zero model count retains the existing global count fallback. Permission failures exclude the channel but do not increase the retry limit. A retry count excludes the initial attempt.
+
 ## Example: Complete Configuration
 
 See `config.example.yaml` for a complete example configuration file.

core/common/config/config.go

@@ -9,6 +9,8 @@ import (
 	"github.com/labring/aiproxy/core/common/env"
 )
 
+const MaxRetryBudgetSeconds = 180
+
 var (
 	disableServe                 atomic.Bool
 	logStorageHours              atomic.Int64 // default 0 means no limit
@@ -22,6 +24,7 @@ var (
 	ipGroupsThreshold            atomic.Int64
 	ipGroupsBanThreshold         atomic.Int64
 	retryTimes                   atomic.Int64
+	retryBudget                  atomic.Int64
 	defaultChannelModels         atomic.Value
 	defaultChannelModelMapping   atomic.Value
 	groupMaxTokenNum             atomic.Int64
@@ -65,6 +68,15 @@ func SetRetryTimes(times int64) {
 	retryTimes.Store(times)
 }
 
+func GetRetryBudget() int64 {
+	return retryBudget.Load()
+}
+
+func SetRetryBudget(seconds int64) {
+	seconds = env.Int64("RETRY_BUDGET", seconds)
+	retryBudget.Store(min(max(seconds, 0), MaxRetryBudgetSeconds))
+}
+
 func GetLogStorageHours() int64 {
 	return logStorageHours.Load()
 }

core/common/config/env.go

@@ -29,6 +29,8 @@ var (
 )
 
 func ReloadEnv() {
+	SetRetryBudget(GetRetryBudget())
+
 	DebugEnabled = env.Bool("DEBUG", false)
 	DebugSQLEnabled = env.Bool("DEBUG_SQL", false)
 	DisableAutoMigrateDB = env.Bool("DISABLE_AUTO_MIGRATE_DB", false)

core/common/config/retry_budget_test.go

@@ -0,0 +1,34 @@
+package config_test
+
+import (
+	"testing"
+
+	"github.com/labring/aiproxy/core/common/config"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRetryBudgetEnvironmentOverrideAndCap(t *testing.T) {
+	t.Setenv("RETRY_BUDGET", "")
+
+	oldBudget := config.GetRetryBudget()
+	t.Cleanup(func() { config.SetRetryBudget(oldBudget) })
+
+	for _, tt := range []struct {
+		env   string
+		value int64
+		want  int64
+	}{
+		{value: 30, want: 30},
+		{env: "60", value: 30, want: 60},
+		{env: "0", value: 30},
+		{env: "181", want: 180},
+		{env: "-1"},
+		{value: 999, want: 180},
+	} {
+		t.Run(tt.env, func(t *testing.T) {
+			t.Setenv("RETRY_BUDGET", tt.env)
+			config.SetRetryBudget(tt.value)
+			require.Equal(t, tt.want, config.GetRetryBudget())
+		})
+	}
+}

core/controller/group.go

@@ -465,6 +465,9 @@ type SaveGroupModelConfigRequest struct {
 	OverrideRetryTimes bool  `json:"override_retry_times"`
 	RetryTimes         int64 `json:"retry_times"`
 
+	OverrideRetryBudget bool  `json:"override_retry_budget"`
+	RetryBudget         int64 `json:"retry_budget"          binding:"gte=0,lte=180" minimum:"0" maximum:"180"`
+
 	OverrideTimeoutConfig bool                `json:"override_timeout_config"`
 	TimeoutConfig         model.TimeoutConfig `json:"timeout_config"`
 
@@ -508,6 +511,9 @@ func (r *SaveGroupModelConfigRequest) ToGroupModelConfig(groupID string) model.G
 		OverrideRetryTimes: r.OverrideRetryTimes,
 		RetryTimes:         r.RetryTimes,
 
+		OverrideRetryBudget: r.OverrideRetryBudget,
+		RetryBudget:         r.RetryBudget,
+
 		OverrideTimeoutConfig: r.OverrideTimeoutConfig,
 		TimeoutConfig:         r.TimeoutConfig,
 

core/controller/relay-budget_test.go

@@ -0,0 +1,212 @@
+//nolint:testpackage
+package controller
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/labring/aiproxy/core/common/config"
+	"github.com/labring/aiproxy/core/middleware"
+	"github.com/labring/aiproxy/core/model"
+	relaycontroller "github.com/labring/aiproxy/core/relay/controller"
+	"github.com/labring/aiproxy/core/relay/meta"
+	"github.com/labring/aiproxy/core/relay/mode"
+	relaymodel "github.com/labring/aiproxy/core/relay/model"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRetryBudgetStartsWithFirstAttempt(t *testing.T) {
+	t.Parallel()
+
+	started := time.Now().Add(-time.Minute)
+	times, deadline := getRetryLimits(model.ModelConfig{RetryBudget: new(int64(30))}, 3, 0, started)
+	require.Equal(t, -1, times)
+	require.Equal(t, started.Add(30*time.Second), deadline)
+
+	c, _ := gin.CreateTestContext(httptest.NewRecorder())
+	c.Request = httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/", nil)
+	bizErr := relaymodel.NewOpenAIError(
+		http.StatusBadGateway,
+		relaymodel.OpenAIError{Message: "upstream failed"},
+	)
+	require.True(t, handleRelayResult(c, bizErr, true, times, deadline))
+	require.True(t, c.Writer.Written())
+}
+
+func TestRetryLoopBudgetAndCount(t *testing.T) {
+	t.Setenv("LOG_STORAGE_HOURS", "")
+
+	previous := config.GetLogStorageHours()
+
+	config.SetLogStorageHours(-1)
+	t.Cleanup(func() { config.SetLogStorageHours(previous) })
+
+	for _, tt := range []struct {
+		name                string
+		times               int
+		budget              time.Duration
+		attemptDuration     time.Duration
+		status              int
+		wantAttempts        int
+		initialBackoff      bool
+		cancelDuringBackoff bool
+		succeed             bool
+	}{
+		{name: "count only", times: 2, attemptDuration: time.Second, status: http.StatusBadGateway, wantAttempts: 2},
+		{name: "budget only", times: -1, budget: 3 * time.Second, attemptDuration: time.Second, status: http.StatusBadGateway, wantAttempts: 3},
+		{name: "count expires first", times: 2, budget: 10 * time.Second, attemptDuration: time.Second, status: http.StatusBadGateway, wantAttempts: 2},
+		{name: "budget expires first", times: 10, budget: 2 * time.Second, attemptDuration: time.Second, status: http.StatusBadGateway, wantAttempts: 2},
+		{name: "deadline expires during backoff", times: -1, budget: 500 * time.Millisecond, status: http.StatusTooManyRequests, initialBackoff: true},
+		{name: "cancellation interrupts backoff", times: 2, status: http.StatusTooManyRequests, initialBackoff: true, cancelDuringBackoff: true},
+		{name: "in flight call completes after budget", times: -1, budget: time.Second, attemptDuration: 2 * time.Second, status: http.StatusBadGateway, wantAttempts: 1, succeed: true},
+		{name: "permission errors obey count", times: 2, budget: time.Minute, status: http.StatusUnauthorized, wantAttempts: 2},
+		{name: "non retryable error stops", times: -1, budget: time.Minute, status: http.StatusBadRequest, wantAttempts: 1},
+		{name: "no count or budget", status: http.StatusBadGateway},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			synctest.Test(t, func(t *testing.T) {
+				ctx, cancel := context.WithCancel(t.Context())
+				defer cancel()
+
+				recorder := httptest.NewRecorder()
+				c, _ := gin.CreateTestContext(recorder)
+				c.Request = httptest.NewRequestWithContext(
+					ctx,
+					http.MethodPost,
+					"/",
+					strings.NewReader(`{}`),
+				)
+				c.Set(middleware.Group, model.GroupCache{})
+				c.Set(middleware.Token, model.TokenCache{})
+				c.Set(middleware.ModelConfig, model.ModelConfig{})
+				c.Set(middleware.RequestModel, "retry-budget-test")
+				c.Set(middleware.GroupBalance, &middleware.GroupBalanceConsumer{})
+				middleware.SetRequestAt(c, time.Now())
+
+				channels := []*model.Channel{
+					{ID: 1, Status: model.ChannelStatusEnabled},
+					{ID: 2, Status: model.ChannelStatusEnabled, BackupOnly: true},
+					{ID: 3, Status: model.ChannelStatusEnabled, BackupOnly: true},
+					{ID: 4, Status: model.ChannelStatusEnabled, BackupOnly: true},
+				}
+				initial := &initialChannel{
+					channel:          channels[0],
+					migratedChannels: channels,
+					preferChannelIDs: []int{1, 2, 3, 4},
+				}
+				initialError := relaymodel.NewOpenAIError(
+					http.StatusBadGateway,
+					relaymodel.OpenAIError{Message: "initial failure"},
+				)
+
+				state := initRetryState(
+					tt.times,
+					initial,
+					NewMetaByContext(c, channels[0], mode.Responses),
+					&relaycontroller.HandleResult{Error: initialError},
+					model.Price{},
+					time.Now(),
+				)
+				if tt.budget > 0 {
+					state.retryDeadline = time.Now().Add(tt.budget)
+				}
+
+				if tt.initialBackoff {
+					state.recordChannelFailure(1, time.Now())
+				}
+
+				if tt.cancelDuringBackoff {
+					go func() {
+						time.Sleep(200 * time.Millisecond)
+						cancel()
+					}()
+				}
+
+				attempts := 0
+				started := time.Now()
+				retryLoop(
+					c,
+					mode.Responses,
+					state,
+					func(c *gin.Context, _ *meta.Meta) *relaycontroller.HandleResult {
+						attempts++
+
+						if !state.retryDeadline.IsZero() {
+							require.True(t, time.Now().Before(state.retryDeadline))
+						}
+
+						time.Sleep(tt.attemptDuration)
+						require.NoError(t, c.Request.Context().Err())
+
+						if tt.succeed {
+							return &relaycontroller.HandleResult{}
+						}
+
+						return &relaycontroller.HandleResult{
+							Error: relaymodel.NewOpenAIError(
+								tt.status,
+								relaymodel.OpenAIError{Message: "retry failure"},
+							),
+						}
+					},
+				)
+				synctest.Wait()
+				require.Equal(t, tt.wantAttempts, attempts)
+				require.Equal(t, tt.times, state.retryTimes)
+
+				if tt.initialBackoff {
+					wantElapsed := tt.budget
+					if tt.cancelDuringBackoff {
+						wantElapsed = 200 * time.Millisecond
+					}
+
+					require.Equal(t, wantElapsed, time.Since(started))
+					require.Equal(t, initialError, state.result.Error)
+				}
+
+				if tt.succeed {
+					require.Nil(t, state.result.Error)
+				} else {
+					require.True(t, c.Writer.Written())
+				}
+			})
+		})
+	}
+}
+
+func TestRetryBudgetRequestBinding(t *testing.T) {
+	t.Parallel()
+
+	for _, value := range []string{"-1", "181", "1.5"} {
+		c, _ := gin.CreateTestContext(httptest.NewRecorder())
+		c.Request = httptest.NewRequestWithContext(
+			t.Context(),
+			http.MethodPost,
+			"/",
+			strings.NewReader(`{"retry_budget":`+value+`}`),
+		)
+
+		var request SaveModelConfigsRequest
+		require.Error(t, c.ShouldBindJSON(&request))
+		c.Request = httptest.NewRequestWithContext(
+			t.Context(),
+			http.MethodPost,
+			"/",
+			strings.NewReader(`{"retry_budget":`+value+`}`),
+		)
+
+		var groupRequest SaveGroupModelConfigRequest
+		require.Error(t, c.ShouldBindJSON(&groupRequest))
+	}
+
+	request := SaveGroupModelConfigRequest{OverrideRetryBudget: true, RetryBudget: 60}
+	group := request.ToGroupModelConfig("test")
+	require.True(t, group.OverrideRetryBudget)
+	require.Equal(t, int64(60), group.RetryBudget)
+}