From cfaba1dd6754d4238e1360247c198a64a313e96c Mon Sep 17 00:00:00 2001 From: Calcium-Ion Date: Sat, 1 Aug 2026 09:35:51 +0800 Subject: [PATCH 01/99] fix(billing): harden tiered retry group-switch billing (#6570) Follow-up to #6518 (issue #6480) addressing three review findings: - Document and lock in arrears semantics for the wallet Reserve top-up: when an auto-group retry lands on a more expensive group, the full reservation delta is deducted unconditionally (balance may go negative), mirroring settlement, so the logged pre-consumed quota always reconciles with the actual balance movement. Genuine DB errors still fail the attempt with update_data_error. Subscription funding keeps its insufficient-quota behavior: subscriptions enforce a hard used<=total cap and do not support arrears. - PriceData.FreeModel is cleared when a retry switches from a free group to a paid one, keeping it consistent with the billing session created at that point. - getChannel refreshes GroupRatioInfo only after channel selection succeeds, and the retry loop records the channel in use_channel before PrepareTieredBillingForSelectedGroup can fail. --- controller/relay.go | 7 +-- service/billing_session.go | 4 ++ service/tiered_settle.go | 12 +++- service/tiered_settle_test.go | 109 ++++++++++++++++++++++++++++++++++ 4 files changed, 127 insertions(+), 5 deletions(-) diff --git a/controller/relay.go b/controller/relay.go index 7e4270867463..8dccfe76dddd 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -199,12 +199,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { newAPIError = channelErr break } + addUsedChannel(c, channel.Id) if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil { newAPIError = billingErr break } - addUsedChannel(c, channel.Id) bodyStorage, bodyErr := common.GetBodyStorage(c) if bodyErr != nil { // Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path) @@ -312,9 +312,6 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service }, nil } channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam) - - info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info) - if err != nil { return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) } @@ -322,6 +319,8 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在（retry）", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) } + info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info) + newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName) if newAPIError != nil { return nil, newAPIError diff --git a/service/billing_session.go b/service/billing_session.go index 96afcf034e30..42b0ffbca8ec 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -232,6 +232,10 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.NewAPIErro func (s *BillingSession) reserveFunding(delta int) error { switch funding := s.funding.(type) { case *WalletFunding: + // 与结算补扣（SettleBilling 正差额 → WalletFunding.Settle）语义一致： + // 全额无条件扣减，余额不足的部分记为欠费（余额可为负），不中断请求， + // 保证日志记录的预扣额度与用户余额的实际变动始终对账一致。 + // DecreaseUserQuota 仅在数据库错误时失败。 if err := model.DecreaseUserQuota(funding.userId, delta, false); err != nil { return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry()) } diff --git a/service/tiered_settle.go b/service/tiered_settle.go index b31d7964f56c..1f3f58fe31c1 100644 --- a/service/tiered_settle.go +++ b/service/tiered_settle.go @@ -132,9 +132,19 @@ func PrepareTieredBillingForSelectedGroup(c *gin.Context, relayInfo *relaycommon types.ErrOptionWithSkipRetry(), ) } - if snap == nil || snap.GroupRatio == 0 { + if snap == nil { return nil } + if snap.GroupRatio == 0 { + // Paid-to-free keeps FreeModel as-is: FreeModel means "pre-consume was + // skipped", which is not true once a session exists, and settlement + // already yields 0 for a zero group ratio. + return nil + } + + // The selected group is paid; clear a FreeModel flag frozen when the + // initial group was free so downstream state stays consistent. + relayInfo.PriceData.FreeModel = false if relayInfo.Billing == nil { return PreConsumeBilling(c, snap.EstimatedQuotaAfterGroup, relayInfo) diff --git a/service/tiered_settle_test.go b/service/tiered_settle_test.go index fe3a30a08bf1..29a3bfb4fdab 100644 --- a/service/tiered_settle_test.go +++ b/service/tiered_settle_test.go @@ -389,6 +389,7 @@ func TestPrepareTieredBillingForSelectedGroupStartsBillingAfterFreeGroup(t *test QuotaPerUnit: testQuotaPerUnit, }, PriceData: types.PriceData{ + FreeModel: true, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.20}, }, } @@ -396,6 +397,7 @@ func TestPrepareTieredBillingForSelectedGroupStartsBillingAfterFreeGroup(t *test require.Nil(t, PrepareTieredBillingForSelectedGroup(ctx, relayInfo)) require.NotNil(t, relayInfo.Billing) + assert.False(t, relayInfo.PriceData.FreeModel, "FreeModel must be cleared after switching to a paid group") assert.Equal(t, 100_000, relayInfo.FinalPreConsumedQuota) assert.Equal(t, 0.20, relayInfo.TieredBillingSnapshot.GroupRatio) assert.Equal(t, 100_000, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) @@ -405,6 +407,113 @@ func TestPrepareTieredBillingForSelectedGroupStartsBillingAfterFreeGroup(t *test assert.Equal(t, 400_000, userQuota) } +func TestPrepareTieredBillingForSelectedGroupPaidToFreeKeepsFreeModelFalse(t *testing.T) { + const expr = `tier("base", p)` + billing := &recordingBillingSettler{preConsumedQuota: 50_000} + relayInfo := &relaycommon.RelayInfo{ + Billing: billing, + FinalPreConsumedQuota: 50_000, + TieredBillingSnapshot: &billingexpr.BillingSnapshot{ + BillingMode: "tiered_expr", + ExprString: expr, + ExprHash: billingexpr.ExprHashString(expr), + GroupRatio: 0.10, + EstimatedQuotaBeforeGroup: 500_000, + EstimatedQuotaAfterGroup: 50_000, + QuotaPerUnit: testQuotaPerUnit, + }, + PriceData: types.PriceData{ + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0}, + }, + } + + require.Nil(t, PrepareTieredBillingForSelectedGroup(nil, relayInfo)) + + // Pre-consume did happen under the paid group, so FreeModel stays false; + // settlement already yields 0 for GroupRatio == 0 and the session refunds. + assert.False(t, relayInfo.PriceData.FreeModel) + assert.Empty(t, billing.reserveTargets) + assert.Equal(t, 50_000, relayInfo.FinalPreConsumedQuota) +} + +func TestPrepareTieredBillingForSelectedGroupTopUpArrearsAllowsNegativeBalance(t *testing.T) { + truncate(t) + + const userID = 701 + // Balance covers the initial 50k pre-consume (already deducted before this + // test's seed) but not the 50k top-up to the more expensive retry group. + // The top-up must NOT abort the request: the full delta is deducted, the + // uncovered 30k becomes arrears (negative balance), mirroring how + // settlement charges a positive delta unconditionally. + seedUser(t, userID, 20_000) + + relayInfo := &relaycommon.RelayInfo{ + UserId: userID, + IsPlayground: true, + FinalPreConsumedQuota: 50_000, + TieredBillingSnapshot: &billingexpr.BillingSnapshot{ + BillingMode: "tiered_expr", + ExprString: `tier("base", p)`, + ExprHash: billingexpr.ExprHashString(`tier("base", p)`), + GroupRatio: 0.10, + EstimatedQuotaBeforeGroup: 500_000, + EstimatedQuotaAfterGroup: 50_000, + QuotaPerUnit: testQuotaPerUnit, + }, + PriceData: types.PriceData{ + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.20}, + }, + } + session := &BillingSession{ + relayInfo: relayInfo, + funding: &WalletFunding{userId: userID, consumed: 50_000}, + preConsumedQuota: 50_000, + } + relayInfo.Billing = session + + require.Nil(t, PrepareTieredBillingForSelectedGroup(nil, relayInfo)) + + // Full reservation recorded; wallet charged the full delta into arrears. + assert.Equal(t, 100_000, session.GetPreConsumedQuota()) + assert.Equal(t, 100_000, relayInfo.FinalPreConsumedQuota) + assert.Equal(t, 100_000, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) + userQuota, err := model.GetUserQuota(userID, false) + require.NoError(t, err) + assert.Equal(t, -30_000, userQuota) + + // Settlement still reconciles against the full reservation: actual 80k + // refunds the 20k over-reserve, landing at seed - (actual - initial) = -10k. + require.NoError(t, session.Settle(80_000)) + userQuota, err = model.GetUserQuota(userID, false) + require.NoError(t, err) + assert.Equal(t, -10_000, userQuota) +} + +func TestBillingSessionReserveWalletTopUpDecrementsBalance(t *testing.T) { + truncate(t) + + const userID = 702 + seedUser(t, userID, 500_000) + + relayInfo := &relaycommon.RelayInfo{ + UserId: userID, + IsPlayground: true, + } + session := &BillingSession{ + relayInfo: relayInfo, + funding: &WalletFunding{userId: userID, consumed: 50_000}, + preConsumedQuota: 50_000, + } + + require.NoError(t, session.Reserve(100_000)) + + assert.Equal(t, 100_000, session.GetPreConsumedQuota()) + assert.Equal(t, 100_000, relayInfo.FinalPreConsumedQuota) + userQuota, err := model.GetUserQuota(userID, false) + require.NoError(t, err) + assert.Equal(t, 450_000, userQuota) +} + func TestTryTieredSettleUsesFinalGroupAfterRetry(t *testing.T) { const expr = `tier("base", p)` tests := []struct { From bd585d78efd418aaf7baa7e34fa48c5536581868 Mon Sep 17 00:00:00 2001 From: Calcium-Ion Date: Sat, 1 Aug 2026 22:39:54 +0800 Subject: [PATCH 02/99] fix(aws): cancel Bedrock requests on client disconnect (#6589) * fix(aws): cancel Bedrock requests on client disconnect * fix(billing): log effective usage billing path --- relay/channel/aws/relay-aws.go | 78 ++++-- relay/channel/aws/relay_aws_test.go | 403 ++++++++++++++++++++++++++++ service/billing_usage.go | 26 +- service/text_quota_test.go | 15 +- 4 files changed, 479 insertions(+), 43 deletions(-) diff --git a/relay/channel/aws/relay-aws.go b/relay/channel/aws/relay-aws.go index c502364c6adf..c4751b5af855 100644 --- a/relay/channel/aws/relay-aws.go +++ b/relay/channel/aws/relay-aws.go @@ -40,11 +40,24 @@ func getAwsErrorStatusCode(err error) int { return http.StatusInternalServerError } -func newAwsInvokeContext() (context.Context, context.CancelFunc) { +func newAwsInvokeContext(parent context.Context) (context.Context, context.CancelFunc) { if common.RelayTimeout <= 0 { - return context.Background(), func() {} + return context.WithCancel(parent) } - return context.WithTimeout(context.Background(), time.Duration(common.RelayTimeout)*time.Second) + return context.WithTimeout(parent, time.Duration(common.RelayTimeout)*time.Second) +} + +func newAwsInvokeError(requestContext context.Context, err error, operation string) *types.NewAPIError { + options := make([]types.NewAPIErrorOptions, 0, 1) + if requestContext.Err() != nil { + options = append(options, types.ErrOptionWithSkipRetry()) + } + return types.NewOpenAIError( + errors.Wrap(err, operation), + types.ErrorCodeAwsInvokeError, + getAwsErrorStatusCode(err), + options..., + ) } func newAwsClient(c *gin.Context, info *relaycommon.RelayInfo) (*bedrockruntime.Client, error) { @@ -215,13 +228,13 @@ func getAwsModelID(requestModel string) string { func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) { - ctx, cancel := newAwsInvokeContext() + requestContext := c.Request.Context() + ctx, cancel := newAwsInvokeContext(requestContext) defer cancel() awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput)) if err != nil { - statusCode := getAwsErrorStatusCode(err) - return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil + return newAwsInvokeError(requestContext, err, "InvokeModel"), nil } claudeInfo := &claude.ClaudeResponseInfo{ @@ -245,13 +258,13 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types } func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) { - ctx, cancel := newAwsInvokeContext() + requestContext := c.Request.Context() + ctx, cancel := newAwsInvokeContext(requestContext) defer cancel() awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput)) if err != nil { - statusCode := getAwsErrorStatusCode(err) - return types.NewOpenAIError(errors.Wrap(err, "InvokeModelWithResponseStream"), types.ErrorCodeAwsInvokeError, statusCode), nil + return newAwsInvokeError(requestContext, err, "InvokeModelWithResponseStream"), nil } stream := awsResp.GetStream() defer stream.Close() @@ -264,23 +277,38 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) ( Usage: &dto.Usage{}, } - for event := range stream.Events() { - switch v := event.(type) { - case *bedrockruntimeTypes.ResponseStreamMemberChunk: - info.SetFirstResponseTime() - respErr := claude.HandleStreamResponseData(c, info, claudeInfo, string(v.Value.Bytes)) - if respErr != nil { - return respErr, nil + events := stream.Events() +streamLoop: + for { + select { + case <-ctx.Done(): + break streamLoop + case event, ok := <-events: + if !ok { + break streamLoop + } + if ctx.Err() != nil { + break streamLoop + } + + switch v := event.(type) { + case *bedrockruntimeTypes.ResponseStreamMemberChunk: + info.SetFirstResponseTime() + respErr := claude.HandleStreamResponseData(c, info, claudeInfo, string(v.Value.Bytes)) + if respErr != nil { + return respErr, nil + } + case *bedrockruntimeTypes.UnknownUnionMember: + fmt.Println("unknown tag:", v.Tag) + return types.NewError(errors.New("unknown response type"), types.ErrorCodeInvalidRequest), nil + default: + fmt.Println("union is nil or unknown type") + return types.NewError(errors.New("nil or unknown response type"), types.ErrorCodeInvalidRequest), nil } - case *bedrockruntimeTypes.UnknownUnionMember: - fmt.Println("unknown tag:", v.Tag) - return types.NewError(errors.New("unknown response type"), types.ErrorCodeInvalidRequest), nil - default: - fmt.Println("union is nil or unknown type") - return types.NewError(errors.New("nil or unknown response type"), types.ErrorCodeInvalidRequest), nil } } + _ = stream.Close() claude.HandleStreamFinalResponse(c, info, claudeInfo) return nil, claudeInfo.Usage } @@ -288,13 +316,13 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) ( // Nova模型处理函数 func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) { - ctx, cancel := newAwsInvokeContext() + requestContext := c.Request.Context() + ctx, cancel := newAwsInvokeContext(requestContext) defer cancel() awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput)) if err != nil { - statusCode := getAwsErrorStatusCode(err) - return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil + return newAwsInvokeError(requestContext, err, "InvokeModel"), nil } // 解析Nova响应 diff --git a/relay/channel/aws/relay_aws_test.go b/relay/channel/aws/relay_aws_test.go index 92745ff40929..22d8373873ed 100644 --- a/relay/channel/aws/relay_aws_test.go +++ b/relay/channel/aws/relay_aws_test.go @@ -2,17 +2,145 @@ package aws import ( "bytes" + "context" + "errors" + "io" "net/http" "net/http/httptest" + "sync" "testing" + "time" "github.com/QuantumNous/new-api/common" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relaykit/dto" + relaytypes "github.com/QuantumNous/new-api/relaykit/types" + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream" + "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream/eventstreamapi" + "github.com/aws/aws-sdk-go-v2/credentials" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +const awsTestModel = "anthropic.claude-3-5-sonnet-20240620-v1:0" + +type awsHTTPClientFunc func(*http.Request) (*http.Response, error) + +func (f awsHTTPClientFunc) Do(request *http.Request) (*http.Response, error) { + return f(request) +} + +type awsNotifyingResponseWriter struct { + *httptest.ResponseRecorder + notifyOn []byte + notified chan int + once sync.Once +} + +func newAwsNotifyingResponseWriter(notifyOn string) *awsNotifyingResponseWriter { + return &awsNotifyingResponseWriter{ + ResponseRecorder: httptest.NewRecorder(), + notifyOn: []byte(notifyOn), + notified: make(chan int, 1), + } +} + +func (w *awsNotifyingResponseWriter) Write(data []byte) (int, error) { + return w.ResponseRecorder.Write(data) +} + +func (w *awsNotifyingResponseWriter) Flush() { + w.ResponseRecorder.Flush() + if bytes.Contains(w.Body.Bytes(), w.notifyOn) { + w.once.Do(func() { + w.notified <- w.Body.Len() + }) + } +} + +func newAwsTestClient(httpClient bedrockruntime.HTTPClient) *bedrockruntime.Client { + return bedrockruntime.New(bedrockruntime.Options{ + Region: "us-east-1", + BaseEndpoint: aws.String("https://bedrock.test"), + Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider( + "access-key", "secret-key", "", + )), + HTTPClient: httpClient, + Retryer: aws.NopRetryer{}, + }) +} + +func newAwsTestContext(writer http.ResponseWriter, requestContext context.Context) *gin.Context { + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(requestContext) + return c +} + +func newAwsTestRelayInfo() *relaycommon.RelayInfo { + return &relaycommon.RelayInfo{ + StartTime: time.Now(), + IsStream: true, + OriginModelName: awsTestModel, + RelayFormat: relaytypes.RelayFormatOpenAI, + ShouldIncludeUsage: true, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: awsTestModel, + }, + } +} + +func newAwsInvokeModelInput() *bedrockruntime.InvokeModelInput { + return &bedrockruntime.InvokeModelInput{ + ModelId: aws.String(awsTestModel), + Body: []byte(`{}`), + Accept: aws.String("application/json"), + ContentType: aws.String("application/json"), + } +} + +func newAwsStreamInput() *bedrockruntime.InvokeModelWithResponseStreamInput { + return &bedrockruntime.InvokeModelWithResponseStreamInput{ + ModelId: aws.String(awsTestModel), + Body: []byte(`{}`), + Accept: aws.String("application/json"), + ContentType: aws.String("application/json"), + } +} + +func writeAwsStreamEvent(writer io.Writer, data string) error { + payload, err := common.Marshal(struct { + Bytes []byte `json:"bytes"` + }{Bytes: []byte(data)}) + if err != nil { + return err + } + + return eventstream.NewEncoder().Encode(writer, eventstream.Message{ + Headers: eventstream.Headers{ + {Name: eventstreamapi.MessageTypeHeader, Value: eventstream.StringValue(eventstreamapi.EventMessageType)}, + {Name: eventstreamapi.EventTypeHeader, Value: eventstream.StringValue("chunk")}, + {Name: eventstreamapi.ContentTypeHeader, Value: eventstream.StringValue("application/json")}, + }, + Payload: payload, + }) +} + +func newAwsStreamResponse(request *http.Request, body io.ReadCloser) *http.Response { + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Header: http.Header{ + "Content-Type": []string{"application/vnd.amazon.eventstream"}, + "X-Amzn-Bedrock-Content-Type": []string{"application/json"}, + }, + Body: body, + Request: request, + } +} + func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testing.T) { t.Parallel() @@ -53,3 +181,278 @@ func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testi require.True(t, ok) require.Equal(t, []any{"computer-use-2025-01-24"}, values) } + +func TestNewAwsInvokeContextInheritsParent(t *testing.T) { + originalRelayTimeout := common.RelayTimeout + t.Cleanup(func() { + common.RelayTimeout = originalRelayTimeout + }) + + tests := []struct { + name string + relayTimeout int + wantDeadline bool + }{ + {name: "without relay timeout", relayTimeout: 0, wantDeadline: false}, + {name: "with relay timeout", relayTimeout: 30, wantDeadline: true}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + common.RelayTimeout = test.relayTimeout + parent, cancelParent := context.WithCancel(context.Background()) + invokeContext, cancelInvoke := newAwsInvokeContext(parent) + defer cancelInvoke() + + _, hasDeadline := invokeContext.Deadline() + assert.Equal(t, test.wantDeadline, hasDeadline) + + cancelParent() + require.ErrorIs(t, invokeContext.Err(), context.Canceled) + }) + } +} + +func TestNewAwsInvokeErrorSkipsRetryOnlyForClientCancellation(t *testing.T) { + canceledContext, cancel := context.WithCancel(context.Background()) + cancel() + + tests := []struct { + name string + requestContext context.Context + err error + wantSkipRetry bool + }{ + { + name: "client context canceled", + requestContext: canceledContext, + err: context.Canceled, + wantSkipRetry: true, + }, + { + name: "relay timeout with live client context", + requestContext: context.Background(), + err: context.DeadlineExceeded, + wantSkipRetry: false, + }, + { + name: "upstream error with live client context", + requestContext: context.Background(), + err: errors.New("upstream failed"), + wantSkipRetry: false, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + err := newAwsInvokeError(test.requestContext, test.err, "InvokeModel") + assert.Equal(t, test.wantSkipRetry, relaytypes.IsSkipRetryError(err)) + }) + } +} + +func TestAwsHandlersCancelSdkRequestAndSkipRetry(t *testing.T) { + originalRelayTimeout := common.RelayTimeout + common.RelayTimeout = 0 + t.Cleanup(func() { + common.RelayTimeout = originalRelayTimeout + }) + + tests := []struct { + name string + request any + handle func(*gin.Context, *relaycommon.RelayInfo, *Adaptor) (*relaytypes.NewAPIError, *dto.Usage) + }{ + {name: "non-stream", request: newAwsInvokeModelInput(), handle: awsHandler}, + {name: "stream", request: newAwsStreamInput(), handle: awsStreamHandler}, + {name: "nova", request: newAwsInvokeModelInput(), handle: handleNovaRequest}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + requestContext, cancelRequest := context.WithCancel(context.Background()) + t.Cleanup(cancelRequest) + + upstreamContexts := make(chan context.Context, 1) + client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) { + upstreamContexts <- request.Context() + <-request.Context().Done() + return nil, request.Context().Err() + })) + adaptor := &Adaptor{AwsClient: client, AwsReq: test.request} + c := newAwsTestContext(httptest.NewRecorder(), requestContext) + info := newAwsTestRelayInfo() + + type handlerResult struct { + err *relaytypes.NewAPIError + usage *dto.Usage + } + results := make(chan handlerResult, 1) + go func() { + err, usage := test.handle(c, info, adaptor) + results <- handlerResult{err: err, usage: usage} + }() + + var upstreamContext context.Context + select { + case upstreamContext = <-upstreamContexts: + case result := <-results: + t.Fatalf("handler returned before issuing AWS request: %v", result.err) + case <-time.After(5 * time.Second): + t.Fatal("AWS request did not start") + } + + cancelRequest() + + var result handlerResult + select { + case result = <-results: + case <-time.After(5 * time.Second): + t.Fatal("handler did not stop after client cancellation") + } + + require.ErrorIs(t, upstreamContext.Err(), context.Canceled) + require.NotNil(t, result.err) + assert.True(t, relaytypes.IsSkipRetryError(result.err)) + assert.Nil(t, result.usage) + }) + } +} + +func TestAwsStreamHandlerUsesFinalUpstreamUsage(t *testing.T) { + originalRelayTimeout := common.RelayTimeout + common.RelayTimeout = 0 + t.Cleanup(func() { + common.RelayTimeout = originalRelayTimeout + }) + + events := []string{ + `{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`, + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`, + `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`, + `{"type":"message_stop"}`, + } + client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) { + var body bytes.Buffer + for _, event := range events { + if err := writeAwsStreamEvent(&body, event); err != nil { + return nil, err + } + } + return newAwsStreamResponse(request, io.NopCloser(bytes.NewReader(body.Bytes()))), nil + })) + adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()} + recorder := httptest.NewRecorder() + c := newAwsTestContext(recorder, context.Background()) + + handlerErr, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor) + + require.Nil(t, handlerErr) + require.NotNil(t, usage) + require.NotNil(t, usage.BillingUsage) + require.NotNil(t, usage.BillingUsage.ClaudeUsage) + assert.Equal(t, 100, usage.BillingUsage.ClaudeUsage.InputTokens) + assert.Equal(t, 423, usage.BillingUsage.ClaudeUsage.OutputTokens) + assert.Contains(t, recorder.Body.String(), "[DONE]") +} + +func TestAwsStreamHandlerStopsAtClientCancellationAndKeepsPartialBillingUsage(t *testing.T) { + originalRelayTimeout := common.RelayTimeout + common.RelayTimeout = 0 + t.Cleanup(func() { + common.RelayTimeout = originalRelayTimeout + }) + + requestContext, cancelRequest := context.WithCancel(context.Background()) + t.Cleanup(cancelRequest) + releaseFinal := make(chan struct{}) + var releaseFinalOnce sync.Once + release := func() { + releaseFinalOnce.Do(func() { + close(releaseFinal) + }) + } + t.Cleanup(release) + + producerResults := make(chan error, 1) + upstreamContexts := make(chan context.Context, 1) + client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) { + upstreamContexts <- request.Context() + reader, writer := io.Pipe() + go func() { + defer writer.Close() + initialEvents := []string{ + `{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`, + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`, + } + for _, event := range initialEvents { + if err := writeAwsStreamEvent(writer, event); err != nil { + producerResults <- err + return + } + } + + <-releaseFinal + producerResults <- writeAwsStreamEvent(writer, `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`) + }() + return newAwsStreamResponse(request, reader), nil + })) + + responseWriter := newAwsNotifyingResponseWriter("partial") + c := newAwsTestContext(responseWriter, requestContext) + adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()} + + type handlerResult struct { + err *relaytypes.NewAPIError + usage *dto.Usage + } + results := make(chan handlerResult, 1) + go func() { + err, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor) + results <- handlerResult{err: err, usage: usage} + }() + + var upstreamContext context.Context + select { + case upstreamContext = <-upstreamContexts: + case <-time.After(5 * time.Second): + t.Fatal("AWS stream request did not start") + } + + var bodyLengthBeforeCancel int + select { + case bodyLengthBeforeCancel = <-responseWriter.notified: + case <-time.After(5 * time.Second): + t.Fatal("partial response was not written") + } + cancelRequest() + + var result handlerResult + select { + case result = <-results: + case <-time.After(5 * time.Second): + t.Fatal("stream handler did not stop after client cancellation") + } + + require.ErrorIs(t, upstreamContext.Err(), context.Canceled) + require.Nil(t, result.err) + require.NotNil(t, result.usage) + require.NotNil(t, result.usage.BillingUsage) + require.NotNil(t, result.usage.BillingUsage.ClaudeUsage) + assert.Equal(t, dto.BillingUsageSourceClaudeMessages, result.usage.BillingUsage.Source) + assert.Equal(t, dto.BillingUsageSemanticAnthropic, result.usage.BillingUsage.Semantic) + assert.Equal(t, 100, result.usage.BillingUsage.ClaudeUsage.InputTokens) + assert.Equal(t, 1, result.usage.BillingUsage.ClaudeUsage.OutputTokens) + assert.Equal(t, bodyLengthBeforeCancel, responseWriter.Body.Len()) + assert.NotContains(t, responseWriter.Body.String(), "[DONE]") + + release() + select { + case producerErr := <-producerResults: + require.Error(t, producerErr) + case <-time.After(5 * time.Second): + t.Fatal("upstream producer did not observe the closed stream") + } +} diff --git a/service/billing_usage.go b/service/billing_usage.go index 2766178d68c4..12656e693708 100644 --- a/service/billing_usage.go +++ b/service/billing_usage.go @@ -25,36 +25,32 @@ func effectiveBillingUsage(usage *dto.Usage) *dto.Usage { } func usageBillingPathForLog(isLocalCountTokens bool, usage *dto.Usage) string { - if isLocalCountTokens { - return usageBillingPathLocal - } - if usage == nil || usage.BillingUsage == nil { + effectiveUsage, ok := usageFromBillingUsage(usage) + if !ok { + if isLocalCountTokens { + return usageBillingPathLocal + } return usageBillingPathUpstream } - source := strings.TrimSpace(usage.BillingUsage.Source) - semantic := strings.TrimSpace(usage.BillingUsage.Semantic) - if strings.EqualFold(source, dto.BillingUsageSourceOAIChat) || - strings.EqualFold(source, dto.BillingUsageSourceOAIResponses) || - strings.EqualFold(semantic, dto.BillingUsageSemanticOpenAI) { + + switch effectiveUsage.UsageSemantic { + case dto.BillingUsageSemanticOpenAI: if usage.BillingUsage.Estimated { return usageBillingPathOpenAIEstimated } return usageBillingPathOpenAI - } - if strings.EqualFold(source, dto.BillingUsageSourceClaudeMessages) || - strings.EqualFold(semantic, dto.BillingUsageSemanticAnthropic) { + case dto.BillingUsageSemanticAnthropic: if usage.BillingUsage.Estimated { return usageBillingPathAnthropicEstimated } return usageBillingPathAnthropic - } - if strings.EqualFold(source, dto.BillingUsageSourceGeminiChat) || - strings.EqualFold(semantic, dto.BillingUsageSemanticGemini) { + case dto.BillingUsageSemanticGemini: if usage.BillingUsage.Estimated { return usageBillingPathGeminiEstimated } return usageBillingPathGemini } + return usageBillingPathUpstream } diff --git a/service/text_quota_test.go b/service/text_quota_test.go index 5a908af2d485..c9e958e7bf2b 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -284,9 +284,18 @@ func TestCalculateTextQuotaSummaryUsesOpenAIBillingUsageBeforeTopLevelUsage(t *t } func TestUsageBillingPathForLog(t *testing.T) { - require.Equal(t, usageBillingPathLocal, usageBillingPathForLog(true, &dto.Usage{ + require.Equal(t, usageBillingPathAnthropic, usageBillingPathForLog(true, &dto.Usage{ BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), })) + invalidBillingUsage := &dto.Usage{ + PromptTokens: 1, + BillingUsage: &dto.BillingUsage{ + Source: dto.BillingUsageSourceClaudeMessages, + Semantic: dto.BillingUsageSemanticAnthropic, + }, + } + require.Equal(t, usageBillingPathLocal, usageBillingPathForLog(true, invalidBillingUsage)) + require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, invalidBillingUsage)) require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, &dto.Usage{})) require.Equal(t, usageBillingPathOpenAI, usageBillingPathForLog(false, &dto.Usage{ BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{PromptTokens: 1}), @@ -297,7 +306,7 @@ func TestUsageBillingPathForLog(t *testing.T) { require.Equal(t, usageBillingPathGemini, usageBillingPathForLog(false, &dto.Usage{ BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{PromptTokenCount: 1}), })) - require.Equal(t, usageBillingPathGeminiEstimated, usageBillingPathForLog(false, &dto.Usage{ + require.Equal(t, usageBillingPathGeminiEstimated, usageBillingPathForLog(true, &dto.Usage{ BillingUsage: dto.NewEstimatedGeminiChatBillingUsage(&dto.Usage{PromptTokens: 1}), })) } @@ -306,7 +315,7 @@ func TestAppendUsageBillingPathForLogWritesAdminInfo(t *testing.T) { other := map[string]interface{}{ "admin_info": map[string]interface{}{}, } - appendUsageBillingPathForLog(other, false, &dto.Usage{ + appendUsageBillingPathForLog(other, true, &dto.Usage{ BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), }) From 0ab02020603d22e5613bc4cf46bfab06f8567769 Mon Sep 17 00:00:00 2001 From: Calcium-Ion Date: Sat, 1 Aug 2026 23:19:01 +0800 Subject: [PATCH 03/99] Feat/auto group (#6590) * feat(token): support custom auto group order * feat(keys): enhance auto group presentation * fix(keys): rework Auto flow border and compact inherited order The Auto group highlight previously tinted the whole control surface with a gradient and animated only a 1px top sweep, which read as a background color rather than a flowing border. Replace it with a border-only effect: an aria-hidden, pointer-events-none overlay whose conic gradient is masked down to a thin ring hugging the rounded perimeter, so the highlight travels around all four edges and corners every 3.2s. The interior stays neutral with a restrained static primary border and glow; prefers-reduced-motion hides the moving layer while keeping the static emphasis. The inherited global Auto order also rendered as spacious two-line rows with circular sequence markers, wasting drawer space. Render it as a compact wrapping strip of one-line chips (index, name, ratio badge) with descriptions kept accessible via title and sr-only text, scrolling only past a much smaller max height. Custom add/remove/reorder editing, empty-array inheritance semantics, and the submit payload are unchanged. * fix(keys): preserve Auto inheritance and unify effects * refactor(keys): temporarily disable AutoGroupBadge in api-key-group-cell --- constant/context_key.go | 1 + controller/model.go | 45 +- controller/model_list_test.go | 70 ++- controller/model_owned_by_test.go | 33 ++ controller/token.go | 132 ++++- controller/token_auto_groups_test.go | 234 ++++++++ controller/token_test.go | 40 ++ i18n/keys.go | 3 + i18n/locales/en.yaml | 3 + i18n/locales/zh-CN.yaml | 3 + i18n/locales/zh-TW.yaml | 3 + middleware/auth.go | 10 + middleware/distributor.go | 2 +- middleware/token_auto_groups_context_test.go | 48 ++ model/option.go | 6 + model/option_auto_group_test.go | 17 + model/token.go | 45 +- model/token_auto_groups_cache_test.go | 57 ++ router/api-router.go | 1 + service/channel_select.go | 5 +- service/channel_select_auto_groups_test.go | 129 +++++ service/group.go | 58 +- service/group_auto_groups_test.go | 72 +++ setting/auto_group.go | 33 ++ setting/auto_group_test.go | 29 + web/src/features/keys/api.ts | 9 + .../__tests__/api-key-group-cell.test.tsx | 236 ++++++++ .../__tests__/api-key-group-combobox.test.tsx | 294 ++++++++++ .../__tests__/api-keys-mutate-drawer.test.tsx | 371 ++++++++++++ .../auto-group-order-editor.test.tsx | 540 ++++++++++++++++++ .../keys/components/api-key-group-cell.tsx | 90 +++ .../components/api-key-group-combobox.tsx | 146 ++--- .../keys/components/api-keys-columns.tsx | 53 +- .../components/api-keys-mutate-drawer.tsx | 216 ++++++- .../components/auto-group-order-editor.tsx | 338 +++++++++++ .../keys/components/auto-group-visuals.tsx | 140 +++++ .../lib/__tests__/auto-group-form.test.ts | 189 ++++++ web/src/features/keys/lib/api-key-form.ts | 63 +- web/src/features/keys/types.ts | 7 + .../drawers/model-mutate-drawer.tsx | 1 + .../system-settings/billing/index.tsx | 1 + .../billing/section-registry.tsx | 1 + .../group-auto-limit-validation.test.ts | 40 ++ .../models/group-ratio-form.tsx | 56 ++ .../models/group-ratio-visual-editor.tsx | 12 +- .../features/system-settings/models/index.tsx | 1 + .../models/ratio-settings-card.tsx | 7 + web/src/features/system-settings/types.ts | 2 + .../system-settings/utils/numeric-field.ts | 5 + web/src/i18n/locales/en.json | 27 +- web/src/i18n/locales/fr.json | 27 +- web/src/i18n/locales/ja.json | 27 +- web/src/i18n/locales/ru.json | 27 +- web/src/i18n/locales/vi.json | 27 +- web/src/i18n/locales/zh-TW.json | 27 +- web/src/i18n/locales/zh.json | 27 +- web/src/styles/index.css | 46 ++ 57 files changed, 3922 insertions(+), 210 deletions(-) create mode 100644 controller/token_auto_groups_test.go create mode 100644 middleware/token_auto_groups_context_test.go create mode 100644 model/option_auto_group_test.go create mode 100644 model/token_auto_groups_cache_test.go create mode 100644 service/channel_select_auto_groups_test.go create mode 100644 service/group_auto_groups_test.go create mode 100644 setting/auto_group_test.go create mode 100644 web/src/features/keys/components/__tests__/api-key-group-cell.test.tsx create mode 100644 web/src/features/keys/components/__tests__/api-key-group-combobox.test.tsx create mode 100644 web/src/features/keys/components/__tests__/api-keys-mutate-drawer.test.tsx create mode 100644 web/src/features/keys/components/__tests__/auto-group-order-editor.test.tsx create mode 100644 web/src/features/keys/components/api-key-group-cell.tsx create mode 100644 web/src/features/keys/components/auto-group-order-editor.tsx create mode 100644 web/src/features/keys/components/auto-group-visuals.tsx create mode 100644 web/src/features/keys/lib/__tests__/auto-group-form.test.ts create mode 100644 web/src/features/system-settings/models/__tests__/group-auto-limit-validation.test.ts diff --git a/constant/context_key.go b/constant/context_key.go index b856bc3dda14..ccb8010f9476 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -19,6 +19,7 @@ const ( ContextKeyTokenModelLimitEnabled ContextKey = "token_model_limit_enabled" ContextKeyTokenModelLimit ContextKey = "token_model_limit" ContextKeyTokenCrossGroupRetry ContextKey = "token_cross_group_retry" + ContextKeyTokenAutoGroups ContextKey = "token_auto_groups" /* channel related keys */ ContextKeyChannelId ContextKey = "channel_id" diff --git a/controller/model.go b/controller/model.go index b32eebd7daac..1d759301bc7e 100644 --- a/controller/model.go +++ b/controller/model.go @@ -20,6 +20,7 @@ import ( "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/gin-gonic/gin" "github.com/samber/lo" ) @@ -190,7 +191,7 @@ func getModelListGroups(c *gin.Context) (modelListGroups, error) { return modelListGroups{ userGroup: userGroup, tokenGroup: tokenGroup, - ownerGroups: service.GetUserAutoGroup(userGroup), + ownerGroups: service.GetRequestAutoGroups(c, userGroup), }, nil } @@ -228,32 +229,28 @@ func ListModels(c *gin.Context, modelType int) { } ownerGroups := groups.ownerGroups modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) + var tokenModelLimit map[string]bool if modelLimitEnable { s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) - var tokenModelLimit map[string]bool if ok { - tokenModelLimit = s.(map[string]bool) - } else { + tokenModelLimit, _ = s.(map[string]bool) + } + if tokenModelLimit == nil { tokenModelLimit = map[string]bool{} } - for allowModel, _ := range tokenModelLimit { - if !acceptUnsetRatioModel { - if !helper.HasModelBillingConfig(allowModel) { - continue - } + } + models := service.GetGroupsEnabledModels(ownerGroups) + for _, modelName := range models { + if modelLimitEnable { + matchingName := ratio_setting.FormatMatchingModelName(modelName) + if !tokenModelLimit[modelName] && !tokenModelLimit[matchingName] { + continue } - userModelNames = append(userModelNames, allowModel) } - } else { - models := service.GetGroupsEnabledModels(ownerGroups) - for _, modelName := range models { - if !acceptUnsetRatioModel { - if !helper.HasModelBillingConfig(modelName) { - continue - } - } - userModelNames = append(userModelNames, modelName) + if !acceptUnsetRatioModel && !helper.HasModelBillingConfig(modelName) { + continue } + userModelNames = append(userModelNames, modelName) } ownerByModel := map[string]string{} @@ -276,11 +273,17 @@ func ListModels(c *gin.Context, modelType int) { Type: "model", } } + firstID := "" + lastID := "" + if len(useranthropicModels) > 0 { + firstID = useranthropicModels[0].ID + lastID = useranthropicModels[len(useranthropicModels)-1].ID + } c.JSON(200, gin.H{ "data": useranthropicModels, - "first_id": useranthropicModels[0].ID, + "first_id": firstID, "has_more": false, - "last_id": useranthropicModels[len(useranthropicModels)-1].ID, + "last_id": lastID, }) case constant.ChannelTypeGemini: userGeminiModels := make([]dto.GeminiModel, len(userOpenAiModels)) diff --git a/controller/model_list_test.go b/controller/model_list_test.go index b1fa9b956ce3..812207b8fd44 100644 --- a/controller/model_list_test.go +++ b/controller/model_list_test.go @@ -402,7 +402,13 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { "zz-token-tiered-visible-model": `tier("base", p * 1 + c * 2)`, "zz-token-tiered-empty-expr-model": "", }) - setupModelListControllerTestDB(t) + db := setupModelListControllerTestDB(t) + require.NoError(t, db.Create(&[]model.Ability{ + {Group: "default", Model: "zz-token-tiered-visible-model", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-token-tiered-empty-expr-model", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-token-tiered-missing-expr-model", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-token-unpriced-model", ChannelId: 1, Enabled: true}, + }).Error) recorder := httptest.NewRecorder() ctx, _ := gin.CreateTestContext(recorder) @@ -425,6 +431,68 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { require.NotContains(t, ids, "zz-token-unpriced-model") } +func TestListModelsTokenLimitUsesResolvedCustomAutoGroups(t *testing.T) { + withSelfUseModeEnabled(t) + originalMax := setting.GetMaxTokenAutoGroups() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups("5")) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) + + db := setupModelListControllerTestDB(t) + require.NoError(t, db.Create(&[]model.Ability{ + {Group: "vip", Model: "zz-vip-allowed", ChannelId: 1, Enabled: true}, + {Group: "vip", Model: "zz-vip-denied", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-default-outside-snapshot", ChannelId: 1, Enabled: true}, + }).Error) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto") + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + common.SetContextKey(ctx, constant.ContextKeyTokenModelLimitEnabled, true) + common.SetContextKey(ctx, constant.ContextKeyTokenModelLimit, map[string]bool{ + "zz-vip-allowed": true, + "zz-default-outside-snapshot": true, + "zz-not-enabled": true, + }) + + ListModels(ctx, constant.ChannelTypeOpenAI) + ids := decodeListModelsResponse(t, recorder) + require.Equal(t, map[string]struct{}{"zz-vip-allowed": {}}, ids) + + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) + emptyRecorder := httptest.NewRecorder() + emptyCtx, _ := gin.CreateTestContext(emptyRecorder) + emptyCtx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + common.SetContextKey(emptyCtx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(emptyCtx, constant.ContextKeyTokenGroup, "auto") + common.SetContextKey(emptyCtx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimitEnabled, true) + common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimit, map[string]bool{"zz-vip-allowed": true}) + + require.NotPanics(t, func() { + ListModels(emptyCtx, constant.ChannelTypeAnthropic) + }) + var anthropicResponse struct { + Data []dto.AnthropicModel `json:"data"` + FirstID string `json:"first_id"` + LastID string `json:"last_id"` + } + require.NoError(t, common.Unmarshal(emptyRecorder.Body.Bytes(), &anthropicResponse)) + require.Empty(t, anthropicResponse.Data) + require.Empty(t, anthropicResponse.FirstID) + require.Empty(t, anthropicResponse.LastID) +} + func TestCheckUpdatePasswordRequiresCurrentPassword(t *testing.T) { db := setupModelListControllerTestDB(t) hashedPassword, err := common.Password2Hash("CurrentPassword123") diff --git a/controller/model_owned_by_test.go b/controller/model_owned_by_test.go index bc2ef32f135c..da9bb1a0ff0e 100644 --- a/controller/model_owned_by_test.go +++ b/controller/model_owned_by_test.go @@ -1,11 +1,14 @@ package controller import ( + "fmt" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -83,3 +86,33 @@ func TestGetModelListGroupsUsesExplicitTokenGroup(t *testing.T) { require.Equal(t, "vip", groups.tokenGroup) require.Equal(t, []string{"vip"}, groups.ownerGroups) } + +func TestGetModelListGroupsUsesFilteredTokenAutoGroupsSnapshot(t *testing.T) { + originalMax := setting.GetMaxTokenAutoGroups() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups("1")) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) + + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto") + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip", "default"}) + + groups, err := getModelListGroups(ctx) + require.NoError(t, err) + require.Equal(t, []string{"vip"}, groups.ownerGroups) + + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) + groups, err = getModelListGroups(ctx) + require.NoError(t, err) + require.Empty(t, groups.ownerGroups) +} diff --git a/controller/token.go b/controller/token.go index 836e9b2952ac..c26d82e3dee1 100644 --- a/controller/token.go +++ b/controller/token.go @@ -7,30 +7,115 @@ import ( "strings" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/gin-gonic/gin" ) -func buildMaskedTokenResponse(token *model.Token) *model.Token { +type tokenAutoGroupsInput struct { + Set bool + Groups []string +} + +func (input *tokenAutoGroupsInput) UnmarshalJSON(data []byte) error { + input.Set = true + if strings.TrimSpace(string(data)) == "null" { + input.Groups = nil + return nil + } + return common.Unmarshal(data, &input.Groups) +} + +type tokenRequest struct { + model.Token + AutoGroups tokenAutoGroupsInput `json:"auto_groups"` +} + +type tokenResponse struct { + *model.Token + AutoGroups []string `json:"auto_groups"` +} + +func buildMaskedTokenResponse(token *model.Token) *tokenResponse { if token == nil { return nil } maskedToken := *token maskedToken.Key = token.GetMaskedKey() - return &maskedToken + autoGroups, err := token.GetAutoGroups() + if err != nil { + common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err)) + autoGroups = nil + } + if len(autoGroups) == 0 { + autoGroups = nil + } + return &tokenResponse{Token: &maskedToken, AutoGroups: autoGroups} } -func buildMaskedTokenResponses(tokens []*model.Token) []*model.Token { - maskedTokens := make([]*model.Token, 0, len(tokens)) +func buildMaskedTokenResponses(tokens []*model.Token) []*tokenResponse { + maskedTokens := make([]*tokenResponse, 0, len(tokens)) for _, token := range tokens { maskedTokens = append(maskedTokens, buildMaskedTokenResponse(token)) } return maskedTokens } +func getTokenRequestUserGroup(c *gin.Context) (string, error) { + if userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup); userGroup != "" { + return userGroup, nil + } + if userGroup := c.GetString("group"); userGroup != "" { + return userGroup, nil + } + return model.GetUserGroup(c.GetInt("id"), false) +} + +func setTokenAutoGroups(c *gin.Context, token *model.Token, groups []string) bool { + if len(groups) == 0 { + if err := token.SetAutoGroups(nil); err != nil { + common.ApiError(c, err) + return false + } + return true + } + + maxCount := setting.GetMaxTokenAutoGroups() + if len(groups) > maxCount { + common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsTooMany, map[string]any{"Max": maxCount}) + return false + } + + userGroup, err := getTokenRequestUserGroup(c) + if err != nil { + common.ApiError(c, err) + return false + } + seen := make(map[string]struct{}, len(groups)) + for _, group := range groups { + if _, ok := seen[group]; ok { + common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": group}) + return false + } + seen[group] = struct{}{} + if !service.IsUserSelectableGroup(userGroup, group) { + common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsInvalid, map[string]any{"Group": group}) + return false + } + } + + if err := token.SetAutoGroups(groups); err != nil { + common.ApiError(c, err) + return false + } + return true +} + func GetAllTokens(c *gin.Context) { userId := c.GetInt("id") pageInfo := common.GetPageQuery(c) @@ -77,6 +162,18 @@ func GetToken(c *gin.Context) { common.ApiSuccess(c, buildMaskedTokenResponse(token)) } +func GetTokenAutoGroups(c *gin.Context) { + userGroup, err := getTokenRequestUserGroup(c) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{ + "groups": service.GetUserAutoGroup(userGroup), + "max_count": setting.GetMaxTokenAutoGroups(), + }) +} + func GetTokenKey(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) userId := c.GetInt("id") @@ -165,12 +262,13 @@ func GetTokenUsage(c *gin.Context) { } func AddToken(c *gin.Context) { - token := model.Token{} - err := c.ShouldBindJSON(&token) + request := tokenRequest{} + err := c.ShouldBindJSON(&request) if err != nil { common.ApiError(c, err) return } + token := request.Token if len(token.Name) > 50 { common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong) return @@ -201,6 +299,14 @@ func AddToken(c *gin.Context) { }) return } + if token.Group == "auto" { + if !setTokenAutoGroups(c, &token, request.AutoGroups.Groups) { + return + } + } else { + token.CrossGroupRetry = false + _ = token.SetAutoGroups(nil) + } key, err := common.GenerateKey() if err != nil { common.ApiErrorI18n(c, i18n.MsgTokenGenerateFailed) @@ -221,6 +327,7 @@ func AddToken(c *gin.Context) { AllowIps: token.AllowIps, Group: token.Group, CrossGroupRetry: token.CrossGroupRetry, + AutoGroups: token.AutoGroups, } err = cleanToken.Insert() if err != nil { @@ -250,12 +357,13 @@ func DeleteToken(c *gin.Context) { func UpdateToken(c *gin.Context) { userId := c.GetInt("id") statusOnly := c.Query("status_only") - token := model.Token{} - err := c.ShouldBindJSON(&token) + request := tokenRequest{} + err := c.ShouldBindJSON(&request) if err != nil { common.ApiError(c, err) return } + token := request.Token if len(token.Name) > 50 { common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong) return @@ -299,6 +407,14 @@ func UpdateToken(c *gin.Context) { cleanToken.AllowIps = token.AllowIps cleanToken.Group = token.Group cleanToken.CrossGroupRetry = token.CrossGroupRetry + if token.Group != "auto" { + cleanToken.CrossGroupRetry = false + _ = cleanToken.SetAutoGroups(nil) + } else if request.AutoGroups.Set { + if !setTokenAutoGroups(c, cleanToken, request.AutoGroups.Groups) { + return + } + } } err = cleanToken.Update() if err != nil { diff --git a/controller/token_auto_groups_test.go b/controller/token_auto_groups_test.go new file mode 100644 index 000000000000..3da2969575c1 --- /dev/null +++ b/controller/token_auto_groups_test.go @@ -0,0 +1,234 @@ +package controller + +import ( + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func configureTokenAutoGroupsTest(t *testing.T, maxCount string, autoGroups string) { + t.Helper() + originalMax := setting.GetMaxTokenAutoGroups() + originalAutoGroups := setting.AutoGroups2JsonString() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups(maxCount)) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(autoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(stringInt(originalMax))) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) +} + +func stringInt(value int) string { + return fmt.Sprintf("%d", value) +} + +func setupTokenAutoGroupsControllerTest(t *testing.T) *model.User { + t.Helper() + db := setupTokenControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.User{})) + user := &model.User{ + Id: 101, + Username: "token-auto-user", + Password: "password", + Group: "default", + Status: common.UserStatusEnabled, + } + require.NoError(t, db.Create(user).Error) + return user +} + +func baseAutoTokenRequest(name string) map[string]any { + return map[string]any{ + "name": name, + "expired_time": -1, + "remain_quota": 0, + "unlimited_quota": true, + "group": "auto", + "cross_group_retry": true, + } +} + +func newTokenAutoGroupsAuthenticatedContext(t *testing.T, method string, target string, body any, userID int) (*gin.Context, *httptest.ResponseRecorder) { + t.Helper() + ctx, recorder := newAuthenticatedContext(t, method, target, body, userID) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + return ctx, recorder +} + +func TestAddTokenEmptyAutoGroupsInheritGlobalAuto(t *testing.T) { + tests := []struct { + name string + includeField bool + value any + }{ + {name: "omitted"}, + {name: "null", includeField: true, value: nil}, + {name: "empty array", includeField: true, value: []string{}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + configureTokenAutoGroupsTest(t, "5", `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + request := baseAutoTokenRequest("create-" + test.name) + if test.includeField { + request["auto_groups"] = test.value + } + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id) + AddToken(ctx) + + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + var token model.Token + require.NoError(t, model.DB.Where("name = ?", request["name"]).First(&token).Error) + assert.Empty(t, token.AutoGroups) + assert.True(t, token.CrossGroupRetry) + payload, err := common.Marshal(buildMaskedTokenResponse(&token)) + require.NoError(t, err) + var responseData map[string]any + require.NoError(t, common.Unmarshal(payload, &responseData)) + assert.Nil(t, responseData["auto_groups"]) + }) + } +} + +func TestAddTokenPersistsOrderedAutoGroupsSnapshot(t *testing.T) { + configureTokenAutoGroupsTest(t, "5", `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + request := baseAutoTokenRequest("ordered-snapshot") + request["auto_groups"] = []string{"vip", "default"} + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id) + AddToken(ctx) + require.True(t, decodeAPIResponse(t, recorder).Success) + + var token model.Token + require.NoError(t, model.DB.Where("name = ?", "ordered-snapshot").First(&token).Error) + assert.JSONEq(t, `["vip","default"]`, token.AutoGroups) + + getCtx, getRecorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodGet, "/api/token/"+stringInt(token.Id), nil, user.Id) + getCtx.Params = append(getCtx.Params, gin.Param{Key: "id", Value: stringInt(token.Id)}) + GetToken(getCtx) + getResponse := decodeAPIResponse(t, getRecorder) + require.True(t, getResponse.Success) + var data struct { + AutoGroups []string `json:"auto_groups"` + } + require.NoError(t, common.Unmarshal(getResponse.Data, &data)) + assert.Equal(t, []string{"vip", "default"}, data.AutoGroups) +} + +func TestUpdateTokenAutoGroupsTriStateAndNonAutoCleanup(t *testing.T) { + tests := []struct { + name string + includeField bool + value any + group string + expectedAutoGroups string + expectedRetry bool + }{ + {name: "omitted preserves", group: "auto", expectedAutoGroups: `["vip","default"]`, expectedRetry: true}, + {name: "null inherits", includeField: true, value: nil, group: "auto", expectedRetry: true}, + {name: "empty inherits", includeField: true, value: []string{}, group: "auto", expectedRetry: true}, + {name: "non auto clears and disables retry", includeField: true, value: []string{"vip"}, group: "default"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + configureTokenAutoGroupsTest(t, "5", `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + token := seedToken(t, model.DB, user.Id, "update-auto", "update-auto-key") + token.Group = "auto" + token.CrossGroupRetry = true + require.NoError(t, token.SetAutoGroups([]string{"vip", "default"})) + require.NoError(t, model.DB.Save(token).Error) + + request := baseAutoTokenRequest("updated-auto") + request["id"] = token.Id + request["status"] = common.TokenStatusEnabled + request["group"] = test.group + if test.includeField { + request["auto_groups"] = test.value + } + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPut, "/api/token/", request, user.Id) + UpdateToken(ctx) + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + + var updated model.Token + require.NoError(t, model.DB.First(&updated, token.Id).Error) + if test.expectedAutoGroups == "" { + assert.Empty(t, updated.AutoGroups) + } else { + assert.JSONEq(t, test.expectedAutoGroups, updated.AutoGroups) + } + assert.Equal(t, test.expectedRetry, updated.CrossGroupRetry) + }) + } +} + +func TestAddTokenRejectsInvalidAutoGroups(t *testing.T) { + tests := []struct { + name string + maxCount string + groups []string + }{ + {name: "over limit", maxCount: "1", groups: []string{"default", "vip"}}, + {name: "duplicate", maxCount: "5", groups: []string{"default", "default"}}, + {name: "auto pseudo group", maxCount: "5", groups: []string{"auto"}}, + {name: "unavailable", maxCount: "5", groups: []string{"missing"}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + configureTokenAutoGroupsTest(t, test.maxCount, `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + request := baseAutoTokenRequest("invalid-" + test.name) + request["auto_groups"] = test.groups + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id) + AddToken(ctx) + + response := decodeAPIResponse(t, recorder) + assert.False(t, response.Success) + var count int64 + require.NoError(t, model.DB.Model(&model.Token{}).Count(&count).Error) + assert.Zero(t, count) + }) + } +} + +func TestGetTokenAutoGroupsReturnsFullFilteredGlobalOrderAndLimit(t *testing.T) { + configureTokenAutoGroupsTest(t, "1", `["vip","missing","default"]`) + user := setupTokenAutoGroupsControllerTest(t) + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodGet, "/api/token/auto-groups", nil, user.Id) + GetTokenAutoGroups(ctx) + + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + var data struct { + Groups []string `json:"groups"` + MaxCount int `json:"max_count"` + } + require.NoError(t, common.Unmarshal(response.Data, &data)) + assert.Equal(t, []string{"vip", "default"}, data.Groups) + assert.Equal(t, 1, data.MaxCount) +} diff --git a/controller/token_test.go b/controller/token_test.go index 12b1cbdd84fb..9cca168ab100 100644 --- a/controller/token_test.go +++ b/controller/token_test.go @@ -273,6 +273,34 @@ func getTokenKeyColumnType(t *testing.T, db *gorm.DB, dialect string) string { } } +func getTokenAutoGroupsColumnType(t *testing.T, db *gorm.DB, dialect string) string { + t.Helper() + + switch dialect { + case "sqlite": + return getSQLiteColumnType(t, db, "tokens", "auto_groups") + case "mysql": + var columnType string + if err := db.Raw(`SELECT DATA_TYPE FROM information_schema.columns + WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`, + "tokens", "auto_groups").Scan(&columnType).Error; err != nil { + t.Fatalf("failed to inspect mysql token auto_groups column: %v", err) + } + return strings.ToLower(columnType) + case "postgres": + var dataType string + if err := db.Raw(`SELECT data_type FROM information_schema.columns + WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`, + "tokens", "auto_groups").Scan(&dataType).Error; err != nil { + t.Fatalf("failed to inspect postgres token auto_groups column: %v", err) + } + return strings.ToLower(dataType) + default: + t.Fatalf("unsupported dialect %q", dialect) + return "" + } +} + func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect string, managedTokensTable *bool) { t.Helper() @@ -314,6 +342,12 @@ func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect strin if got := getTokenKeyColumnType(t, db, dialect); got != "varchar(128)" { t.Fatalf("expected migrated key column type varchar(128), got %q", got) } + if !db.Migrator().HasColumn(&model.Token{}, "auto_groups") { + t.Fatal("expected migration to add auto_groups column") + } + if got := getTokenAutoGroupsColumnType(t, db, dialect); got != "text" { + t.Fatalf("expected migrated auto_groups column type text, got %q", got) + } var migratedToken model.Token if err := db.First(&migratedToken, "name = ?", "legacy-token").Error; err != nil { @@ -325,6 +359,9 @@ func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect strin if migratedToken.Name != "legacy-token" { t.Fatalf("expected migrated token name to be preserved, got %q", migratedToken.Name) } + if migratedToken.AutoGroups != "" { + t.Fatalf("expected legacy token to inherit global Auto groups, got %q", migratedToken.AutoGroups) + } inserted := model.Token{ UserId: 8, @@ -362,6 +399,9 @@ func TestTokenAutoMigrateUsesVarchar128KeyColumn(t *testing.T) { if got := getTokenKeyColumnType(t, db, "sqlite"); got != "varchar(128)" { t.Fatalf("expected key column type varchar(128), got %q", got) } + if got := getSQLiteColumnType(t, db, "tokens", "auto_groups"); got != "text" { + t.Fatalf("expected auto_groups column type text, got %q", got) + } } func TestTokenMigrationFromChar48ToVarchar128(t *testing.T) { diff --git a/i18n/keys.go b/i18n/keys.go index 8e9a4b5694d4..64a835e1a942 100644 --- a/i18n/keys.go +++ b/i18n/keys.go @@ -55,6 +55,9 @@ const ( MsgTokenExhausted = "token.exhausted" MsgTokenStatusUnavailable = "token.status_unavailable" MsgTokenDbError = "token.db_error" + MsgTokenAutoGroupsTooMany = "token.auto_groups_too_many" + MsgTokenAutoGroupsDuplicate = "token.auto_groups_duplicate" + MsgTokenAutoGroupsInvalid = "token.auto_groups_invalid" ) // Redemption related messages diff --git a/i18n/locales/en.yaml b/i18n/locales/en.yaml index 3f1fd03cb090..c533daecc32d 100644 --- a/i18n/locales/en.yaml +++ b/i18n/locales/en.yaml @@ -47,6 +47,9 @@ token.expired: "This token has expired" token.exhausted: "This token quota is exhausted TokenStatusExhausted[sk-{{.Prefix}}***{{.Suffix}}]" token.status_unavailable: "This token status is unavailable" token.db_error: "Invalid token, database query error, please contact administrator" +token.auto_groups_too_many: "A token can select at most {{.Max}} Auto groups" +token.auto_groups_duplicate: "Auto group {{.Group}} is duplicated" +token.auto_groups_invalid: "Auto group {{.Group}} is unavailable or unauthorized" # Redemption messages redemption.name_length: "Redemption code name length must be between 1-20" diff --git a/i18n/locales/zh-CN.yaml b/i18n/locales/zh-CN.yaml index fe982e59a0f3..a2f5275be9a8 100644 --- a/i18n/locales/zh-CN.yaml +++ b/i18n/locales/zh-CN.yaml @@ -48,6 +48,9 @@ token.expired: "该令牌已过期" token.exhausted: "该令牌额度已用尽 TokenStatusExhausted[sk-{{.Prefix}}***{{.Suffix}}]" token.status_unavailable: "该令牌状态不可用" token.db_error: "无效的令牌，数据库查询出错，请联系管理员" +token.auto_groups_too_many: "每个令牌最多可选择 {{.Max}} 个 Auto 分组" +token.auto_groups_duplicate: "Auto 分组 {{.Group}} 重复" +token.auto_groups_invalid: "Auto 分组 {{.Group}} 不可用或无权访问" # Redemption messages redemption.name_length: "兑换码名称长度必须在1-20之间" diff --git a/i18n/locales/zh-TW.yaml b/i18n/locales/zh-TW.yaml index 27759d07f378..84ebd57ed587 100644 --- a/i18n/locales/zh-TW.yaml +++ b/i18n/locales/zh-TW.yaml @@ -48,6 +48,9 @@ token.expired: "該令牌已過期" token.exhausted: "該令牌額度已用盡 TokenStatusExhausted[sk-{{.Prefix}}***{{.Suffix}}]" token.status_unavailable: "該令牌狀態不可用" token.db_error: "無效的令牌，資料庫查詢出錯，請聯繫管理員" +token.auto_groups_too_many: "每個令牌最多可選擇 {{.Max}} 個 Auto 分組" +token.auto_groups_duplicate: "Auto 分組 {{.Group}} 重複" +token.auto_groups_invalid: "Auto 分組 {{.Group}} 不可用或無權存取" # Redemption messages redemption.name_length: "兌換碼名稱長度必須在1-20之間" diff --git a/middleware/auth.go b/middleware/auth.go index 2ad09a7a8d2d..4e1436f38466 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -502,6 +502,16 @@ func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) e } common.SetContextKey(c, constant.ContextKeyTokenGroup, token.Group) common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, token.CrossGroupRetry) + if token.AutoGroups != "" { + autoGroups, err := token.GetAutoGroups() + if err != nil { + common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err)) + autoGroups = []string{} + common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, autoGroups) + } else if len(autoGroups) > 0 { + common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, autoGroups) + } + } if len(parts) > 1 { if model.IsAdmin(token.UserId) { c.Set("specific_channel_id", parts[1]) diff --git a/middleware/distributor.go b/middleware/distributor.go index bde639ddee5b..7decf0e28728 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -109,7 +109,7 @@ func Distribute() func(c *gin.Context) { channelSupportsRequestPath(preferred, c.Request.URL.Path, modelRequest.Model) { if usingGroup == "auto" { userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) - autoGroups := service.GetUserAutoGroup(userGroup) + autoGroups := service.GetRequestAutoGroups(c, userGroup) for _, g := range autoGroups { if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { selectGroup = g diff --git a/middleware/token_auto_groups_context_test.go b/middleware/token_auto_groups_context_test.go new file mode 100644 index 000000000000..ac507d5b3782 --- /dev/null +++ b/middleware/token_auto_groups_context_test.go @@ -0,0 +1,48 @@ +package middleware + +import ( + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newTokenAutoGroupsContext() *gin.Context { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + return ctx +} + +func TestSetupContextForTokenPreservesCustomAutoGroupsOrder(t *testing.T) { + ctx := newTokenAutoGroupsContext() + token := &model.Token{Id: 1, UserId: 2, AutoGroups: `["vip","default"]`} + + require.NoError(t, SetupContextForToken(ctx, token)) + value, ok := common.GetContextKey(ctx, constant.ContextKeyTokenAutoGroups) + require.True(t, ok) + assert.Equal(t, []string{"vip", "default"}, value) +} + +func TestSetupContextForTokenTreatsStoredEmptyArrayAsInheritance(t *testing.T) { + ctx := newTokenAutoGroupsContext() + token := &model.Token{Id: 1, UserId: 2, AutoGroups: `[]`} + + require.NoError(t, SetupContextForToken(ctx, token)) + _, ok := common.GetContextKey(ctx, constant.ContextKeyTokenAutoGroups) + assert.False(t, ok) +} + +func TestSetupContextForTokenMalformedAutoGroupsFailsClosed(t *testing.T) { + ctx := newTokenAutoGroupsContext() + token := &model.Token{Id: 1, UserId: 2, AutoGroups: `not-json`} + + require.NoError(t, SetupContextForToken(ctx, token)) + value, ok := common.GetContextKey(ctx, constant.ContextKeyTokenAutoGroups) + require.True(t, ok) + assert.Equal(t, []string{}, value) +} diff --git a/model/option.go b/model/option.go index 7ab64e0ded80..e7fda5231be7 100644 --- a/model/option.go +++ b/model/option.go @@ -120,6 +120,7 @@ func InitOptionMap() { common.OptionMap["Chats"] = setting.Chats2JsonString() common.OptionMap["AutoGroups"] = setting.AutoGroups2JsonString() common.OptionMap["DefaultUseAutoGroup"] = strconv.FormatBool(setting.DefaultUseAutoGroup) + common.OptionMap["MaxTokenAutoGroups"] = strconv.Itoa(setting.GetMaxTokenAutoGroups()) common.OptionMap["PayMethods"] = operation_setting.PayMethods2JsonString() common.OptionMap["GitHubClientId"] = "" common.OptionMap["GitHubClientSecret"] = "" @@ -208,6 +209,9 @@ func validateOptionValue(key string, value string) error { if key == operation_setting.ToolPriceOptionKey { return operation_setting.ValidateToolPricesJSON(value) } + if key == "MaxTokenAutoGroups" { + return setting.ValidateMaxTokenAutoGroups(value) + } return nil } @@ -413,6 +417,8 @@ func updateOptionMap(key string, value string) (err error) { err = setting.UpdateChatsByJsonString(value) case "AutoGroups": err = setting.UpdateAutoGroupsByJsonString(value) + case "MaxTokenAutoGroups": + err = setting.UpdateMaxTokenAutoGroups(value) case "CustomCallbackAddress": operation_setting.CustomCallbackAddress = value case "EpayId": diff --git a/model/option_auto_group_test.go b/model/option_auto_group_test.go new file mode 100644 index 000000000000..c1a8f168343c --- /dev/null +++ b/model/option_auto_group_test.go @@ -0,0 +1,17 @@ +package model + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestValidateOptionValueRejectsInvalidMaxTokenAutoGroups(t *testing.T) { + for _, value := range []string{"", "0", "-1", "1.5", "invalid"} { + t.Run(value, func(t *testing.T) { + assert.Error(t, validateOptionValue("MaxTokenAutoGroups", value)) + }) + } + require.NoError(t, validateOptionValue("MaxTokenAutoGroups", "999999")) +} diff --git a/model/token.go b/model/token.go index 5d62258e7920..5aa8b3d56e0e 100644 --- a/model/token.go +++ b/model/token.go @@ -28,9 +28,34 @@ type Token struct { UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota Group string `json:"group" gorm:"default:''"` CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试，仅auto分组有效 + AutoGroups string `json:"-" gorm:"type:text"` DeletedAt gorm.DeletedAt `gorm:"index"` } +func (token *Token) GetAutoGroups() ([]string, error) { + if token.AutoGroups == "" { + return nil, nil + } + var groups []string + if err := common.UnmarshalJsonStr(token.AutoGroups, &groups); err != nil { + return nil, err + } + return groups, nil +} + +func (token *Token) SetAutoGroups(groups []string) error { + if len(groups) == 0 { + token.AutoGroups = "" + return nil + } + data, err := common.Marshal(groups) + if err != nil { + return err + } + token.AutoGroups = string(data) + return nil +} + func (token *Token) Clean() { token.Key = "" } @@ -291,18 +316,16 @@ func (token *Token) Insert() error { // Update Make sure your token's fields is completed, because this will update non-zero values func (token *Token) Update() (err error) { - defer func() { - if shouldUpdateRedis(true, err) { - gopool.Go(func() { - err := cacheSetToken(*token) - if err != nil { - common.SysLog("failed to update token cache: " + err.Error()) - } - }) - } - }() err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", - "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error + "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry", "auto_groups").Updates(token).Error + if shouldUpdateRedis(true, err) { + if cacheErr := cacheSetToken(*token); cacheErr != nil { + common.SysLog("failed to update token cache: " + cacheErr.Error()) + if deleteErr := cacheDeleteToken(token.Key); deleteErr != nil { + common.SysLog("failed to invalidate token cache after update: " + deleteErr.Error()) + } + } + } return err } diff --git a/model/token_auto_groups_cache_test.go b/model/token_auto_groups_cache_test.go new file mode 100644 index 000000000000..2018502d6200 --- /dev/null +++ b/model/token_auto_groups_cache_test.go @@ -0,0 +1,57 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTokenAutoGroupsRoundTripThroughRedisHashCache(t *testing.T) { + useUserCacheMiniRedis(t) + token := Token{ + Id: 42, + UserId: 7, + Key: "token-auto-groups-cache-key", + Name: "auto-cache", + Group: "auto", + AutoGroups: `["vip","default"]`, + } + + require.NoError(t, cacheSetToken(token)) + cached, err := cacheGetTokenByKey(token.Key) + require.NoError(t, err) + assert.Equal(t, token.AutoGroups, cached.AutoGroups) + groups, err := cached.GetAutoGroups() + require.NoError(t, err) + assert.Equal(t, []string{"vip", "default"}, groups) +} + +func TestTokenUpdateSynchronouslyNarrowsPreheatedAutoGroupsCache(t *testing.T) { + truncateTables(t) + useUserCacheMiniRedis(t) + token := Token{ + UserId: 7, + Key: "token-auto-groups-update-cache-key", + Name: "auto-cache-update", + Status: common.TokenStatusEnabled, + ExpiredTime: -1, + UnlimitedQuota: true, + Group: "auto", + CrossGroupRetry: true, + AutoGroups: `["default","vip"]`, + } + require.NoError(t, token.Insert()) + require.NoError(t, cacheSetToken(token)) + + preheated, err := cacheGetTokenByKey(token.Key) + require.NoError(t, err) + assert.JSONEq(t, `["default","vip"]`, preheated.AutoGroups) + + require.NoError(t, token.SetAutoGroups([]string{"vip"})) + require.NoError(t, token.Update()) + immediate, err := cacheGetTokenByKey(token.Key) + require.NoError(t, err) + assert.JSONEq(t, `["vip"]`, immediate.AutoGroups) +} diff --git a/router/api-router.go b/router/api-router.go index 80fd65178c44..907cf1ed2885 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -238,6 +238,7 @@ func SetApiRouter(router *gin.Engine) { { tokenRoute.GET("/", controller.GetAllTokens) tokenRoute.GET("/search", middleware.SearchRateLimit(), controller.SearchTokens) + tokenRoute.GET("/auto-groups", controller.GetTokenAutoGroups) tokenRoute.GET("/:id", controller.GetToken) tokenRoute.POST("/:id/key", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.GetTokenKey) tokenRoute.POST("/", controller.AddToken) diff --git a/service/channel_select.go b/service/channel_select.go index 24c4e252bfb3..0ab88dc84ff2 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -7,7 +7,6 @@ import ( "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/model" - "github.com/QuantumNous/new-api/setting" "github.com/gin-gonic/gin" ) @@ -88,10 +87,10 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup) if param.TokenGroup == "auto" { - if len(setting.GetAutoGroups()) == 0 { + autoGroups := GetRequestAutoGroups(param.Ctx, userGroup) + if len(autoGroups) == 0 { return nil, selectGroup, errors.New("auto groups is not enabled") } - autoGroups := GetUserAutoGroup(userGroup) // startGroupIndex: the group index to start searching from // startGroupIndex: 开始搜索的分组索引 diff --git a/service/channel_select_auto_groups_test.go b/service/channel_select_auto_groups_test.go new file mode 100644 index 000000000000..e8454b389471 --- /dev/null +++ b/service/channel_select_auto_groups_test.go @@ -0,0 +1,129 @@ +package service + +import ( + "fmt" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupChannelSelectAutoGroupsTest(t *testing.T) *gorm.DB { + t.Helper() + + originalDB := model.DB + originalMemoryCacheEnabled := common.MemoryCacheEnabled + originalRetryTimes := common.RetryTimes + originalAutoGroups := setting.AutoGroups2JsonString() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalGroupRatios := ratio_setting.GroupRatio2JSONString() + originalMaxTokenAutoGroups := setting.GetMaxTokenAutoGroups() + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{})) + model.DB = db + common.MemoryCacheEnabled = true + common.RetryTimes = 0 + + require.NoError(t, setting.UpdateAutoGroupsByJsonString(`[]`)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":2}`)) + require.NoError(t, setting.UpdateMaxTokenAutoGroups("2")) + + t.Cleanup(func() { + model.DB = originalDB + common.MemoryCacheEnabled = originalMemoryCacheEnabled + common.RetryTimes = originalRetryTimes + require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatios)) + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMaxTokenAutoGroups))) + + if originalMemoryCacheEnabled && originalDB != nil && + originalDB.Migrator().HasTable(&model.Channel{}) && originalDB.Migrator().HasTable(&model.Ability{}) { + model.InitChannelCache() + } + sqlDB, err := db.DB() + if err == nil { + require.NoError(t, sqlDB.Close()) + } + }) + + return db +} + +func createChannelSelectAutoGroupsChannel(t *testing.T, db *gorm.DB, id int, group, modelName string) { + t.Helper() + priority := int64(0) + weight := uint(100) + require.NoError(t, db.Create(&model.Channel{ + Id: id, + Type: constant.ChannelTypeOpenAI, + Key: fmt.Sprintf("key-%d", id), + Status: common.ChannelStatusEnabled, + Name: fmt.Sprintf("channel-%d", id), + Weight: &weight, + Models: modelName, + Group: group, + Priority: &priority, + }).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: group, + Model: modelName, + ChannelId: id, + Enabled: true, + Priority: &priority, + Weight: weight, + }).Error) +} + +func TestCacheGetRandomSatisfiedChannelUsesTokenAutoGroupsWhenGlobalAutoIsEmpty(t *testing.T) { + db := setupChannelSelectAutoGroupsTest(t) + const modelName = "auto-groups-runtime-model" + createChannelSelectAutoGroupsChannel(t, db, 2101, "vip", modelName) + createChannelSelectAutoGroupsChannel(t, db, 2102, "default", modelName) + model.InitChannelCache() + + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip", "default"}) + common.SetContextKey(ctx, constant.ContextKeyTokenCrossGroupRetry, true) + + retry := 0 + param := &RetryParam{ + Ctx: ctx, + TokenGroup: "auto", + ModelName: modelName, + RequestPath: "/v1/chat/completions", + Retry: &retry, + } + + first, selectedGroup, err := CacheGetRandomSatisfiedChannel(param) + require.NoError(t, err) + require.NotNil(t, first) + assert.Equal(t, 2101, first.Id) + assert.Equal(t, "vip", selectedGroup) + assert.Equal(t, "vip", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup)) + assert.Empty(t, setting.GetAutoGroups(), "the selection must not depend on the global Auto list") + + param.IncreaseRetry() + second, selectedGroup, err := CacheGetRandomSatisfiedChannel(param) + require.NoError(t, err) + require.NotNil(t, second) + assert.Equal(t, 2102, second.Id) + assert.Equal(t, "default", selectedGroup) + assert.Equal(t, "default", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup)) +} diff --git a/service/group.go b/service/group.go index 8cb359bcff63..e792083e4a09 100644 --- a/service/group.go +++ b/service/group.go @@ -3,9 +3,12 @@ package service import ( "strings" + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" ) func GetUserUsableGroups(userGroup string) map[string]string { @@ -42,18 +45,67 @@ func GroupInUserUsableGroups(userGroup, groupName string) bool { return ok } +func IsUserSelectableGroup(userGroup, groupName string) bool { + if groupName == "" || groupName == "auto" { + return false + } + return GroupInUserUsableGroups(userGroup, groupName) && ratio_setting.ContainsGroupRatio(groupName) +} + // GetUserAutoGroup 根据用户分组获取自动分组设置 func GetUserAutoGroup(userGroup string) []string { - groups := GetUserUsableGroups(userGroup) autoGroups := make([]string, 0) + seen := make(map[string]struct{}) for _, group := range setting.GetAutoGroups() { - if _, ok := groups[group]; ok { - autoGroups = append(autoGroups, group) + if !IsUserSelectableGroup(userGroup, group) { + continue + } + if _, ok := seen[group]; ok { + continue } + seen[group] = struct{}{} + autoGroups = append(autoGroups, group) } return autoGroups } +// FilterUserTokenAutoGroups applies current permissions before the current +// per-token limit. It intentionally does not fall back to the global Auto list. +func FilterUserTokenAutoGroups(userGroup string, groups []string) []string { + maxCount := setting.GetMaxTokenAutoGroups() + filtered := make([]string, 0, min(len(groups), maxCount)) + seen := make(map[string]struct{}) + for _, group := range groups { + if !IsUserSelectableGroup(userGroup, group) { + continue + } + if _, ok := seen[group]; ok { + continue + } + seen[group] = struct{}{} + filtered = append(filtered, group) + if len(filtered) == maxCount { + break + } + } + return filtered +} + +// GetRequestAutoGroups resolves the ordered Auto groups for the current token. +// The absence of the context value means that the token inherits the complete +// global Auto list; a present (even empty) value is an explicit token snapshot. +func GetRequestAutoGroups(c *gin.Context, userGroup string) []string { + value, ok := common.GetContextKey(c, constant.ContextKeyTokenAutoGroups) + if !ok { + return GetUserAutoGroup(userGroup) + } + groups, ok := value.([]string) + if !ok { + return []string{} + } + return FilterUserTokenAutoGroups(userGroup, groups) +} + // GetGroupsEnabledModels 按 groups 顺序获取各分组启用的模型并去重 func GetGroupsEnabledModels(groups []string) []string { seen := make(map[string]struct{}) diff --git a/service/group_auto_groups_test.go b/service/group_auto_groups_test.go new file mode 100644 index 000000000000..1f138ad4bf1d --- /dev/null +++ b/service/group_auto_groups_test.go @@ -0,0 +1,72 @@ +package service + +import ( + "fmt" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func configureRequestAutoGroupsTest(t *testing.T) { + t.Helper() + originalMax := setting.GetMaxTokenAutoGroups() + originalAutoGroups := setting.AutoGroups2JsonString() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups("2")) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["vip","default","svip"]`)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP","svip":"SVIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) +} + +func newRequestAutoGroupsContext() *gin.Context { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + return ctx +} + +func TestGetRequestAutoGroupsInheritedListIsNotLimited(t *testing.T) { + configureRequestAutoGroupsTest(t) + ctx := newRequestAutoGroupsContext() + + groups := GetRequestAutoGroups(ctx, "default") + + assert.Equal(t, []string{"vip", "default", "svip"}, groups) +} + +func TestGetRequestAutoGroupsFiltersBeforeApplyingCurrentLimit(t *testing.T) { + configureRequestAutoGroupsTest(t) + ctx := newRequestAutoGroupsContext() + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"revoked", "vip", "default", "svip"}) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(`[]`)) + + groups := GetRequestAutoGroups(ctx, "default") + + assert.Equal(t, []string{"vip", "default"}, groups) + require.NoError(t, setting.UpdateMaxTokenAutoGroups("1")) + assert.Equal(t, []string{"vip"}, GetRequestAutoGroups(ctx, "default")) +} + +func TestGetRequestAutoGroupsDoesNotFallBackAfterPermissionChange(t *testing.T) { + configureRequestAutoGroupsTest(t) + ctx := newRequestAutoGroupsContext() + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) + + groups := GetRequestAutoGroups(ctx, "default") + + assert.Empty(t, groups) +} diff --git a/setting/auto_group.go b/setting/auto_group.go index 9261286bca93..3b509b6026e6 100644 --- a/setting/auto_group.go +++ b/setting/auto_group.go @@ -1,15 +1,27 @@ package setting import ( + "fmt" + "strconv" + "sync/atomic" + "github.com/QuantumNous/new-api/common" ) +const DefaultMaxTokenAutoGroups = 5 + var autoGroups = []string{ "default", } var DefaultUseAutoGroup = false +var maxTokenAutoGroups atomic.Int64 + +func init() { + maxTokenAutoGroups.Store(DefaultMaxTokenAutoGroups) +} + func ContainsAutoGroup(group string) bool { for _, autoGroup := range autoGroups { if autoGroup == group { @@ -35,3 +47,24 @@ func AutoGroups2JsonString() string { func GetAutoGroups() []string { return autoGroups } + +func GetMaxTokenAutoGroups() int { + return int(maxTokenAutoGroups.Load()) +} + +func ValidateMaxTokenAutoGroups(value string) error { + maxCount, err := strconv.Atoi(value) + if err != nil || maxCount <= 0 { + return fmt.Errorf("MaxTokenAutoGroups must be a positive integer") + } + return nil +} + +func UpdateMaxTokenAutoGroups(value string) error { + if err := ValidateMaxTokenAutoGroups(value); err != nil { + return err + } + maxCount, _ := strconv.Atoi(value) + maxTokenAutoGroups.Store(int64(maxCount)) + return nil +} diff --git a/setting/auto_group_test.go b/setting/auto_group_test.go new file mode 100644 index 000000000000..414c169ed706 --- /dev/null +++ b/setting/auto_group_test.go @@ -0,0 +1,29 @@ +package setting + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestUpdateMaxTokenAutoGroupsAcceptsAnyPositiveInteger(t *testing.T) { + original := GetMaxTokenAutoGroups() + t.Cleanup(func() { + require.NoError(t, UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", original))) + }) + + require.NoError(t, UpdateMaxTokenAutoGroups("123456")) + assert.Equal(t, 123456, GetMaxTokenAutoGroups()) +} + +func TestUpdateMaxTokenAutoGroupsRejectsInvalidValuesWithoutChangingState(t *testing.T) { + original := GetMaxTokenAutoGroups() + for _, value := range []string{"", "0", "-1", "1.5", "not-a-number"} { + t.Run(value, func(t *testing.T) { + assert.Error(t, UpdateMaxTokenAutoGroups(value)) + assert.Equal(t, original, GetMaxTokenAutoGroups()) + }) + } +} diff --git a/web/src/features/keys/api.ts b/web/src/features/keys/api.ts index df3cc5ff74bc..0f90490c7852 100644 --- a/web/src/features/keys/api.ts +++ b/web/src/features/keys/api.ts @@ -25,6 +25,7 @@ import type { GetApiKeysResponse, SearchApiKeysParams, ApiKeyFormData, + TokenAutoGroupsConfig, } from './types' // ============================================================================ @@ -60,6 +61,14 @@ export async function getApiKey(id: number): Promise> { return res.data } +// Get the current user's global Auto order and the per-token selection limit. +export async function getTokenAutoGroups(): Promise< + ApiResponse +> { + const res = await api.get('/api/token/auto-groups') + return res.data +} + // Create a new API key export async function createApiKey( data: ApiKeyFormData diff --git a/web/src/features/keys/components/__tests__/api-key-group-cell.test.tsx b/web/src/features/keys/components/__tests__/api-key-group-cell.test.tsx new file mode 100644 index 000000000000..5cb64ae57f7b --- /dev/null +++ b/web/src/features/keys/components/__tests__/api-key-group-cell.test.tsx @@ -0,0 +1,236 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { after, describe, test } from 'node:test' + +import { Window } from 'happy-dom' + +const domWindow = new Window() +const domGlobals = [ + 'window', + 'document', + 'navigator', + 'HTMLElement', + 'HTMLButtonElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const + +for (const key of domGlobals) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +const { act } = await import('react') +const { createRoot } = await import('react-dom/client') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { TooltipProvider } = await import('@/components/ui/tooltip') +const { ApiKeyGroupCell } = await import('../api-key-group-cell') + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { + en: { + translation: { + Auto: 'Auto', + 'Cross-group': 'Cross-group', + Ratio: 'Ratio', + 'Automatically selects the best available group with circuit breaker mechanism': + 'Automatically selects the best available group with circuit breaker mechanism', + }, + }, + }, +}) + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +function CellHarness(props: { + group: string + ratio?: number | string + crossGroupRetry?: boolean + shouldReduceMotion?: boolean +}) { + return ( + + + + + + ) +} + +describe('API key group table cell', () => { + after(() => { + domWindow.close() + }) + + test('renders two unclipped rings and a localized Auto ratio when API data uses a nonlocalized string', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render( + + ) + ) + + const badgeCell = container.querySelector( + '[data-api-key-group-cell="auto"]' + ) + assert.ok(badgeCell) + assert.equal(badgeCell.classList.contains('overflow-visible'), true) + assert.equal(badgeCell.classList.contains('overflow-hidden'), false) + + const frames = container.querySelectorAll('[data-auto-group-frame]') + const movingRings = container.querySelectorAll( + '[data-auto-group-flow-border]' + ) + assert.equal(frames.length, 2) + assert.equal(movingRings.length, 2) + for (const frame of frames) { + assert.equal(frame.classList.contains('relative'), true) + assert.equal(frame.classList.contains('overflow-visible'), true) + assert.equal(frame.classList.contains('rounded-4xl'), true) + assert.equal(frame.classList.contains('p-px'), true) + } + + const ratio = container.querySelector( + '[data-auto-group-effect="ratio"]' + ) + assert.ok(ratio) + assert.equal(ratio.textContent, 'Auto Ratio') + assert.equal(ratio.textContent?.includes('x'), false) + assert.equal(container.textContent?.includes('自动'), false) + assert.equal(container.textContent?.includes('Cross-group'), true) + + const crossGroupBadge = [ + ...container.querySelectorAll('[data-slot="status-badge"]'), + ].find((badge) => badge.textContent === 'Cross-group') + assert.ok(crossGroupBadge) + assert.equal(crossGroupBadge.closest('[data-auto-group-frame]'), null) + + await act(async () => root.unmount()) + container.remove() + }) + + test('keeps static Auto frames but omits both moving layers for reduced motion', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render() + ) + + assert.equal( + container.querySelectorAll('[data-auto-group-frame]').length, + 2 + ) + assert.equal( + container.querySelectorAll('[data-auto-group-flow-border]').length, + 0 + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('shows only the Auto badge when ratio data is unavailable', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render() + ) + + assert.equal( + container.querySelectorAll('[data-auto-group-frame]').length, + 1 + ) + assert.equal( + container.querySelectorAll('[data-auto-group-flow-border]').length, + 1 + ) + assert.equal( + container.querySelector('[data-auto-group-effect="ratio"]'), + null + ) + assert.equal(container.textContent?.includes('Auto'), true) + assert.equal(container.textContent?.includes('Ratio'), false) + + await act(async () => root.unmount()) + container.remove() + }) + + test('narrows normal group ratios to numbers and never applies Auto rings', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render( + + ) + ) + + assert.equal(container.textContent?.includes('vip'), true) + assert.equal(container.textContent?.includes('自动'), false) + assert.equal(container.querySelector('[data-auto-group-frame]'), null) + assert.equal(container.querySelector('[data-auto-group-flow-border]'), null) + + await act(async () => + root.render( + + ) + ) + + assert.equal(container.textContent?.includes('3x'), true) + assert.equal(container.querySelector('[data-auto-group-frame]'), null) + + await act(async () => root.unmount()) + container.remove() + }) +}) diff --git a/web/src/features/keys/components/__tests__/api-key-group-combobox.test.tsx b/web/src/features/keys/components/__tests__/api-key-group-combobox.test.tsx new file mode 100644 index 000000000000..5c3b6525e481 --- /dev/null +++ b/web/src/features/keys/components/__tests__/api-key-group-combobox.test.tsx @@ -0,0 +1,294 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { after, describe, test } from 'node:test' + +import { Window } from 'happy-dom' + +const domWindow = new Window() +const domGlobals = [ + 'window', + 'document', + 'navigator', + 'HTMLElement', + 'HTMLButtonElement', + 'HTMLInputElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'KeyboardEvent', + 'PointerEvent', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const + +for (const key of domGlobals) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +let shouldReduceMotion = false +const reducedMotionMediaQuery = domWindow.matchMedia('(prefers-reduced-motion)') +Object.defineProperty(reducedMotionMediaQuery, 'matches', { + configurable: true, + get: () => shouldReduceMotion, +}) +Object.defineProperty(domWindow, 'matchMedia', { + configurable: true, + value: () => reducedMotionMediaQuery, +}) + +function setReducedMotion(value: boolean) { + shouldReduceMotion = value + reducedMotionMediaQuery.dispatchEvent(new domWindow.Event('change')) +} + +const { act, useState } = await import('react') +const { createRoot } = await import('react-dom/client') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { ApiKeyGroupCombobox } = await import('../api-key-group-combobox') + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { + en: { + translation: { + Auto: 'Auto', + Ratio: 'Ratio', + 'Search...': 'Search...', + 'No group found.': 'No group found.', + 'Select a group': 'Select a group', + }, + }, + }, +}) + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +const options = [ + { + value: 'auto', + label: 'auto', + desc: 'Global automatic routing', + ratio: '自动', + }, + { value: 'default', label: 'default', desc: 'User group', ratio: 1 }, + { value: 'vip', label: 'vip', desc: 'Priority group', ratio: 3 }, +] + +function Harness(props: { initialValue: string }) { + const [value, setValue] = useState(props.initialValue) + + return ( + + + {value} + + ) +} + +function getTrigger(container: ParentNode): HTMLButtonElement { + const trigger = container.querySelector( + 'button[role="combobox"]' + ) + assert.ok(trigger) + return trigger +} + +function getCommandItem(label: string): HTMLElement { + const item = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ].find((candidate) => candidate.textContent?.includes(label)) + assert.ok(item) + return item +} + +describe('API key group combobox Auto effect', () => { + after(() => { + domWindow.close() + }) + + test('rings the selected Auto trigger and its localized ratio without rendering the API ratio text', async () => { + setReducedMotion(false) + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const trigger = getTrigger(container) + assert.equal(trigger.getAttribute('aria-expanded'), 'false') + assert.equal(trigger.dataset.autoGroupEffect, 'trigger') + assert.equal(trigger.classList.contains('bg-linear-to-r'), false) + assert.equal(trigger.classList.contains('overflow-hidden'), false) + assert.equal(trigger.classList.contains('overflow-visible'), true) + + const triggerFlowBorder = trigger.querySelector( + '[data-auto-group-flow-border]' + ) + assert.ok(triggerFlowBorder) + assert.equal(triggerFlowBorder.getAttribute('aria-hidden'), 'true') + assert.equal( + triggerFlowBorder.classList.contains('pointer-events-none'), + true + ) + assert.equal( + triggerFlowBorder.classList.contains('auto-group-flow-border'), + true + ) + + const triggerRatio = trigger.querySelector( + '[data-auto-group-effect="ratio"]' + ) + assert.ok(triggerRatio) + assert.equal(triggerRatio.textContent, 'Auto Ratio') + assert.equal(triggerRatio.textContent?.includes('Auto'), true) + assert.equal(triggerRatio.textContent?.includes('x'), false) + assert.equal(trigger.textContent?.includes('自动'), false) + assert.equal(triggerRatio.classList.contains('relative'), true) + assert.equal(triggerRatio.classList.contains('overflow-visible'), true) + assert.equal(triggerRatio.classList.contains('rounded-4xl'), true) + assert.ok(triggerRatio.querySelector('[data-auto-group-flow-border]')) + + await act(async () => trigger.click()) + assert.equal(trigger.getAttribute('aria-expanded'), 'true') + + const autoOption = getCommandItem('Global automatic routing') + assert.equal(autoOption.dataset.autoGroupEffect, 'option') + assert.equal(autoOption.getAttribute('aria-selected'), 'true') + assert.equal(autoOption.classList.contains('bg-linear-to-r'), false) + assert.equal(autoOption.classList.contains('overflow-visible'), true) + assert.ok(autoOption.querySelector('[data-auto-group-flow-border]')) + const optionRatio = autoOption.querySelector( + '[data-auto-group-effect="ratio"]' + ) + assert.ok(optionRatio) + assert.equal(optionRatio.textContent, 'Auto Ratio') + assert.ok(optionRatio.querySelector('[data-auto-group-flow-border]')) + + const defaultOption = getCommandItem('User group') + assert.equal(defaultOption.hasAttribute('data-auto-group-effect'), false) + assert.equal( + defaultOption.querySelector('[data-auto-group-flow-border]'), + null + ) + assert.equal(defaultOption.textContent?.includes('1x Ratio'), true) + assert.equal( + defaultOption.querySelector('[data-auto-group-effect="ratio"]'), + null + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('keeps search and selection behavior while leaving normal groups unstyled', async () => { + setReducedMotion(false) + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const trigger = getTrigger(container) + await act(async () => trigger.click()) + + const searchInput = document.querySelector( + 'input[placeholder="Search..."]' + ) + assert.ok(searchInput) + await act(async () => { + const valueSetter = Object.getOwnPropertyDescriptor( + domWindow.HTMLInputElement.prototype, + 'value' + )?.set + assert.ok(valueSetter) + valueSetter.call(searchInput, 'vip') + searchInput.dispatchEvent( + new domWindow.Event('input', { bubbles: true }) as unknown as Event + ) + }) + + const visibleOptions = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ] + assert.equal( + visibleOptions.some((option) => + option.textContent?.includes('Global automatic routing') + ), + false + ) + const vipOption = getCommandItem('Priority group') + await act(async () => vipOption.click()) + + assert.equal( + container.querySelector('[data-testid="selected-group"]')?.textContent, + 'vip' + ) + assert.equal(trigger.getAttribute('aria-expanded'), 'false') + assert.equal(trigger.hasAttribute('data-auto-group-effect'), false) + assert.equal(trigger.querySelector('[data-auto-group-flow-border]'), null) + + await act(async () => root.unmount()) + container.remove() + }) + + test('preserves the static Auto treatment but omits moving layers for reduced motion', async () => { + setReducedMotion(true) + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const trigger = getTrigger(container) + assert.equal(trigger.dataset.autoGroupEffect, 'trigger') + assert.equal(trigger.querySelector('[data-auto-group-flow-border]'), null) + assert.ok(trigger.querySelector('[data-auto-group-effect="ratio"]')) + + await act(async () => trigger.click()) + const autoOption = getCommandItem('Global automatic routing') + assert.equal(autoOption.dataset.autoGroupEffect, 'option') + assert.equal( + autoOption.querySelector('[data-auto-group-flow-border]'), + null + ) + assert.ok(autoOption.querySelector('[data-auto-group-effect="ratio"]')) + + await act(async () => root.unmount()) + container.remove() + setReducedMotion(false) + }) +}) diff --git a/web/src/features/keys/components/__tests__/api-keys-mutate-drawer.test.tsx b/web/src/features/keys/components/__tests__/api-keys-mutate-drawer.test.tsx new file mode 100644 index 000000000000..238d0b21d4d0 --- /dev/null +++ b/web/src/features/keys/components/__tests__/api-keys-mutate-drawer.test.tsx @@ -0,0 +1,371 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { after, afterEach, describe, test } from 'node:test' + +import { Window } from 'happy-dom' + +const domWindow = new Window() +const domGlobals = [ + 'window', + 'document', + 'navigator', + 'HTMLElement', + 'HTMLButtonElement', + 'HTMLInputElement', + 'HTMLFormElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'KeyboardEvent', + 'PointerEvent', + 'MouseEvent', + 'FocusEvent', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const + +for (const key of domGlobals) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +const { act } = await import('react') +const { createRoot } = await import('react-dom/client') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { QueryClient, QueryClientProvider } = + await import('@tanstack/react-query') +const { api } = await import('@/lib/api') +const { ApiKeysProvider } = await import('../api-keys-provider') +const { ApiKeysMutateDrawer } = await import('../api-keys-mutate-drawer') + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { en: { translation: {} } }, +}) + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +type ApiMethod = (url: string, data?: unknown) => Promise<{ data: unknown }> +type MockableApi = { + get: ApiMethod + post: ApiMethod +} +type RenderedDrawer = { + host: HTMLDivElement + queryClient: InstanceType + root: ReturnType +} + +const apiClient = api as unknown as MockableApi +const originalGet = apiClient.get +const originalPost = apiClient.post +let renderedDrawer: RenderedDrawer | null = null + +function installApiFixtures(createdPayloads: Array>) { + apiClient.get = async (url) => { + switch (url) { + case '/api/status': + return { data: { data: { default_use_auto_group: true } } } + case '/api/user/models': + return { data: { success: true, data: [] } } + case '/api/user/self/groups': + return { + data: { + success: true, + data: { + auto: { desc: 'Automatic routing', ratio: 'auto' }, + default: { desc: 'Standard access', ratio: 1 }, + vip: { desc: 'Priority access', ratio: 2 }, + }, + }, + } + case '/api/token/auto-groups': + return { + data: { + success: true, + data: { groups: ['vip', 'default'], max_count: 3 }, + }, + } + default: + throw new Error(`Unexpected GET ${url}`) + } + } + apiClient.post = async (url, data) => { + assert.equal(url, '/api/token/') + assert.ok(data && typeof data === 'object') + createdPayloads.push(data as Record) + return { data: { success: true, data: {} } } + } +} + +async function waitForCondition( + condition: () => boolean, + failureMessage: string +): Promise { + if (condition()) return + + await new Promise((resolve, reject) => { + const observer = new MutationObserver(() => { + if (!condition()) return + clearTimeout(timeoutId) + observer.disconnect() + resolve() + }) + const timeoutId = setTimeout(() => { + observer.disconnect() + reject(new Error(`${failureMessage}: ${document.body.textContent}`)) + }, 1500) + + observer.observe(document, { + attributes: true, + childList: true, + characterData: true, + subtree: true, + }) + }) +} + +async function renderCreateDrawer(): Promise { + const host = document.createElement('div') + document.body.append(host) + const root = createRoot(host) + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }) + const freshAt = Date.now() + 60_000 + queryClient.setQueryData( + ['status'], + { default_use_auto_group: true }, + { updatedAt: freshAt } + ) + queryClient.setQueryData( + ['user-models'], + { success: true, data: [] }, + { updatedAt: freshAt } + ) + queryClient.setQueryData( + ['user-groups'], + { + success: true, + data: { + auto: { desc: 'Automatic routing', ratio: 'auto' }, + default: { desc: 'Standard access', ratio: 1 }, + vip: { desc: 'Priority access', ratio: 2 }, + }, + }, + { updatedAt: freshAt } + ) + queryClient.setQueryData( + ['token-auto-groups'], + { + success: true, + data: { groups: ['vip', 'default'], max_count: 3 }, + }, + { updatedAt: freshAt } + ) + renderedDrawer = { host, queryClient, root } + + await act(async () => + root.render( + + + + undefined} /> + + + + ) + ) + await act(async () => + waitForCondition(() => { + const saveButton = findButton('Save changes', false) + return saveButton !== null && !saveButton.disabled + }, 'API key drawer did not finish initializing') + ) +} + +function findButton(text: string, required: true): HTMLButtonElement +function findButton(text: string, required: false): HTMLButtonElement | null +function findButton(text: string, required = true): HTMLButtonElement | null { + const button = [ + ...document.querySelectorAll('button'), + ].find((candidate) => candidate.textContent?.includes(text)) + if (required) assert.ok(button, `Expected button containing "${text}"`) + return button ?? null +} + +function getControlByLabel(labelText: string): T { + const label = [...document.querySelectorAll('label')].find( + (candidate) => candidate.textContent?.trim() === labelText + ) + assert.ok(label, `Expected label "${labelText}"`) + assert.ok(label.htmlFor) + const control = + label.control ?? + label + .closest('[data-slot="form-item"]') + ?.querySelector( + '[data-slot="form-control"], input, textarea, button[role="combobox"], [role="group"]' + ) + assert.ok(control) + return control as T +} + +async function changeInput(input: HTMLInputElement, value: string) { + await act(async () => { + const valueSetter = Object.getOwnPropertyDescriptor( + domWindow.HTMLInputElement.prototype, + 'value' + )?.set + assert.ok(valueSetter) + valueSetter.call(input, value) + input.dispatchEvent( + new domWindow.Event('input', { bubbles: true }) as unknown as Event + ) + }) +} + +async function selectComboboxOption( + trigger: HTMLButtonElement, + optionDescription: string +) { + await act(async () => trigger.click()) + const option = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ].find((candidate) => candidate.textContent?.includes(optionDescription)) + assert.ok(option, `Expected option containing "${optionDescription}"`) + await act(async () => option.click()) +} + +afterEach(async () => { + apiClient.get = originalGet + apiClient.post = originalPost + domWindow.localStorage.clear() + if (renderedDrawer) { + await act(async () => renderedDrawer?.root.unmount()) + renderedDrawer.queryClient.clear() + renderedDrawer.host.remove() + renderedDrawer = null + } + document.body.replaceChildren() +}) + +after(() => { + domWindow.close() +}) + +describe('API keys mutate drawer Auto group integration', () => { + test('inherits the root Auto order and sends an empty override for every batch-created key', async () => { + const createdPayloads: Array> = [] + installApiFixtures(createdPayloads) + await renderCreateDrawer() + + const groupTrigger = getControlByLabel('Group') + assert.equal(groupTrigger.textContent?.includes('auto'), true) + assert.equal( + document.body.textContent?.includes( + 'Using the complete global Auto order (2 groups)' + ), + true + ) + assert.deepEqual( + [ + ...document.querySelectorAll('[data-slot="global-auto-order-name"]'), + ].map((item) => item.textContent), + ['vip', 'default'] + ) + assert.equal(findButton('Restore global Auto', true).disabled, true) + + await changeInput(getControlByLabel('Name'), 'batch') + await changeInput(getControlByLabel('Quantity'), '2') + await act(async () => findButton('Save changes', true).click()) + await act(async () => + waitForCondition( + () => createdPayloads.length === 2, + 'batch API keys were not created' + ) + ) + + assert.equal(createdPayloads.length, 2) + assert.equal(createdPayloads[0]?.name, 'batch') + for (const payload of createdPayloads) { + assert.equal(payload.group, 'auto') + assert.deepEqual(payload.auto_groups, []) + assert.equal(payload.cross_group_retry, true) + } + }) + + test('preserves an unsaved custom order and mode after Auto to ordinary to Auto changes', async () => { + const createdPayloads: Array> = [] + installApiFixtures(createdPayloads) + await renderCreateDrawer() + + const autoOrderControl = getControlByLabel('Auto group order') + const addGroupTrigger = autoOrderControl.querySelector( + 'button[role="combobox"]' + ) + assert.ok(addGroupTrigger) + await selectComboboxOption(addGroupTrigger, 'Priority access') + + assert.ok(document.querySelector('button[aria-label="Remove vip"]')) + assert.equal( + document.body.textContent?.includes('1 / 3 groups selected'), + true + ) + assert.equal(findButton('Restore global Auto', true).disabled, false) + + const groupTrigger = getControlByLabel('Group') + await selectComboboxOption(groupTrigger, 'Standard access') + assert.equal( + document.querySelector('button[aria-label="Remove vip"]'), + null + ) + await selectComboboxOption(groupTrigger, 'Automatic routing') + + assert.ok(document.querySelector('button[aria-label="Remove vip"]')) + assert.equal( + document.body.textContent?.includes('1 / 3 groups selected'), + true + ) + assert.equal(findButton('Restore global Auto', true).disabled, false) + + await changeInput(getControlByLabel('Name'), 'custom') + await act(async () => findButton('Save changes', true).click()) + await act(async () => + waitForCondition( + () => createdPayloads.length === 1, + 'custom-order API key was not created' + ) + ) + assert.deepEqual(createdPayloads[0]?.auto_groups, ['vip']) + }) +}) diff --git a/web/src/features/keys/components/__tests__/auto-group-order-editor.test.tsx b/web/src/features/keys/components/__tests__/auto-group-order-editor.test.tsx new file mode 100644 index 000000000000..f37f50755f2d --- /dev/null +++ b/web/src/features/keys/components/__tests__/auto-group-order-editor.test.tsx @@ -0,0 +1,540 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { after, describe, test } from 'node:test' + +import { Window } from 'happy-dom' + +const domWindow = new Window() +const domGlobals = [ + 'window', + 'document', + 'navigator', + 'HTMLElement', + 'HTMLButtonElement', + 'HTMLInputElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'KeyboardEvent', + 'PointerEvent', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const + +for (const key of domGlobals) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +const { act, useState } = await import('react') +const { createRoot } = await import('react-dom/client') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { AutoGroupOrderEditor } = await import('../auto-group-order-editor') + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { + en: { + translation: { + '{{count}} / {{max}} groups selected': + '{{count}} / {{max}} groups selected', + 'Add Auto group': 'Add Auto group', + 'Auto group order': 'Auto group order', + 'Drag {{group}} to reorder': 'Drag {{group}} to reorder', + 'Inherit global Auto order': 'Inherit global Auto order', + 'Maximum {{max}} groups selected': 'Maximum {{max}} groups selected', + 'Move {{group}} down': 'Move {{group}} down', + 'Move {{group}} up': 'Move {{group}} up', + 'No available groups in the global Auto order.': + 'No available groups in the global Auto order.', + 'No valid custom Auto groups remain. Add a group or restore global Auto.': + 'No valid custom Auto groups remain. Add a group or restore global Auto.', + 'No custom groups. Saving will inherit the complete global Auto order.': + 'No custom groups. Saving will inherit the complete global Auto order.', + 'Remove {{group}}': 'Remove {{group}}', + 'Restore global Auto': 'Restore global Auto', + Ratio: 'Ratio', + 'Search...': 'Search...', + 'No group found.': 'No group found.', + 'Select a group': 'Select a group', + 'Using the complete global Auto order ({{count}} groups)': + 'Using the complete global Auto order ({{count}} groups)', + }, + }, + }, +}) + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +const globalOptions = [ + { value: 'vip', label: 'VIP', desc: 'Priority access', ratio: 3 }, + { value: 'default', label: 'Default', desc: 'Standard access', ratio: 1 }, + { value: 'team', label: 'Team', desc: 'Shared access', ratio: 2 }, +] + +function Harness(props: { initialGroups?: string[] }) { + const [groups, setGroups] = useState( + props.initialGroups ?? ['default', 'vip'] + ) + const [mode, setMode] = useState<'inherit' | 'custom'>('custom') + return ( + + { + setGroups(value.groups) + setMode(value.mode) + }} + /> + {groups.join(',')} + {mode} + + ) +} + +function InheritanceHarness(props: { globalOptions?: typeof globalOptions }) { + const [groups, setGroups] = useState([]) + const [mode, setMode] = useState<'inherit' | 'custom'>('inherit') + + return ( + + { + setGroups(value.groups) + setMode(value.mode) + }} + /> + {groups.join(',')} + {mode} + + ) +} + +function CustomEmptyHarness() { + const [groups, setGroups] = useState([]) + const [mode, setMode] = useState<'inherit' | 'custom'>('custom') + + return ( + + { + setGroups(value.groups) + setMode(value.mode) + }} + /> + {groups.join(',')} + {mode} + + ) +} + +function findButton(container: ParentNode, label: string): HTMLButtonElement { + const button = container.querySelector( + `button[aria-label="${label}"]` + ) + assert.ok(button) + return button +} + +describe('Auto group order editor', () => { + after(() => { + domWindow.close() + }) + + test('enforces the limit and exposes accessible reorder controls', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const addButton = container.querySelector( + 'button[role="combobox"]' + ) + assert.ok(addButton) + assert.equal(addButton.disabled, true) + assert.equal(container.textContent?.includes('2 / 2 groups selected'), true) + assert.ok( + container.querySelector('[role="group"][aria-label="Auto group order"]') + ) + assert.equal( + findButton(container, 'Drag default to reorder').type, + 'button' + ) + + await act(async () => findButton(container, 'Move default down').click()) + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'vip,default' + ) + + await act(async () => { + findButton(container, 'Drag vip to reorder').dispatchEvent( + new domWindow.KeyboardEvent('keydown', { + key: 'ArrowDown', + bubbles: true, + }) as unknown as KeyboardEvent + ) + }) + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'default,vip' + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('adds and removes groups, then restores inheritance as an empty value', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + await act(async () => findButton(container, 'Remove vip').click()) + + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'default' + ) + const addButton = container.querySelector( + 'button[role="combobox"]' + ) + assert.ok(addButton) + assert.equal(addButton.disabled, false) + + await act(async () => addButton.click()) + const teamOption = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ].find((option) => option.textContent?.includes('team')) + assert.ok(teamOption) + await act(async () => teamOption.click()) + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'default,team' + ) + assert.equal(addButton.disabled, true) + + const restoreButton = [...container.querySelectorAll('button')].find( + (button) => button.textContent?.includes('Restore global Auto') + ) + assert.ok(restoreButton) + await act(async () => restoreButton.click()) + + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + '' + ) + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'inherit' + ) + assert.equal( + container.textContent?.includes( + 'Using the complete global Auto order (3 groups)' + ), + true + ) + + const inheritedItems = container.querySelectorAll( + '[data-slot="global-auto-order"] > li' + ) + assert.deepEqual( + [...inheritedItems].map( + (item) => + item.querySelector('[data-slot="global-auto-order-name"]') + ?.textContent + ), + ['VIP', 'Default', 'Team'] + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('shows the complete inherited order with metadata beyond the custom limit', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + assert.equal( + container.textContent?.includes( + 'Using the complete global Auto order (3 groups)' + ), + true + ) + assert.equal( + container.textContent?.includes('0 / 2 groups selected'), + false + ) + + const order = container.querySelector( + '[data-slot="global-auto-order"]' + ) + assert.ok(order) + assert.equal(order.classList.contains('overflow-y-auto'), true) + assert.equal(order.classList.contains('flex-wrap'), true) + + const items = [...order.querySelectorAll('li')] + assert.equal(items.length, 3) + assert.equal( + order.querySelectorAll('[data-slot="global-auto-order-connector"]') + .length, + 2 + ) + assert.deepEqual( + items.map((item) => ({ + index: item.querySelector('[data-slot="global-auto-order-index"]') + ?.textContent, + name: item.querySelector('[data-slot="global-auto-order-name"]') + ?.textContent, + title: item + .querySelector('[data-slot="global-auto-order-chip"]') + ?.getAttribute('title'), + description: item.querySelector( + '[data-slot="global-auto-order-description"]' + )?.textContent, + ratio: item.querySelector('[data-slot="badge"]')?.textContent, + })), + [ + { + index: '1', + name: 'VIP', + title: 'Priority access', + description: 'Priority access', + ratio: '3x Ratio', + }, + { + index: '2', + name: 'Default', + title: 'Standard access', + description: 'Standard access', + ratio: '1x Ratio', + }, + { + index: '3', + name: 'Team', + title: 'Shared access', + description: 'Shared access', + ratio: '2x Ratio', + }, + ] + ) + + for (const item of items) { + const chip = item.querySelector('[data-slot="global-auto-order-chip"]') + assert.ok(chip) + const description = item.querySelector( + '[data-slot="global-auto-order-description"]' + ) + assert.ok(description) + assert.equal(description.classList.contains('sr-only'), true) + } + + assert.equal( + items[0]?.querySelector('[data-slot="global-auto-order-connector"]'), + null + ) + for (const item of items.slice(1)) { + const connector = item.querySelector( + '[data-slot="global-auto-order-connector"]' + ) + assert.ok(connector) + assert.equal(connector.getAttribute('aria-hidden'), 'true') + } + + assert.equal(container.querySelector('[aria-label^="Drag "]'), null) + assert.equal(container.querySelector('[aria-label^="Move "]'), null) + assert.equal(container.querySelector('[aria-label^="Remove "]'), null) + + const restoreButton = [...container.querySelectorAll('button')].find( + (button) => button.textContent?.includes('Restore global Auto') + ) + assert.ok(restoreButton) + assert.equal(restoreButton.disabled, true) + + await act(async () => root.unmount()) + container.remove() + }) + + test('shows an explicit empty state when the global Auto order has no groups', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render() + ) + + assert.equal( + container.textContent?.includes( + 'Using the complete global Auto order (0 groups)' + ), + true + ) + assert.equal( + container.textContent?.includes( + 'No available groups in the global Auto order.' + ), + true + ) + assert.equal( + container.querySelector('[data-slot="global-auto-order"]'), + null + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('keeps an empty custom order distinct from global inheritance', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'custom' + ) + assert.equal( + container.textContent?.includes( + 'No valid custom Auto groups remain. Add a group or restore global Auto.' + ), + true + ) + assert.equal( + container.querySelector('[data-slot="global-auto-order"]'), + null + ) + + const restoreButton = [...container.querySelectorAll('button')].find( + (button) => button.textContent?.includes('Restore global Auto') + ) + assert.ok(restoreButton) + assert.equal(restoreButton.disabled, false) + await act(async () => restoreButton.click()) + + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'inherit' + ) + assert.ok(container.querySelector('[data-slot="global-auto-order"]')) + + await act(async () => root.unmount()) + container.remove() + }) + + test('adding a group from inheritance explicitly creates a custom order', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const addButton = container.querySelector( + 'button[role="combobox"]' + ) + assert.ok(addButton) + await act(async () => addButton.click()) + const vipOption = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ].find((option) => option.textContent?.includes('VIP')) + assert.ok(vipOption) + await act(async () => vipOption.click()) + + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'custom' + ) + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'vip' + ) + assert.equal( + container.querySelector('[data-slot="global-auto-order"]'), + null + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('removing the last custom group does not silently enable inheritance', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + await act(async () => findButton(container, 'Remove default').click()) + + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + '' + ) + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'custom' + ) + assert.equal( + container.textContent?.includes( + 'No valid custom Auto groups remain. Add a group or restore global Auto.' + ), + true + ) + + await act(async () => root.unmount()) + container.remove() + }) +}) diff --git a/web/src/features/keys/components/api-key-group-cell.tsx b/web/src/features/keys/components/api-key-group-cell.tsx new file mode 100644 index 000000000000..21a1bc7d9a04 --- /dev/null +++ b/web/src/features/keys/components/api-key-group-cell.tsx @@ -0,0 +1,90 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { useTranslation } from 'react-i18next' + +import { BadgeCell, TruncatedCell } from '@/components/data-table' +import { GroupBadge } from '@/components/group-badge' +import { StatusBadge } from '@/components/status-badge' +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from '@/components/ui/tooltip' + +import { + // AutoGroupBadge, + GroupRatioBadge, + type GroupRatio, +} from './auto-group-visuals' + +type ApiKeyGroupCellProps = { + crossGroupRetry: boolean + group: string + ratio?: GroupRatio + shouldReduceMotion: boolean +} + +export function ApiKeyGroupCell(props: ApiKeyGroupCellProps) { + const { t } = useTranslation() + + if (props.group !== 'auto') { + const ratio = typeof props.ratio === 'number' ? props.ratio : undefined + return ( + + + + ) + } + + return ( + + + } + > + + {/**/} + + + + + {t( + 'Automatically selects the best available group with circuit breaker mechanism' + )} + + + + ) +} diff --git a/web/src/features/keys/components/api-key-group-combobox.tsx b/web/src/features/keys/components/api-key-group-combobox.tsx index 2593eff016eb..a5dada76d147 100644 --- a/web/src/features/keys/components/api-key-group-combobox.tsx +++ b/web/src/features/keys/components/api-key-group-combobox.tsx @@ -20,7 +20,6 @@ import { Check, ChevronsUpDown } from 'lucide-react' import { useMemo, useState } from 'react' import { useTranslation } from 'react-i18next' -import { Badge } from '@/components/ui/badge' import { Button } from '@/components/ui/button' import { Command, @@ -35,8 +34,15 @@ import { PopoverContent, PopoverTrigger, } from '@/components/ui/popover' +import { useMediaQuery } from '@/hooks' import { cn } from '@/lib/utils' +import { + AUTO_GROUP_FRAME_CLASS_NAME, + AutoGroupFlowBorder, + GroupRatioBadge, +} from './auto-group-visuals' + export type ApiKeyGroupOption = { value: string label: string @@ -52,50 +58,6 @@ type ApiKeyGroupComboboxProps = { disabled?: boolean } -function formatGroupRatio( - ratio: ApiKeyGroupOption['ratio'], - ratioLabel: string -) { - if (ratio === undefined || ratio === null || ratio === '') return null - return `${ratio}x ${ratioLabel}` -} - -function getRatioBadgeClassName(ratio: ApiKeyGroupOption['ratio']) { - if (typeof ratio !== 'number') { - return 'border-emerald-200 bg-emerald-50 text-emerald-700 dark:border-emerald-900/60 dark:bg-emerald-950/40 dark:text-emerald-300' - } - - if (ratio > 5) { - return 'border-rose-200 bg-rose-50 text-rose-700 dark:border-rose-900/60 dark:bg-rose-950/40 dark:text-rose-300' - } - if (ratio > 3) { - return 'border-orange-200 bg-orange-50 text-orange-700 dark:border-orange-900/60 dark:bg-orange-950/40 dark:text-orange-300' - } - if (ratio > 1) { - return 'border-blue-200 bg-blue-50 text-blue-700 dark:border-blue-900/60 dark:bg-blue-950/40 dark:text-blue-300' - } - return 'border-emerald-200 bg-emerald-50 text-emerald-700 dark:border-emerald-900/60 dark:bg-emerald-950/40 dark:text-emerald-300' -} - -function GroupRatioBadge({ ratio }: { ratio: ApiKeyGroupOption['ratio'] }) { - const { t } = useTranslation() - const label = formatGroupRatio(ratio, t('Ratio')) - - if (!label) return null - - return ( - - {label} - - ) -} - export function ApiKeyGroupCombobox({ options, value, @@ -106,7 +68,9 @@ export function ApiKeyGroupCombobox({ const { t } = useTranslation() const [open, setOpen] = useState(false) const [searchValue, setSearchValue] = useState('') + const shouldReduceMotion = useMediaQuery('(prefers-reduced-motion: reduce)') const selectedOption = options.find((option) => option.value === value) + const isAutoSelected = selectedOption?.value === 'auto' const filteredOptions = useMemo(() => { const search = searchValue.trim().toLowerCase() @@ -138,11 +102,22 @@ export function ApiKeyGroupCombobox({ variant='outline' role='combobox' aria-expanded={open} + data-auto-group-effect={isAutoSelected ? 'trigger' : undefined} disabled={disabled} - className='border-input bg-muted/40 hover:bg-muted/55 hover:text-foreground active:bg-background data-popup-open:border-ring data-popup-open:bg-background data-popup-open:ring-ring/20 h-auto min-h-14 w-full justify-between gap-2 rounded-lg px-3 py-2 text-start shadow-none transition-[background-color,border-color,box-shadow] duration-150 data-popup-open:ring-[3px] sm:min-h-20 sm:gap-3 sm:px-4 sm:py-3' + className={cn( + 'border-input bg-muted/40 hover:bg-muted/55 hover:text-foreground active:bg-background data-popup-open:border-ring data-popup-open:bg-background data-popup-open:ring-ring/20 relative h-auto min-h-14 w-full justify-between gap-2 rounded-lg px-3 py-2 text-start shadow-none transition-[background-color,border-color,box-shadow] duration-150 data-popup-open:ring-[3px] sm:min-h-20 sm:gap-3 sm:px-4 sm:py-3', + isAutoSelected && + cn( + AUTO_GROUP_FRAME_CLASS_NAME, + 'hover:border-primary/55 data-popup-open:border-primary/55 data-popup-open:ring-primary/20' + ) + )} /> } > + {isAutoSelected && ( + + )} @@ -155,10 +130,17 @@ export function ApiKeyGroupCombobox({ )} - + - +