From 9c15b043e35b9f1edc8e7064ce9a9c3ec5bf67ab Mon Sep 17 00:00:00 2001 From: zijiren233 Date: Sun, 6 Sep 2026 19:06:46 +0800 Subject: [PATCH] feat: add model and group retry time budgets --- config.example.yaml | 3 + config.md | 17 ++ core/common/config/config.go | 12 + core/common/config/env.go | 2 + core/common/config/retry_budget_test.go | 34 +++ core/controller/group.go | 6 + core/controller/relay-budget_test.go | 212 ++++++++++++++++++ core/controller/relay-controller.go | 123 ++++++---- core/controller/relay-controller_test.go | 7 +- core/docs/docs.go | 34 +++ core/docs/swagger.json | 34 +++ core/docs/swagger.yaml | 27 +++ core/model/groupmodel.go | 16 ++ core/model/modelconfig.go | 39 ++++ core/model/option.go | 15 ++ core/model/retry_budget_test.go | 141 ++++++++++++ web/public/locales/en/translation.json | 3 + web/public/locales/zh/translation.json | 3 + .../group/components/GroupModelConfigsTab.tsx | 43 +++- .../model/components/BuiltinModelsDialog.tsx | 1 + .../feature/model/components/ModelForm.tsx | 29 +++ web/src/types/group.ts | 4 + web/src/types/model.ts | 2 + web/src/validation/model.test.ts | 11 +- web/src/validation/model.ts | 5 +- 25 files changed, 772 insertions(+), 51 deletions(-) create mode 100644 core/common/config/retry_budget_test.go create mode 100644 core/controller/relay-budget_test.go create mode 100644 core/model/retry_budget_test.go diff --git a/config.example.yaml b/config.example.yaml index 9c20bf0ad..2aec8d8eb 100644 --- a/config.example.yaml +++ b/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 diff --git a/config.md b/config.md index d62b6aebb..b1ca5f5d1 100644 --- a/config.md +++ b/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. diff --git a/core/common/config/config.go b/core/common/config/config.go index c3bead458..a37bde441 100644 --- a/core/common/config/config.go +++ b/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() } diff --git a/core/common/config/env.go b/core/common/config/env.go index d2ecf84a2..c52517c1b 100644 --- a/core/common/config/env.go +++ b/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) diff --git a/core/common/config/retry_budget_test.go b/core/common/config/retry_budget_test.go new file mode 100644 index 000000000..0110fbb8c --- /dev/null +++ b/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()) + }) + } +} diff --git a/core/controller/group.go b/core/controller/group.go index 4716ea7d2..ad0eb5bcb 100644 --- a/core/controller/group.go +++ b/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, diff --git a/core/controller/relay-budget_test.go b/core/controller/relay-budget_test.go new file mode 100644 index 000000000..9a062bef0 --- /dev/null +++ b/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) +} diff --git a/core/controller/relay-controller.go b/core/controller/relay-controller.go index 9d6461ece..c385ba588 100644 --- a/core/controller/relay-controller.go +++ b/core/controller/relay-controller.go @@ -339,14 +339,16 @@ func relay(c *gin.Context, mode mode.Mode, relayController RelayController) { } // First attempt + firstAttemptAt := time.Now() + retryTimes, retryDeadline := getRetryLimits( + mc, + config.GetRetryTimes(), + config.GetRetryBudget(), + firstAttemptAt, + ) result, retry := RelayHelper(c, meta, relayController.Handler) - retryTimes := int(config.GetRetryTimes()) - if mc.RetryTimes > 0 { - retryTimes = int(mc.RetryTimes) - } - - if handleRelayResult(c, result.Error, retry, retryTimes) { + if handleRelayResult(c, result.Error, retry, retryTimes, retryDeadline) { recordResult( c, meta, @@ -369,6 +371,7 @@ func relay(c *gin.Context, mode mode.Mode, relayController RelayController) { price, time.Now(), ) + retryState.retryDeadline = retryDeadline // Retry loop retryLoop(c, mode, retryState, relayController.Handler) @@ -549,6 +552,7 @@ type retryState struct { channelSelectionState retryTimes int + retryDeadline time.Time designatedChannel *model.Channel preferChannelIDs []int ignoreChannelIDs map[int64]struct{} @@ -574,11 +578,32 @@ const ( relayRetryMaxJitter = time.Second ) +func getRetryLimits( + mc model.ModelConfig, + defaultRetryTimes, defaultRetryBudget int64, + startedAt time.Time, +) (int, time.Time) { + retryTimes, budget := mc.RetryLimits(defaultRetryTimes, defaultRetryBudget) + + var deadline time.Time + if budget > 0 { + deadline = startedAt.Add(budget) + } + + return int(retryTimes), deadline +} + +func (s *retryState) canRetry(attempts int, now time.Time) bool { + return (s.retryTimes < 0 || attempts < s.retryTimes) && + (s.retryDeadline.IsZero() || now.Before(s.retryDeadline)) +} + func handleRelayResult( c *gin.Context, bizErr adaptor.Error, retry bool, retryTimes int, + retryDeadline time.Time, ) (done bool) { if bizErr == nil { return true @@ -586,6 +611,7 @@ func handleRelayResult( if !retry || retryTimes == 0 || + (!retryDeadline.IsZero() && !time.Now().Before(retryDeadline)) || c.Request.Context().Err() != nil { ErrorWithRequestID(c, bizErr) return true @@ -693,34 +719,38 @@ func (s *retryState) remainingRelayDelay( func retryLoop(c *gin.Context, mode mode.Mode, state *retryState, relayController RelayHandler) { log := common.GetLogger(c) - // retryTimes can grow when permission failures add more eligible-channel attempts + // The budget limits scheduling and backoff, while in-flight requests keep their own timeout. + ctx := c.Request.Context() + if !state.retryDeadline.IsZero() { + var cancel context.CancelFunc + + ctx, cancel = context.WithDeadline(ctx, state.retryDeadline) + defer cancel() + } + i := 0 - for { - newChannel, err := getRetryChannel(c.Request.Context(), state) + for state.canRetry(i, time.Now()) && ctx.Err() == nil { + newChannel, err := getRetryChannel(ctx, state) if err == nil { err = prepareRetry(c) } + if err == nil { + err = relayDelay(ctx, state, newChannel.ID) + } + if err != nil { - if !errors.Is(err, ErrChannelsExhausted) { + if !errors.Is(err, ErrChannelsExhausted) && ctx.Err() == nil { log.Errorf("prepare retry failed: %+v", err) } - // when the last request has not recorded the result, record the result - if state.meta != nil && state.result != nil { - recordResult( - c, - state.meta, - state.price, - state.result, - i, - true, - middleware.GetRequestMetadata(c), - ) - } break } + + if !state.canRetry(i, time.Now()) || ctx.Err() != nil { + break + } // when the last request has not recorded the result, record the result if state.meta != nil && state.result != nil { recordResult( @@ -732,21 +762,17 @@ func retryLoop(c *gin.Context, mode mode.Mode, state *retryState, relayControlle false, middleware.GetRequestMetadata(c), ) - state.meta = nil - state.result = nil } log.Data["retry"] = strconv.Itoa(i + 1) - log.Warnf("using channel %s (type: %d, id: %d) to retry (remain times %d)", + log.Warnf("using channel %s (type: %d, id: %d) to retry (attempt %d)", newChannel.Name, newChannel.Type, newChannel.ID, - state.retryTimes-i, + i+1, ) - relayDelay(state, newChannel.ID) - state.meta = NewMetaByContext( c, newChannel, @@ -759,6 +785,8 @@ func retryLoop(c *gin.Context, mode mode.Mode, state *retryState, relayControlle var retry bool state.result, retry = RelayHelper(c, state.meta, relayController) + i++ + if state.result.Error != nil && shouldBackoffStatus(state.result.Error.StatusCode()) { state.recordChannelFailure(newChannel.ID, time.Now()) } @@ -770,23 +798,21 @@ func retryLoop(c *gin.Context, mode mode.Mode, state *retryState, relayControlle state.failedChannelIDs[int64(newChannel.ID)] = struct{}{} } - if done || i == state.retryTimes-1 { - recordResult( - c, - state.meta, - state.price, - state.result, - i+1, - true, - middleware.GetRequestMetadata(c), - ) - + if done { break } - - i++ } + recordResult( + c, + state.meta, + state.price, + state.result, + i, + true, + middleware.GetRequestMetadata(c), + ) + if state.result.Error != nil { ErrorWithRequestID(c, state.result.Error) } @@ -829,7 +855,6 @@ func handleRetryResult( } state.ignoreChannelIDs[int64(newChannel.ID)] = struct{}{} - state.retryTimes++ } return false @@ -840,15 +865,23 @@ func shouldBackoffStatus(statusCode int) bool { statusCode == http.StatusServiceUnavailable } -func relayDelay(state *retryState, channelID int) { +func relayDelay(ctx context.Context, state *retryState, channelID int) error { jitter := time.Duration(rand.Int64N(int64(relayRetryMaxJitter))) delay := state.remainingRelayDelay(channelID, time.Now(), jitter) if delay <= 0 { - return + return ctx.Err() } - time.Sleep(delay) + timer := time.NewTimer(delay) + defer timer.Stop() + + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return ctx.Err() + } } func RelayNotImplemented(c *gin.Context) { diff --git a/core/controller/relay-controller_test.go b/core/controller/relay-controller_test.go index a911f6433..dd87f428d 100644 --- a/core/controller/relay-controller_test.go +++ b/core/controller/relay-controller_test.go @@ -174,13 +174,14 @@ func TestHandleRelayResultDecidesRetryLifecycle(t *testing.T) { tt.bizErr, tt.retry, tt.retryTimes, + time.Time{}, )) }) } canceled, cancel := context.WithCancel(context.Background()) cancel() - assert.True(t, handleRelayResult(newContext(canceled), err, true, 2)) + assert.True(t, handleRelayResult(newContext(canceled), err, true, 2, time.Time{})) } func TestInitRetryStateRecordsInitialFailure(t *testing.T) { @@ -297,7 +298,7 @@ func TestHandleRetryResultUpdatesAutomaticRetryState(t *testing.T) { assert.Empty(t, state.ignoreChannelIDs) }) - t.Run("permission failure is hard filtered and extends retry budget", func(t *testing.T) { + t.Run("permission failure is hard filtered within retry limit", func(t *testing.T) { t.Parallel() state := newState(noPermissionError) @@ -305,7 +306,7 @@ func TestHandleRetryResultUpdatesAutomaticRetryState(t *testing.T) { done := handleRetryResult(newContext(), true, channel, state) assert.False(t, done) - assert.Equal(t, 3, state.retryTimes) + assert.Equal(t, 2, state.retryTimes) assert.Contains(t, state.ignoreChannelIDs, int64(channel.ID)) }) diff --git a/core/docs/docs.go b/core/docs/docs.go index d3c867fbc..9c27694d1 100644 --- a/core/docs/docs.go +++ b/core/docs/docs.go @@ -10080,6 +10080,12 @@ const docTemplate = `{ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "description": "Seconds; nil inherits the global budget, zero disables it.", + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, @@ -10988,6 +10994,9 @@ const docTemplate = `{ "override_response_body_storage_max_size": { "type": "boolean" }, + "override_retry_budget": { + "type": "boolean" + }, "override_retry_times": { "type": "boolean" }, @@ -11009,6 +11018,11 @@ const docTemplate = `{ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, @@ -11076,6 +11090,12 @@ const docTemplate = `{ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "description": "Seconds; nil inherits the global budget, zero disables it.", + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, @@ -13161,6 +13181,9 @@ const docTemplate = `{ "override_response_body_storage_max_size": { "type": "boolean" }, + "override_retry_budget": { + "type": "boolean" + }, "override_retry_times": { "type": "boolean" }, @@ -13182,6 +13205,11 @@ const docTemplate = `{ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, @@ -13682,6 +13710,12 @@ const docTemplate = `{ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "description": "Seconds; nil inherits the global budget, zero disables it.", + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, diff --git a/core/docs/swagger.json b/core/docs/swagger.json index dd264d041..90e3c2ce9 100644 --- a/core/docs/swagger.json +++ b/core/docs/swagger.json @@ -10071,6 +10071,12 @@ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "description": "Seconds; nil inherits the global budget, zero disables it.", + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, @@ -10979,6 +10985,9 @@ "override_response_body_storage_max_size": { "type": "boolean" }, + "override_retry_budget": { + "type": "boolean" + }, "override_retry_times": { "type": "boolean" }, @@ -11000,6 +11009,11 @@ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, @@ -11067,6 +11081,12 @@ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "description": "Seconds; nil inherits the global budget, zero disables it.", + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, @@ -13152,6 +13172,9 @@ "override_response_body_storage_max_size": { "type": "boolean" }, + "override_retry_budget": { + "type": "boolean" + }, "override_retry_times": { "type": "boolean" }, @@ -13173,6 +13196,11 @@ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, @@ -13673,6 +13701,12 @@ "response_body_storage_max_size": { "type": "integer" }, + "retry_budget": { + "description": "Seconds; nil inherits the global budget, zero disables it.", + "type": "integer", + "maximum": 180, + "minimum": 0 + }, "retry_times": { "type": "integer" }, diff --git a/core/docs/swagger.yaml b/core/docs/swagger.yaml index 216f38e28..e2b7e0d71 100644 --- a/core/docs/swagger.yaml +++ b/core/docs/swagger.yaml @@ -140,6 +140,11 @@ definitions: type: integer response_body_storage_max_size: type: integer + retry_budget: + description: Seconds; nil inherits the global budget, zero disables it. + maximum: 180 + minimum: 0 + type: integer retry_times: type: integer rpm: @@ -733,6 +738,8 @@ definitions: type: boolean override_response_body_storage_max_size: type: boolean + override_retry_budget: + type: boolean override_retry_times: type: boolean override_summary_claude_long_context: @@ -747,6 +754,10 @@ definitions: type: integer response_body_storage_max_size: type: integer + retry_budget: + maximum: 180 + minimum: 0 + type: integer retry_times: type: integer rpm: @@ -792,6 +803,11 @@ definitions: type: integer response_body_storage_max_size: type: integer + retry_budget: + description: Seconds; nil inherits the global budget, zero disables it. + maximum: 180 + minimum: 0 + type: integer retry_times: type: integer rpm: @@ -2242,6 +2258,8 @@ definitions: type: boolean override_response_body_storage_max_size: type: boolean + override_retry_budget: + type: boolean override_retry_times: type: boolean override_summary_claude_long_context: @@ -2256,6 +2274,10 @@ definitions: type: integer response_body_storage_max_size: type: integer + retry_budget: + maximum: 180 + minimum: 0 + type: integer retry_times: type: integer rpm: @@ -2590,6 +2612,11 @@ definitions: type: integer response_body_storage_max_size: type: integer + retry_budget: + description: Seconds; nil inherits the global budget, zero disables it. + maximum: 180 + minimum: 0 + type: integer retry_times: type: integer rpm: diff --git a/core/model/groupmodel.go b/core/model/groupmodel.go index 7c6693d1a..076151493 100644 --- a/core/model/groupmodel.go +++ b/core/model/groupmodel.go @@ -2,7 +2,9 @@ package model import ( "errors" + "fmt" + "github.com/labring/aiproxy/core/common/config" log "github.com/sirupsen/logrus" "gorm.io/gorm" ) @@ -12,6 +14,10 @@ const ( ) var groupModelConfigZeroValueUpdateFields = []string{ + "override_retry_times", + "retry_times", + "override_retry_budget", + "retry_budget", "override_max_image_generation_count", "max_image_generation_count", "override_max_video_generation_seconds", @@ -35,6 +41,9 @@ type GroupModelConfig 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 TimeoutConfig `json:"timeout_config,omitempty" gorm:"embedded"` @@ -68,6 +77,13 @@ func (g *GroupModelConfig) BeforeSave(_ *gorm.DB) (err error) { return errors.New("model is required") } + if g.RetryBudget < 0 || g.RetryBudget > config.MaxRetryBudgetSeconds { + return fmt.Errorf( + "retry_budget must be between 0 and %d seconds", + config.MaxRetryBudgetSeconds, + ) + } + if err := g.Price.ValidateConditionalPrices(); err != nil { return err } diff --git a/core/model/modelconfig.go b/core/model/modelconfig.go index d9e15640a..7d42b392b 100644 --- a/core/model/modelconfig.go +++ b/core/model/modelconfig.go @@ -11,6 +11,7 @@ import ( "github.com/bytedance/sonic" "github.com/go-viper/mapstructure/v2" "github.com/labring/aiproxy/core/common" + "github.com/labring/aiproxy/core/common/config" "github.com/labring/aiproxy/core/relay/mode" "gorm.io/gorm" ) @@ -38,6 +39,7 @@ type ModelConfig struct { TPM int64 ` json:"tpm,omitempty" yaml:"tpm,omitempty"` Price Price `gorm:"embedded" json:"price,omitempty" yaml:"price,omitempty"` RetryTimes int64 ` json:"retry_times,omitempty" yaml:"retry_times,omitempty"` + RetryBudget *int64 ` json:"retry_budget,omitempty" yaml:"retry_budget,omitempty" binding:"omitempty,gte=0,lte=180" minimum:"0" maximum:"180"` // Seconds; nil inherits the global budget, zero disables it. TimeoutConfig TimeoutConfig `gorm:"embedded" json:"timeout_config,omitempty" yaml:"timeout_config,omitempty"` ForceSaveDetail bool ` json:"force_save_detail,omitempty" yaml:"force_save_detail,omitempty"` MaxImageGenerationCount int ` json:"max_image_generation_count,omitempty" yaml:"max_image_generation_count,omitempty"` @@ -49,6 +51,7 @@ type ModelConfig struct { SummaryServiceTier bool ` json:"summary_service_tier,omitempty" yaml:"summary_service_tier,omitempty"` SummaryClaudeLongContext bool ` json:"summary_claude_long_context,omitempty" yaml:"summary_claude_long_context,omitempty"` DisableResolutionFuzzyMatch bool ` json:"disable_resolution_fuzzy_match,omitempty" yaml:"disable_resolution_fuzzy_match,omitempty"` + retryTimesOverridden bool `gorm:"-" json:"-" yaml:"-"` } func (c *ModelConfig) BeforeSave(_ *gorm.DB) (err error) { @@ -56,6 +59,14 @@ func (c *ModelConfig) BeforeSave(_ *gorm.DB) (err error) { return errors.New("model is required") } + if c.RetryBudget != nil && + (*c.RetryBudget < 0 || *c.RetryBudget > config.MaxRetryBudgetSeconds) { + return fmt.Errorf( + "retry_budget must be between 0 and %d seconds", + config.MaxRetryBudgetSeconds, + ) + } + if err := c.Price.ValidateConditionalPrices(); err != nil { return err } @@ -73,6 +84,29 @@ func NewDefaultModelConfig(model string) ModelConfig { } } +func (c *ModelConfig) RetryLimits(defaultTimes, defaultBudget int64) (int64, time.Duration) { + seconds := defaultBudget + if c.RetryBudget != nil { + seconds = *c.RetryBudget + } + + seconds = min(max(seconds, 0), config.MaxRetryBudgetSeconds) + + times := max(defaultTimes, 0) + if c.RetryTimes > 0 { + times = c.RetryTimes + } else if seconds > 0 && (c.RetryBudget != nil || c.retryTimesOverridden) { + // A local budget without a local count enables retries for the entire budget. + times = 0 + } + + if seconds > 0 && times == 0 { + times = -1 + } + + return times, time.Duration(seconds) * time.Second +} + func (c *ModelConfig) RequestTimeout() time.Duration { return timeoutSecond(c.TimeoutConfig.RequestTimeout) } @@ -143,6 +177,11 @@ func (c *ModelConfig) LoadFromGroupModelConfig(groupModelConfig GroupModelConfig if groupModelConfig.OverrideRetryTimes { newC.RetryTimes = groupModelConfig.RetryTimes + newC.retryTimesOverridden = true + } + + if groupModelConfig.OverrideRetryBudget { + newC.RetryBudget = new(groupModelConfig.RetryBudget) } if groupModelConfig.OverrideTimeoutConfig { diff --git a/core/model/option.go b/core/model/option.go index 8b5c4212a..b75f614b0 100644 --- a/core/model/option.go +++ b/core/model/option.go @@ -80,6 +80,7 @@ func initOptionMap() error { ) optionMap["DisableServe"] = strconv.FormatBool(config.GetDisableServe()) optionMap["RetryTimes"] = strconv.FormatInt(config.GetRetryTimes(), 10) + optionMap["RetryBudget"] = strconv.FormatInt(config.GetRetryBudget(), 10) defaultChannelModelsJSON, err := sonic.Marshal(config.GetDefaultChannelModels()) if err != nil { @@ -401,6 +402,20 @@ func updateOption(key, value string, isInit bool) (err error) { } config.SetRetryTimes(retryTimes) + case "RetryBudget": + retryBudget, err := strconv.ParseInt(value, 10, 64) + if err != nil { + return err + } + + if retryBudget < 0 || retryBudget > config.MaxRetryBudgetSeconds { + return fmt.Errorf( + "retry budget must be between 0 and %d seconds", + config.MaxRetryBudgetSeconds, + ) + } + + config.SetRetryBudget(retryBudget) case "GroupConsumeLevelRatio": var newGroupRpmRatio map[string]float64 diff --git a/core/model/retry_budget_test.go b/core/model/retry_budget_test.go new file mode 100644 index 000000000..3b3806c6e --- /dev/null +++ b/core/model/retry_budget_test.go @@ -0,0 +1,141 @@ +package model_test + +import ( + "math" + "path/filepath" + "testing" + "time" + + "github.com/bytedance/sonic" + "github.com/labring/aiproxy/core/model" + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +func TestModelRetryLimits(t *testing.T) { + t.Parallel() + + for _, tt := range []struct { + name string + model model.ModelConfig + group model.GroupModelConfig + globalTimes int64 + globalBudget int64 + wantTimes int64 + wantBudget time.Duration + }{ + {name: "no retry configuration"}, + {name: "global count", globalTimes: 3, wantTimes: 3}, + {name: "global budget only", globalBudget: 60, wantTimes: -1, wantBudget: time.Minute}, + {name: "both global limits", globalTimes: 3, globalBudget: 60, wantTimes: 3, wantBudget: time.Minute}, + {name: "model count", model: model.ModelConfig{RetryTimes: 2}, globalTimes: 3, wantTimes: 2}, + {name: "model budget only ignores global count", model: model.ModelConfig{RetryBudget: new(int64(60))}, globalTimes: 3, wantTimes: -1, wantBudget: time.Minute}, + {name: "both model limits", model: model.ModelConfig{RetryTimes: 2, RetryBudget: new(int64(60))}, globalTimes: 3, wantTimes: 2, wantBudget: time.Minute}, + {name: "model count with inherited budget", model: model.ModelConfig{RetryTimes: 2}, globalTimes: 3, globalBudget: 30, wantTimes: 2, wantBudget: 30 * time.Second}, + {name: "model disables inherited budget", model: model.ModelConfig{RetryBudget: new(int64(0))}, globalTimes: 3, globalBudget: 30, wantTimes: 3}, + {name: "group overrides budget", model: model.ModelConfig{RetryBudget: new(int64(60))}, group: model.GroupModelConfig{OverrideRetryBudget: true, RetryBudget: 20}, globalTimes: 3, wantTimes: -1, wantBudget: 20 * time.Second}, + {name: "group disables budget", model: model.ModelConfig{RetryTimes: 2, RetryBudget: new(int64(60))}, group: model.GroupModelConfig{OverrideRetryBudget: true}, globalBudget: 90, wantTimes: 2}, + {name: "group inherits model count", model: model.ModelConfig{RetryTimes: 2}, group: model.GroupModelConfig{OverrideRetryBudget: true, RetryBudget: 20}, wantTimes: 2, wantBudget: 20 * time.Second}, + {name: "group clears count for budget only", model: model.ModelConfig{RetryTimes: 2}, group: model.GroupModelConfig{OverrideRetryTimes: true}, globalTimes: 3, globalBudget: 30, wantTimes: -1, wantBudget: 30 * time.Second}, + {name: "group sets both limits", group: model.GroupModelConfig{OverrideRetryTimes: true, RetryTimes: 4, OverrideRetryBudget: true, RetryBudget: 40}, wantTimes: 4, wantBudget: 40 * time.Second}, + {name: "disabled group overrides are ignored", model: model.ModelConfig{RetryTimes: 2, RetryBudget: new(int64(60))}, group: model.GroupModelConfig{RetryTimes: 4, RetryBudget: 40}, wantTimes: 2, wantBudget: time.Minute}, + {name: "runtime cap avoids duration overflow", model: model.ModelConfig{RetryBudget: new(int64(math.MaxInt64))}, wantTimes: -1, wantBudget: 3 * time.Minute}, + {name: "global runtime cap", globalBudget: math.MaxInt64, wantTimes: -1, wantBudget: 3 * time.Minute}, + } { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + effective := tt.model.LoadFromGroupModelConfig(tt.group) + times, budget := effective.RetryLimits(tt.globalTimes, tt.globalBudget) + require.Equal(t, tt.wantTimes, times) + require.Equal(t, tt.wantBudget, budget) + }) + } +} + +func TestRetryBudgetSerializationAndValidation(t *testing.T) { + t.Parallel() + + for _, budget := range []int64{-1, 0, 1, 180, 181} { + mc := model.ModelConfig{Model: "test", RetryBudget: new(budget)} + + gc := model.GroupModelConfig{Model: "test", RetryBudget: budget} + if budget < 0 || budget > 180 { + require.Error(t, mc.BeforeSave(nil)) + require.Error(t, gc.BeforeSave(nil)) + } else { + require.NoError(t, mc.BeforeSave(nil)) + require.NoError(t, gc.BeforeSave(nil)) + } + } + + for _, budget := range []*int64{nil, new(int64(0)), new(int64(60))} { + original := model.ModelConfig{Model: "test", RetryBudget: budget} + jsonData, err := sonic.Marshal(original) + require.NoError(t, err) + + var decoded model.ModelConfig + require.NoError(t, sonic.Unmarshal(jsonData, &decoded)) + require.Equal(t, budget, decoded.RetryBudget) + + //nolint:musttag // ModelConfig also carries runtime state. + yamlData, err := yaml.Marshal(original) + require.NoError(t, err) + + decoded = model.ModelConfig{} + //nolint:musttag // Verify the actual model configuration round trip. + require.NoError(t, yaml.Unmarshal(yamlData, &decoded)) + require.Equal(t, budget, decoded.RetryBudget) + } +} + +func TestRetryBudgetPersistence(t *testing.T) { + db, err := model.OpenSQLite(filepath.Join(t.TempDir(), "retry.db")) + require.NoError(t, err) + + previousDB := model.DB + model.DB = db + t.Cleanup(func() { model.DB = previousDB }) + require.NoError(t, db.AutoMigrate(&model.ModelConfig{}, &model.GroupModelConfig{})) + + for _, budget := range []*int64{new(int64(60)), new(int64(0)), nil} { + mc := model.ModelConfig{Model: "test", RetryBudget: budget} + require.NoError(t, model.SaveModelConfig(mc)) + + loaded, err := model.GetModelConfig("test") + require.NoError(t, err) + require.Equal(t, budget, loaded.RetryBudget) + } + + for _, batch := range []bool{false, true} { + initial := model.GroupModelConfig{ + GroupID: "test", Model: "test", OverrideRetryTimes: true, RetryTimes: 3, + OverrideRetryBudget: true, RetryBudget: 60, + } + require.NoError(t, model.SaveGroupModelConfig(initial)) + + for _, override := range []bool{true, false} { + update := model.GroupModelConfig{ + GroupID: "test", + Model: "test", + OverrideRetryTimes: override, + OverrideRetryBudget: override, + } + if batch { + require.NoError( + t, + model.UpdateGroupModelConfigs("test", []model.GroupModelConfig{update}), + ) + } else { + require.NoError(t, model.UpdateGroupModelConfig(update)) + } + + loaded, err := model.GetGroupModelConfig("test", "test") + require.NoError(t, err) + require.Equal(t, override, loaded.OverrideRetryTimes) + require.Equal(t, override, loaded.OverrideRetryBudget) + require.Zero(t, loaded.RetryTimes) + require.Zero(t, loaded.RetryBudget) + } + } +} diff --git a/web/public/locales/en/translation.json b/web/public/locales/en/translation.json index a277a2b1b..26300a83f 100644 --- a/web/public/locales/en/translation.json +++ b/web/public/locales/en/translation.json @@ -457,6 +457,8 @@ "tpmPlaceholder": "Enter TPM limit", "retryTimes": "Retry Times", "retryTimesPlaceholder": "Number of retry attempts", + "retryBudget": "Retry Budget (seconds)", + "retryBudgetPlaceholder": "Global default; 0 disables the budget", "timeout": "Timeout (seconds)", "timeoutPlaceholder": "Request timeout in seconds", "streamTimeout": "Stream Timeout (seconds)", @@ -1292,6 +1294,7 @@ "overrideLimit": "Override Limit", "overrideLimitDesc": "Override the global RPM/TPM limits for this model", "overrideRetryTimes": "Override Retry", + "overrideRetryBudget": "Override Retry Budget", "overrideRetryTimesDesc": "Override the global retry times for this model", "overrideTimeoutConfig": "Override Timeout Config", "overrideTimeoutConfigDesc": "Override the request and stream timeout settings for this model", diff --git a/web/public/locales/zh/translation.json b/web/public/locales/zh/translation.json index 4a220d81b..485827e43 100644 --- a/web/public/locales/zh/translation.json +++ b/web/public/locales/zh/translation.json @@ -446,6 +446,8 @@ "tpmPlaceholder": "输入TPM限制", "retryTimes": "重试次数", "retryTimesPlaceholder": "重试尝试次数", + "retryBudget": "重试耗时预算（秒）", + "retryBudgetPlaceholder": "继承全局默认值；0 关闭预算", "timeout": "超时时间（秒）", "timeoutPlaceholder": "请求超时时间（秒）", "streamTimeout": "流式超时时间（秒）", @@ -1280,6 +1282,7 @@ "overrideLimit": "覆盖限制", "overrideLimitDesc": "覆盖此模型的全局 RPM/TPM 限制", "overrideRetryTimes": "覆盖重试", + "overrideRetryBudget": "覆盖重试耗时预算", "overrideRetryTimesDesc": "覆盖此模型的全局重试次数", "overrideTimeoutConfig": "覆盖超时配置", "overrideTimeoutConfigDesc": "覆盖此模型的请求超时和流式超时设置", diff --git a/web/src/feature/group/components/GroupModelConfigsTab.tsx b/web/src/feature/group/components/GroupModelConfigsTab.tsx index 33f0721dc..0bcd5db91 100644 --- a/web/src/feature/group/components/GroupModelConfigsTab.tsx +++ b/web/src/feature/group/components/GroupModelConfigsTab.tsx @@ -48,7 +48,7 @@ import { PriceFormFields } from '@/components/price/PriceFormFields' import { PriceDisplay } from '@/components/price/PriceDisplay' import { Combobox } from '@/components/ui/combobox' import { toast } from 'sonner' -import { priceSchema } from '@/validation/model' +import { priceSchema, retryBudgetSchema } from '@/validation/model' interface GroupModelConfigsTabProps { groupId: string @@ -70,6 +70,8 @@ const getDefaultConfig = (): Omit => ({ tpm: 0, override_retry_times: false, retry_times: 0, + override_retry_budget: false, + retry_budget: 0, override_timeout_config: false, timeout_config: {}, override_force_save_detail: false, @@ -163,6 +165,8 @@ export function GroupModelConfigsTab({ groupId }: GroupModelConfigsTabProps) { const [formTpm, setFormTpm] = useState(0) const [formOverrideRetryTimes, setFormOverrideRetryTimes] = useState(false) const [formRetryTimes, setFormRetryTimes] = useState(0) + const [formOverrideRetryBudget, setFormOverrideRetryBudget] = useState(false) + const [formRetryBudget, setFormRetryBudget] = useState(0) const [formOverrideTimeoutConfig, setFormOverrideTimeoutConfig] = useState(false) const [formTimeoutConfig, setFormTimeoutConfig] = useState({}) const [formOverrideForceSaveDetail, setFormOverrideForceSaveDetail] = useState(false) @@ -231,6 +235,8 @@ export function GroupModelConfigsTab({ groupId }: GroupModelConfigsTabProps) { setFormTpm(config.tpm) setFormOverrideRetryTimes(config.override_retry_times) setFormRetryTimes(config.retry_times) + setFormOverrideRetryBudget(config.override_retry_budget ?? false) + setFormRetryBudget(config.retry_budget ?? 0) setFormOverrideTimeoutConfig(config.override_timeout_config) setFormTimeoutConfig(config.timeout_config || {}) setFormOverrideForceSaveDetail(config.override_force_save_detail) @@ -259,6 +265,8 @@ export function GroupModelConfigsTab({ groupId }: GroupModelConfigsTabProps) { setFormTpm(defaults.tpm!) setFormOverrideRetryTimes(defaults.override_retry_times!) setFormRetryTimes(defaults.retry_times!) + setFormOverrideRetryBudget(defaults.override_retry_budget!) + setFormRetryBudget(defaults.retry_budget!) setFormOverrideTimeoutConfig(defaults.override_timeout_config!) setFormTimeoutConfig(defaults.timeout_config || {}) setFormOverrideForceSaveDetail(defaults.override_force_save_detail!) @@ -309,6 +317,8 @@ export function GroupModelConfigsTab({ groupId }: GroupModelConfigsTabProps) { setFormTpm(config.tpm) setFormOverrideRetryTimes(config.override_retry_times) setFormRetryTimes(config.retry_times) + setFormOverrideRetryBudget(config.override_retry_budget ?? false) + setFormRetryBudget(config.retry_budget ?? 0) setFormOverrideTimeoutConfig(config.override_timeout_config) setFormTimeoutConfig(config.timeout_config || {}) setFormOverrideForceSaveDetail(config.override_force_save_detail) @@ -349,6 +359,14 @@ export function GroupModelConfigsTab({ groupId }: GroupModelConfigsTabProps) { } } + if (formOverrideRetryBudget) { + const result = retryBudgetSchema.safeParse(formRetryBudget) + if (!result.success) { + toast.error(result.error.issues[0]?.message || t('error.validationDescription')) + return + } + } + const maxImageGenerationCountConfig = (() => { if (supportImageGenerationCountLimit) { return { @@ -419,6 +437,8 @@ export function GroupModelConfigsTab({ groupId }: GroupModelConfigsTabProps) { tpm: formTpm, override_retry_times: formOverrideRetryTimes, retry_times: formRetryTimes, + override_retry_budget: formOverrideRetryBudget, + retry_budget: formRetryBudget, override_timeout_config: formOverrideTimeoutConfig, ...(formOverrideTimeoutConfig && { timeout_config: formTimeoutConfig }), override_force_save_detail: formOverrideForceSaveDetail, @@ -606,6 +626,7 @@ export function GroupModelConfigsTab({ groupId }: GroupModelConfigsTabProps) { ]} /> )} +
+ + +
+ + {formOverrideRetryBudget && ( +
+ + setFormRetryBudget(Number(e.target.value))} + /> +
+ )} +
diff --git a/web/src/feature/model/components/BuiltinModelsDialog.tsx b/web/src/feature/model/components/BuiltinModelsDialog.tsx index cda13501b..3be79caec 100644 --- a/web/src/feature/model/components/BuiltinModelsDialog.tsx +++ b/web/src/feature/model/components/BuiltinModelsDialog.tsx @@ -70,6 +70,7 @@ const toModelSaveRequest = (model: ModelConfig): ModelSaveRequest => { rpm: model.rpm, tpm: model.tpm, retry_times: model.retry_times, + retry_budget: model.retry_budget, timeout_config: model.timeout_config, force_save_detail: model.force_save_detail, max_image_generation_count: model.max_image_generation_count, diff --git a/web/src/feature/model/components/ModelForm.tsx b/web/src/feature/model/components/ModelForm.tsx index 2b5cc0def..6b2c4159c 100644 --- a/web/src/feature/model/components/ModelForm.tsx +++ b/web/src/feature/model/components/ModelForm.tsx @@ -107,6 +107,7 @@ const MANAGED_MODEL_KEYS = new Set([ 'rpm', 'tpm', 'retry_times', + 'retry_budget', 'timeout_config', 'force_save_detail', 'max_image_generation_count', @@ -271,6 +272,7 @@ const buildChangeSummaries = ( addScalar('rpm', labels.rpm, next.rpm) addScalar('tpm', labels.tpm, next.tpm) addScalar('retry_times', labels.retryTimes, next.retry_times) + addScalar('retry_budget', labels.retryBudget, next.retry_budget) addScalar('force_save_detail', labels.forceSaveDetail, next.force_save_detail) addScalar('summary_service_tier', labels.summaryServiceTier, next.summary_service_tier) addScalar('summary_claude_long_context', labels.summaryClaudeLongContext, next.summary_claude_long_context) @@ -328,6 +330,7 @@ interface ModelFormProps { rpm?: number tpm?: number retry_times?: number + retry_budget?: number timeout_config?: ModelConfig['timeout_config'] timeout?: number stream_timeout?: number @@ -415,6 +418,7 @@ export function ModelForm({ rpm: defaultValues.rpm, tpm: defaultValues.tpm, retry_times: defaultValues.retry_times, + retry_budget: defaultValues.retry_budget, timeout: defaultValues.timeout, stream_timeout: defaultValues.stream_timeout ?? defaultValues.timeout_config?.stream_request_timeout, force_save_detail: defaultValues.force_save_detail ?? false, @@ -460,6 +464,7 @@ export function ModelForm({ rpm: t("model.dialog.rpm"), tpm: t("model.dialog.tpm"), retryTimes: t("model.dialog.retryTimes"), + retryBudget: t("model.dialog.retryBudget"), forceSaveDetail: t("model.dialog.forceSaveDetail"), summaryServiceTier: t("model.dialog.recordServiceTier"), summaryClaudeLongContext: t("model.dialog.recordClaudeLongContext"), @@ -916,6 +921,7 @@ export function ModelForm({ ...(data.rpm !== undefined && { rpm: Number(data.rpm) }), ...(data.tpm !== undefined && { tpm: Number(data.tpm) }), ...(data.retry_times !== undefined && { retry_times: Number(data.retry_times) }), + ...(data.retry_budget !== undefined && { retry_budget: Number(data.retry_budget) }), ...(mergedTimeoutConfig && { timeout_config: mergedTimeoutConfig }), ...(data.force_save_detail !== undefined && { force_save_detail: data.force_save_detail }), ...(supportImageGenerationCountLimit && data.max_image_generation_count !== undefined && { @@ -1187,6 +1193,29 @@ export function ModelForm({ )} /> + ( + + {t("model.dialog.retryBudget")} + + field.onChange(e.target.value ? Number(e.target.value) : undefined)} + /> + + + + )} + /> + {/* Timeout Field */} override_retry_times?: boolean retry_times?: number + override_retry_budget?: boolean + retry_budget?: number override_timeout_config?: boolean timeout_config?: TimeoutConfig override_force_save_detail?: boolean diff --git a/web/src/types/model.ts b/web/src/types/model.ts index 0ea4ae526..a1569b22f 100644 --- a/web/src/types/model.ts +++ b/web/src/types/model.ts @@ -80,6 +80,7 @@ export interface ModelConfig { rpm?: number tpm?: number retry_times?: number + retry_budget?: number timeout_config?: TimeoutConfig force_save_detail?: boolean max_image_generation_count?: number @@ -192,6 +193,7 @@ export interface ModelCreateRequest { rpm?: number tpm?: number retry_times?: number + retry_budget?: number timeout_config?: TimeoutConfig force_save_detail?: boolean max_image_generation_count?: number diff --git a/web/src/validation/model.test.ts b/web/src/validation/model.test.ts index 91852ff23..c3c187711 100644 --- a/web/src/validation/model.test.ts +++ b/web/src/validation/model.test.ts @@ -1,6 +1,15 @@ import { describe, expect, it } from 'vitest' -import { priceSchema } from './model' +import { priceSchema, retryBudgetSchema } from './model' + +describe('retry budget validation', () => { + it.each([0, 1, 60, 180])('accepts %s seconds', value => { + expect(retryBudgetSchema.safeParse(value).success).toBe(true) + }) + it.each([-1, 181, 1.5, Infinity, NaN])('rejects %s seconds', value => { + expect(retryBudgetSchema.safeParse(value).success).toBe(false) + }) +}) describe('daily conditional pricing validation', () => { it('accepts multiple non-overlapping daily peak ranges', () => { diff --git a/web/src/validation/model.ts b/web/src/validation/model.ts index 31629e305..5fdb5c817 100644 --- a/web/src/validation/model.ts +++ b/web/src/validation/model.ts @@ -223,6 +223,8 @@ const modelConfigSchema = z.object({ support_voices: z.array(z.string()).optional(), }).optional() +export const retryBudgetSchema = z.number().int().min(0).max(180) + export const modelCreateSchema = z.object({ model: z.string().min(1, 'Model name is required'), config: modelConfigSchema, @@ -231,7 +233,8 @@ export const modelCreateSchema = z.object({ exclude_from_tests: z.boolean().optional(), rpm: z.number().nonnegative('RPM must be a non-negative number').optional(), tpm: z.number().nonnegative('TPM must be a non-negative number').optional(), - retry_times: z.number().nonnegative('Retry times must be a non-negative number').optional(), + retry_times: z.number().int().nonnegative('Retry times must be a non-negative number').optional(), + retry_budget: retryBudgetSchema.optional(), timeout: z.number().nonnegative('Timeout must be a non-negative number').optional(), stream_timeout: z.number().nonnegative('Stream timeout must be a non-negative number').optional(), force_save_detail: z.boolean().optional(),