diff --git a/controller/relay.go b/controller/relay.go index b6eb5735..aa71153d 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -195,6 +195,10 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } addUsedChannel(c, channel.Id) + if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil { + maxAPIError = billingErr + break + } bodyStorage, bodyErr := common.GetBodyStorage(c) if bodyErr != nil { // Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path) @@ -567,6 +571,10 @@ func RelayTask(c *gin.Context) { } addUsedChannel(c, channel.Id) + if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil { + taskErr = service.TaskErrorWrapperLocal(billingErr.Err, "prepare_tiered_billing_failed", billingErr.StatusCode) + break + } bodyStorage, bodyErr := common.GetBodyStorage(c) if bodyErr != nil { if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) { diff --git a/dto/openai_request.go b/dto/openai_request.go index 16cf20f9..2e09795e 100644 --- a/dto/openai_request.go +++ b/dto/openai_request.go @@ -90,6 +90,7 @@ type GeneralOpenAIRequest struct { // Ali Qwen Params VlHighResolutionImages json.RawMessage `json:"vl_high_resolution_images,omitempty"` EnableThinking json.RawMessage `json:"enable_thinking,omitempty"` + ThinkingBudget json.RawMessage `json:"thinking_budget,omitempty"` ChatTemplateKwargs json.RawMessage `json:"chat_template_kwargs,omitempty"` EnableSearch json.RawMessage `json:"enable_search,omitempty"` // ollama Params @@ -108,6 +109,14 @@ type GeneralOpenAIRequest struct { ReasoningSplit json.RawMessage `json:"reasoning_split,omitempty"` } +func (r GeneralOpenAIRequest) MarshalJSON() ([]byte, error) { + type Alias GeneralOpenAIRequest + if !IsQwenThinkingBudgetModel(r.Model) { + r.ThinkingBudget = nil + } + return common.Marshal((*Alias)(&r)) +} + func (r *GeneralOpenAIRequest) GetTokenCountMeta() *types.TokenCountMeta { var tokenCountMeta types.TokenCountMeta var texts = make([]string, 0) @@ -223,6 +232,14 @@ func IsOpenAIGPT5Model(modelName string) bool { return strings.HasPrefix(modelName, "gpt-5") } +func IsQwenThinkingBudgetModel(modelName string) bool { + normalized := strings.ToLower(strings.TrimSpace(modelName)) + return strings.HasPrefix(normalized, "qwen") || + strings.Contains(normalized, "/qwen") || + strings.HasPrefix(normalized, "qwq") || + strings.Contains(normalized, "/qwq") +} + func (r *GeneralOpenAIRequest) GetSystemRoleName() string { if IsOpenAIReasoningOModel(r.Model) { if !strings.HasPrefix(r.Model, "o1-mini") && !strings.HasPrefix(r.Model, "o1-preview") { @@ -882,10 +899,19 @@ type OpenAIResponsesRequest struct { Prompt json.RawMessage `json:"prompt,omitempty"` // qwen EnableThinking json.RawMessage `json:"enable_thinking,omitempty"` + ThinkingBudget json.RawMessage `json:"thinking_budget,omitempty"` // perplexity Preset json.RawMessage `json:"preset,omitempty"` } +func (r OpenAIResponsesRequest) MarshalJSON() ([]byte, error) { + type Alias OpenAIResponsesRequest + if !IsQwenThinkingBudgetModel(r.Model) { + r.ThinkingBudget = nil + } + return common.Marshal((*Alias)(&r)) +} + func (r *OpenAIResponsesRequest) GetTokenCountMeta() *types.TokenCountMeta { var fileMeta = make([]*types.FileMeta, 0) var texts = make([]string, 0) diff --git a/dto/openai_request_zero_value_test.go b/dto/openai_request_zero_value_test.go index 2ffbb01b..2b7a4739 100644 --- a/dto/openai_request_zero_value_test.go +++ b/dto/openai_request_zero_value_test.go @@ -1,6 +1,7 @@ package dto import ( + "encoding/json" "testing" "github.com/MAX-API-Next/MAX-API/common" @@ -50,6 +51,55 @@ func TestGeneralOpenAIRequestPreserveExplicitZeroValues(t *testing.T) { require.True(t, gjson.GetBytes(encoded, "return_related_questions").Exists()) } +func TestGeneralOpenAIRequestPreserveQwenThinkingBudget(t *testing.T) { + raw := []byte(`{ + "model":"qwen-plus", + "thinking_budget":0 + }`) + + var req GeneralOpenAIRequest + err := common.Unmarshal(raw, &req) + require.NoError(t, err) + + encoded, err := common.Marshal(req) + require.NoError(t, err) + + value := gjson.GetBytes(encoded, "thinking_budget") + require.True(t, value.Exists()) + require.Equal(t, int64(0), value.Int()) +} + +func TestGeneralOpenAIRequestDropsThinkingBudgetForNonQwenModel(t *testing.T) { + req := GeneralOpenAIRequest{ + Model: "gpt-4.1", + ThinkingBudget: json.RawMessage(`128`), + } + + encoded, err := common.Marshal(req) + require.NoError(t, err) + + require.False(t, gjson.GetBytes(encoded, "thinking_budget").Exists()) +} + +func TestIsQwenThinkingBudgetModel(t *testing.T) { + tests := []struct { + model string + want bool + }{ + {model: "qwen-plus", want: true}, + {model: "QWQ-32B", want: true}, + {model: "accounts/fireworks/models/qwen3-235b-a22b", want: true}, + {model: "deepseek-r1", want: false}, + {model: "gpt-4.1", want: false}, + } + + for _, tt := range tests { + t.Run(tt.model, func(t *testing.T) { + require.Equal(t, tt.want, IsQwenThinkingBudgetModel(tt.model)) + }) + } +} + func TestResolveMaxTokens(t *testing.T) { legacy := uint(64) completion := uint(128) @@ -116,6 +166,32 @@ func TestOpenAIResponsesRequestPreserveExplicitZeroValues(t *testing.T) { require.True(t, gjson.GetBytes(encoded, "top_p").Exists()) } +func TestOpenAIResponsesRequestPreserveQwenThinkingBudget(t *testing.T) { + req := OpenAIResponsesRequest{ + Model: "qwen-plus", + ThinkingBudget: json.RawMessage(`0`), + } + + encoded, err := common.Marshal(req) + require.NoError(t, err) + + value := gjson.GetBytes(encoded, "thinking_budget") + require.True(t, value.Exists()) + require.Equal(t, int64(0), value.Int()) +} + +func TestOpenAIResponsesRequestDropsThinkingBudgetForNonQwenModel(t *testing.T) { + req := OpenAIResponsesRequest{ + Model: "gpt-4.1", + ThinkingBudget: json.RawMessage(`128`), + } + + encoded, err := common.Marshal(req) + require.NoError(t, err) + + require.False(t, gjson.GetBytes(encoded, "thinking_budget").Exists()) +} + func TestGeneralOpenAIRequestGetSystemRoleName(t *testing.T) { tests := []struct { name string diff --git a/model/log.go b/model/log.go index 71f980af..5ecc4250 100644 --- a/model/log.go +++ b/model/log.go @@ -363,7 +363,7 @@ func formatUserLogOther(other string, stripAuditContent bool) string { delete(otherMap, "admin_info") delete(otherMap, "audit_info") // delete(otherMap, "reject_reason") - delete(otherMap, "stream_status") + // Keep stream_status visible so users can diagnose interrupted streams. if auditInfo != nil { if stripAuditContent { stripAuditContentFields(auditInfo) diff --git a/model/log_test.go b/model/log_test.go index 06679274..6b4421df 100644 --- a/model/log_test.go +++ b/model/log_test.go @@ -1093,7 +1093,9 @@ func TestFormatUserLogDetailExposesOnlyUserAuditContent(t *testing.T) { require.Equal(t, 42, log.Id) require.Empty(t, log.ChannelName) require.NotContains(t, other, "admin_info") - require.NotContains(t, other, "stream_status") + streamStatus, ok := other["stream_status"].(map[string]interface{}) + require.True(t, ok) + require.Equal(t, "error", streamStatus["status"]) auditInfo, ok := other["audit_info"].(map[string]interface{}) require.True(t, ok) diff --git a/relay/channel/deepseek/adaptor.go b/relay/channel/deepseek/adaptor.go index 443539ec..336e9242 100644 --- a/relay/channel/deepseek/adaptor.go +++ b/relay/channel/deepseek/adaptor.go @@ -68,6 +68,8 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { switch info.RelayMode { case constant.RelayModeCompletions: return fmt.Sprintf("%s/completions", fimBaseUrl), nil + case constant.RelayModeResponses: + return fmt.Sprintf("%s/responses", info.ChannelBaseUrl), nil default: return fmt.Sprintf("%s/v1/chat/completions", info.ChannelBaseUrl), nil } @@ -158,9 +160,33 @@ func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.Rela return nil, errors.New("not implemented") } -func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) { - // TODO implement me - return nil, errors.New("not implemented") +func (a *Adaptor) ConvertOpenAIResponsesRequest(_ *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) { + applyDeepSeekV4ResponsesThinkingSuffix(info, &request) + return request, nil +} + +func applyDeepSeekV4ResponsesThinkingSuffix(info *relaycommon.RelayInfo, request *dto.OpenAIResponsesRequest) { + modelName := request.Model + if info != nil && info.ChannelMeta != nil && info.UpstreamModelName != "" { + modelName = info.UpstreamModelName + } + baseModel, thinkingType, effort, ok := reasoning.ParseDeepSeekV4ThinkingSuffix(modelName) + if ok { + if thinkingType == "disabled" { + effort = "none" + } + request.Model = baseModel + if request.Reasoning == nil { + request.Reasoning = &dto.Reasoning{} + } + request.Reasoning.Effort = effort + if info != nil && info.ChannelMeta != nil { + info.UpstreamModelName = baseModel + } + } + if info != nil && request.Reasoning != nil { + info.ReasoningEffort = request.Reasoning.Effort + } } func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) { diff --git a/relay/channel/deepseek/adaptor_test.go b/relay/channel/deepseek/adaptor_test.go new file mode 100644 index 00000000..2978776c --- /dev/null +++ b/relay/channel/deepseek/adaptor_test.go @@ -0,0 +1,53 @@ +package deepseek + +import ( + "testing" + + "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/relay/constant" + "github.com/stretchr/testify/require" +) + +func TestGetRequestURLSupportsResponses(t *testing.T) { + adaptor := &Adaptor{} + url, err := adaptor.GetRequestURL(&relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelBaseUrl: "https://api.deepseek.com", + }, + RelayMode: constant.RelayModeResponses, + }) + + require.NoError(t, err) + require.Equal(t, "https://api.deepseek.com/responses", url) +} + +func TestConvertOpenAIResponsesRequestAppliesDeepSeekV4ThinkingSuffix(t *testing.T) { + info := &relaycommon.RelayInfo{} + converted, err := (&Adaptor{}).ConvertOpenAIResponsesRequest(nil, info, dto.OpenAIResponsesRequest{ + Model: "deepseek-v4-pro-none", + }) + + require.NoError(t, err) + request, ok := converted.(dto.OpenAIResponsesRequest) + require.True(t, ok) + require.Equal(t, "deepseek-v4-pro", request.Model) + require.NotNil(t, request.Reasoning) + require.Equal(t, "none", request.Reasoning.Effort) + require.Equal(t, "none", info.ReasoningEffort) +} + +func TestConvertOpenAIResponsesRequestKeepsPlainModel(t *testing.T) { + info := &relaycommon.RelayInfo{} + converted, err := (&Adaptor{}).ConvertOpenAIResponsesRequest(nil, info, dto.OpenAIResponsesRequest{ + Model: "deepseek-chat", + Reasoning: &dto.Reasoning{Effort: "low"}, + }) + + require.NoError(t, err) + request, ok := converted.(dto.OpenAIResponsesRequest) + require.True(t, ok) + require.Equal(t, "deepseek-chat", request.Model) + require.Equal(t, "low", request.Reasoning.Effort) + require.Equal(t, "low", info.ReasoningEffort) +} diff --git a/service/openaicompat/chat_to_responses.go b/service/openaicompat/chat_to_responses.go index 68d7cc9c..171353ea 100644 --- a/service/openaicompat/chat_to_responses.go +++ b/service/openaicompat/chat_to_responses.go @@ -386,6 +386,8 @@ func ChatCompletionsRequestToResponsesRequest(req *dto.GeneralOpenAIRequest) (*d ParallelToolCalls: parallelToolCallsRaw, Store: req.Store, Metadata: req.Metadata, + EnableThinking: req.EnableThinking, + ThinkingBudget: req.ThinkingBudget, } if req.MaxTokens != nil || req.MaxCompletionTokens != nil { out.MaxOutputTokens = lo.ToPtr(maxOutputTokens) diff --git a/service/openaicompat/responses_compat_test.go b/service/openaicompat/responses_compat_test.go index c9597935..c3bbfb3a 100644 --- a/service/openaicompat/responses_compat_test.go +++ b/service/openaicompat/responses_compat_test.go @@ -44,6 +44,32 @@ func TestResponsesRequestToChatCompletionsRequestInstructionsScalarInputAndZeroT assert.Equal(t, "medium", got.ReasoningEffort) } +func TestChatCompletionsRequestToResponsesRequestPreservesQwenThinkingBudget(t *testing.T) { + got, err := ChatCompletionsRequestToResponsesRequest(&dto.GeneralOpenAIRequest{ + Model: "qwen-plus", + Messages: []dto.Message{{Role: "user", Content: "hello"}}, + EnableThinking: mustCompatRawMessage(t, false), + ThinkingBudget: mustCompatRawMessage(t, 0), + }) + require.NoError(t, err) + + assert.Equal(t, "false", string(got.EnableThinking)) + assert.Equal(t, "0", string(got.ThinkingBudget)) +} + +func TestResponsesRequestToChatCompletionsRequestPreservesQwenThinkingBudget(t *testing.T) { + got, err := ResponsesRequestToChatCompletionsRequest(&dto.OpenAIResponsesRequest{ + Model: "qwen-plus", + Input: mustCompatRawMessage(t, "hello"), + EnableThinking: mustCompatRawMessage(t, true), + ThinkingBudget: mustCompatRawMessage(t, 0), + }) + require.NoError(t, err) + + assert.Equal(t, "true", string(got.EnableThinking)) + assert.Equal(t, "0", string(got.ThinkingBudget)) +} + func TestResponsesRequestToChatCompletionsRequestFunctionCallConversation(t *testing.T) { got, err := ResponsesRequestToChatCompletionsRequest(&dto.OpenAIResponsesRequest{ Model: "gpt-test", diff --git a/service/openaicompat/responses_request_to_chat.go b/service/openaicompat/responses_request_to_chat.go index bb74f880..3c7d91f7 100644 --- a/service/openaicompat/responses_request_to_chat.go +++ b/service/openaicompat/responses_request_to_chat.go @@ -63,6 +63,7 @@ func ResponsesRequestToChatCompletionsRequest(req *dto.OpenAIResponsesRequest) ( SafetyIdentifier: req.SafetyIdentifier, PromptCacheRetention: req.PromptCacheRetention, EnableThinking: req.EnableThinking, + ThinkingBudget: req.ThinkingBudget, } if req.Reasoning != nil { out.ReasoningEffort = req.Reasoning.Effort diff --git a/service/tiered_settle.go b/service/tiered_settle.go index e914dade..24a19dfb 100644 --- a/service/tiered_settle.go +++ b/service/tiered_settle.go @@ -1,9 +1,14 @@ package service import ( + "fmt" + "net/http" + "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" ) // TieredResultWrapper wraps billingexpr.TieredResult for use at the service layer. @@ -102,6 +107,69 @@ func BuildRealtimeTieredTokenParams(usage *dto.RealtimeUsage, usedVars map[strin return BuildTieredTokenParams(normalized, false, usedVars) } +func refreshTieredBillingGroup(relayInfo *relaycommon.RelayInfo) (*billingexpr.BillingSnapshot, error) { + if relayInfo == nil { + return nil, nil + } + snap := relayInfo.TieredBillingSnapshot + if snap == nil || snap.BillingMode != "tiered_expr" { + return nil, nil + } + + groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio + if snap.GroupRatio == groupRatio { + return snap, nil + } + + estimatedQuotaAfterGroup := snap.EstimatedQuotaBeforeGroup * groupRatio + estimatedQuota, err := billingexpr.QuotaRoundStrict(estimatedQuotaAfterGroup) + if err != nil { + return nil, err + } + snap.GroupRatio = groupRatio + snap.EstimatedQuotaAfterGroup = estimatedQuota + return snap, nil +} + +// PrepareTieredBillingForSelectedGroup refreshes routing-dependent tiered +// billing after the final group for this upstream attempt has been selected. +func PrepareTieredBillingForSelectedGroup(c *gin.Context, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { + snap, err := refreshTieredBillingGroup(relayInfo) + if err != nil { + return types.NewErrorWithStatusCode( + err, + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + if snap == nil { + return nil + } + if snap.GroupRatio == 0 { + return nil + } + + relayInfo.PriceData.FreeModel = false + + if relayInfo.Billing == nil { + if c == nil { + return types.NewErrorWithStatusCode( + fmt.Errorf("tiered billing switched from free to paid group without request context"), + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + return PreConsumeBilling(c, snap.EstimatedQuotaAfterGroup, relayInfo) + } + if err := relayInfo.Billing.Reserve(snap.EstimatedQuotaAfterGroup); err != nil { + return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry()) + } + relayInfo.FinalPreConsumedQuota = relayInfo.Billing.GetPreConsumedQuota() + return nil +} + // TryTieredSettle checks if the request uses tiered_expr billing and, if so, // computes the actual quota using the frozen BillingSnapshot. Returns: // - ok=true, quota, result when tiered billing applies diff --git a/service/tiered_settle_test.go b/service/tiered_settle_test.go index 83c4ba47..30ad0b14 100644 --- a/service/tiered_settle_test.go +++ b/service/tiered_settle_test.go @@ -3,12 +3,15 @@ package service import ( "math" "math/rand" + "net/http" "sync" "testing" "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" "github.com/shopspring/decimal" "github.com/stretchr/testify/require" ) @@ -336,6 +339,129 @@ func TestTryTieredSettle_GroupRatioZero(t *testing.T) { } } +func TestPrepareTieredBillingForSelectedGroupUpdatesReservation(t *testing.T) { + const expr = `tier("base", p)` + billing := &recordingBillingSettler{preConsumed: 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.20}, + }, + } + + require.Nil(t, PrepareTieredBillingForSelectedGroup(nil, relayInfo)) + require.Equal(t, []int{100_000}, billing.reserves) + require.Equal(t, 100_000, billing.preConsumed) + require.Equal(t, 100_000, relayInfo.FinalPreConsumedQuota) + require.Equal(t, 0.20, relayInfo.TieredBillingSnapshot.GroupRatio) + require.Equal(t, 100_000, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) +} + +func TestPrepareTieredBillingForSelectedGroupPaidToFreeDoesNotReserve(t *testing.T) { + const expr = `tier("base", p)` + billing := &recordingBillingSettler{preConsumed: 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)) + require.Empty(t, billing.reserves) + require.Equal(t, 50_000, relayInfo.FinalPreConsumedQuota) + require.Equal(t, 0.0, relayInfo.TieredBillingSnapshot.GroupRatio) + require.Equal(t, 0, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) +} + +func TestPrepareTieredBillingForSelectedGroupFreeToPaidInitializesBilling(t *testing.T) { + truncate(t) + const ( + userID = 921 + tokenID = 922 + tokenKey = "tiered-free-to-paid-token" + initialQuota = 200_000 + expectedQuota = 100_000 + ) + seedUser(t, userID, initialQuota) + seedToken(t, tokenID, userID, tokenKey, initialQuota) + + ctx, _ := gin.CreateTestContext(nil) + relayInfo := &relaycommon.RelayInfo{ + RequestId: "tiered-free-to-paid-request", + UserId: userID, + TokenId: tokenID, + TokenKey: tokenKey, + OriginModelName: "tiered-free-to-paid-model", + UsingGroup: "paid", + UserSetting: dto.UserSetting{ + BillingPreference: "wallet_only", + }, + PriceData: types.PriceData{ + FreeModel: true, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.20}, + }, + TieredBillingSnapshot: &billingexpr.BillingSnapshot{ + BillingMode: "tiered_expr", + ExprString: `tier("base", p)`, + ExprHash: billingexpr.ExprHashString(`tier("base", p)`), + GroupRatio: 0, + EstimatedQuotaBeforeGroup: 500_000, + EstimatedQuotaAfterGroup: 0, + QuotaPerUnit: testQuotaPerUnit, + }, + } + + require.Nil(t, PrepareTieredBillingForSelectedGroup(ctx, relayInfo)) + require.NotNil(t, relayInfo.Billing) + require.Equal(t, expectedQuota, relayInfo.FinalPreConsumedQuota) + require.Equal(t, expectedQuota, relayInfo.Billing.GetPreConsumedQuota()) + require.False(t, relayInfo.PriceData.FreeModel) + require.Equal(t, 0.20, relayInfo.TieredBillingSnapshot.GroupRatio) + require.Equal(t, expectedQuota, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) + require.EqualValues(t, initialQuota-expectedQuota, getUserQuota(t, userID)) + require.Equal(t, initialQuota-expectedQuota, getTokenRemainQuota(t, tokenID)) +} + +func TestPrepareTieredBillingForSelectedGroupFreeToPaidRequiresContext(t *testing.T) { + const expr = `tier("base", p)` + relayInfo := makeRelayInfo(expr, 0, 1_000_000, 0) + relayInfo.PriceData = types.PriceData{ + FreeModel: true, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.20}, + } + + err := PrepareTieredBillingForSelectedGroup(nil, relayInfo) + + require.NotNil(t, err) + require.Equal(t, types.ErrorCodeModelPriceError, err.GetErrorCode()) + require.Equal(t, http.StatusBadRequest, err.StatusCode) + require.True(t, types.IsSkipRetryError(err)) + require.ErrorContains(t, err, "without request context") + require.Nil(t, relayInfo.Billing) + require.Zero(t, relayInfo.FinalPreConsumedQuota) +} + // --------------------------------------------------------------------------- // Ratio mode (negative tests) — TryTieredSettle must return false // --------------------------------------------------------------------------- diff --git a/web/default/src/features/auth/lib/oauth-callback-mode.test.ts b/web/default/src/features/auth/lib/oauth-callback-mode.test.ts new file mode 100644 index 00000000..3956ced4 --- /dev/null +++ b/web/default/src/features/auth/lib/oauth-callback-mode.test.ts @@ -0,0 +1,149 @@ +import { defaultParseSearch } from '@tanstack/react-router' +import assert from 'node:assert/strict' +import { afterEach, describe, test } from 'node:test' +import { + oauthCallbackSearchSchema, + resolveOAuthCallbackMode, +} from './oauth-callback-mode' + +const originalLocalStorageDescriptor = Object.getOwnPropertyDescriptor( + globalThis, + 'localStorage' +) + +function setLocalStorage(value: Storage): void { + Object.defineProperty(globalThis, 'localStorage', { + configurable: true, + value, + }) +} + +function fakeStorage(initial: Record = {}): Storage { + const data = new Map(Object.entries(initial)) + return { + get length() { + return data.size + }, + clear: () => data.clear(), + getItem: (key: string) => data.get(key) ?? null, + key: (index: number) => Array.from(data.keys())[index] ?? null, + removeItem: (key: string) => void data.delete(key), + setItem: (key: string, value: string) => void data.set(key, value), + } +} + +afterEach(() => { + if (originalLocalStorageDescriptor) { + Object.defineProperty( + globalThis, + 'localStorage', + originalLocalStorageDescriptor + ) + } else { + Reflect.deleteProperty(globalThis, 'localStorage') + } +}) + +describe('resolveOAuthCallbackMode', () => { + test('matching bind state is treated as a bind flow', () => { + const state = 'bind-state' + setLocalStorage( + fakeStorage({ + [`oauth:binding:state:github:${state}`]: String(Date.now()), + }) + ) + + assert.equal(resolveOAuthCallbackMode('github', state), 'bind') + }) + + test('callback without a local bind marker stays a login flow', () => { + setLocalStorage(fakeStorage()) + + assert.equal( + resolveOAuthCallbackMode('github', 'foreign-opener-state'), + 'login' + ) + }) + + test('bind marker for another provider stays a login flow', () => { + const state = 'bind-state' + setLocalStorage( + fakeStorage({ + [`oauth:binding:state:github:${state}`]: String(Date.now()), + }) + ) + + assert.equal(resolveOAuthCallbackMode('oidc', state), 'login') + }) + + test('explicit bind hint is treated as a bind flow', () => { + setLocalStorage(fakeStorage()) + + assert.equal(resolveOAuthCallbackMode('custom', undefined, true), 'bind') + }) + + test('stale bind marker falls back to login', () => { + const state = 'stale-state' + setLocalStorage( + fakeStorage({ + [`oauth:binding:state:github:${state}`]: String( + Date.now() - 11 * 60 * 1000 + ), + }) + ) + + assert.equal(resolveOAuthCallbackMode('github', state), 'login') + }) + + test('storage access failure falls back to login', () => { + Object.defineProperty(globalThis, 'localStorage', { + configurable: true, + get() { + throw new Error('blocked') + }, + }) + + assert.equal(resolveOAuthCallbackMode('github', 'bind-state'), 'login') + }) +}) + +describe('oauthCallbackSearchSchema', () => { + test('preserves supported callback parameters', () => { + assert.deepEqual( + oauthCallbackSearchSchema.parse( + defaultParseSearch( + '?code=authorization-code&state=callback-state&redirect=/dashboard&bind=true' + ) + ), + { + code: 'authorization-code', + state: 'callback-state', + redirect: '/dashboard', + bind: true, + } + ) + }) + + test('normalizes an explicit false bind hint', () => { + assert.deepEqual(oauthCallbackSearchSchema.parse({ bind: 'false' }), { + bind: false, + }) + }) + + test('drops unsupported callback parameter shapes', () => { + assert.deepEqual( + oauthCallbackSearchSchema.parse({ + code: ['authorization-code'], + state: { value: 'callback-state' }, + redirect: true, + bind: 'yes', + }), + { + code: undefined, + state: undefined, + redirect: undefined, + bind: undefined, + } + ) + }) +}) diff --git a/web/default/src/features/auth/lib/oauth-callback-mode.ts b/web/default/src/features/auth/lib/oauth-callback-mode.ts new file mode 100644 index 00000000..b8bce5fb --- /dev/null +++ b/web/default/src/features/auth/lib/oauth-callback-mode.ts @@ -0,0 +1,23 @@ +import { z } from 'zod' +import { isOAuthBindingState } from '@/lib/oauth' + +export type OAuthCallbackMode = 'login' | 'bind' + +const oauthBindSearchSchema = z + .union([z.boolean(), z.enum(['true', 'false'])]) + .transform((value) => value === true || value === 'true') + +export const oauthCallbackSearchSchema = z.object({ + code: z.string().optional().catch(undefined), + state: z.string().optional().catch(undefined), + redirect: z.string().optional().catch(undefined), + bind: oauthBindSearchSchema.optional().catch(undefined), +}) + +export function resolveOAuthCallbackMode( + provider?: string, + state?: string, + bindHint = false +): OAuthCallbackMode { + return bindHint || isOAuthBindingState(provider, state) ? 'bind' : 'login' +} diff --git a/web/default/src/features/channels/components/dialogs/status-code-risk-dialog.test.tsx b/web/default/src/features/channels/components/dialogs/status-code-risk-dialog.test.tsx new file mode 100644 index 00000000..e9e6d868 --- /dev/null +++ b/web/default/src/features/channels/components/dialogs/status-code-risk-dialog.test.tsx @@ -0,0 +1,82 @@ +/* +Copyright (C) 2023-2026 MAX-API-Next + +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 https://github.com/MAX-API-Next/MAX-API/issues +*/ +import { createReactTestEnvironment } from '@/test/react' +import assert from 'node:assert/strict' +import { after, before, describe, test } from 'node:test' +import { StatusCodeRiskConfirmationContent } from './status-code-risk-dialog' + +const testEnv = createReactTestEnvironment() + +before(() => testEnv.setup()) + +after(() => testEnv.teardown()) + +describe('StatusCodeRiskDialog', () => { + test('keeps risk notices and invokes both available actions', async () => { + let cancelCount = 0 + let confirmCount = 0 + const view = await testEnv.render( + 500']} + onCancel={() => { + cancelCount += 1 + }} + onConfirm={() => { + confirmCount += 1 + }} + /> + ) + + try { + assert.match(view.container.textContent || '', /200 -> 500/) + for (let index = 1; index <= 4; index += 1) { + assert.match( + view.container.textContent || '', + new RegExp(`High-risk status code retry risk check ${index}`) + ) + } + + assert.equal(view.container.querySelector('input'), null) + assert.equal(view.container.querySelector('[type="checkbox"]'), null) + + const buttons = Array.from(view.container.querySelectorAll('button')) + const cancelButton = buttons.find( + (button) => button.textContent === 'Cancel' + ) + const confirmButton = buttons.find( + (button) => button.textContent === 'I confirm enabling high-risk retry' + ) + + assert.ok(cancelButton) + assert.ok(confirmButton) + assert.equal(confirmButton.disabled, false) + + await view.click(cancelButton) + assert.equal(cancelCount, 1) + assert.equal(confirmCount, 0) + + await view.click(confirmButton) + + assert.equal(cancelCount, 1) + assert.equal(confirmCount, 1) + } finally { + await view.unmount() + } + }) +}) diff --git a/web/default/src/features/channels/components/dialogs/status-code-risk-dialog.tsx b/web/default/src/features/channels/components/dialogs/status-code-risk-dialog.tsx index 0470ab36..d716f993 100644 --- a/web/default/src/features/channels/components/dialogs/status-code-risk-dialog.tsx +++ b/web/default/src/features/channels/components/dialogs/status-code-risk-dialog.tsx @@ -16,11 +16,11 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useState } from 'react' -import { AlertTriangle } from 'lucide-react' +import type { ReactElement } from 'react' +import { Alert02Icon } from '@hugeicons/core-free-icons' +import { HugeiconsIcon } from '@hugeicons/react' import { useTranslation } from 'react-i18next' import { Button } from '@/components/ui/button' -import { Checkbox } from '@/components/ui/checkbox' import { Dialog, DialogContent, @@ -29,8 +29,6 @@ import { DialogHeader, DialogTitle, } from '@/components/ui/dialog' -import { Input } from '@/components/ui/input' -import { Label } from '@/components/ui/label' interface StatusCodeRiskDialogProps { open: boolean @@ -39,56 +37,77 @@ interface StatusCodeRiskDialogProps { onConfirm: () => void } -const CHECKLIST_KEYS = [ +const RISK_NOTICE_KEYS = [ 'High-risk status code retry risk check 1', 'High-risk status code retry risk check 2', 'High-risk status code retry risk check 3', 'High-risk status code retry risk check 4', ] as const -export function StatusCodeRiskDialog({ - open, - onOpenChange, - detailItems, - onConfirm, -}: StatusCodeRiskDialogProps) { - const { t } = useTranslation() - const [checkedItems, setCheckedItems] = useState>(new Set()) - const [confirmText, setConfirmText] = useState('') - - const requiredText = t('High-risk status code retry confirmation text') - const allChecked = checkedItems.size === CHECKLIST_KEYS.length - const textMatches = confirmText.trim() === requiredText.trim() - const canConfirm = allChecked && textMatches +interface StatusCodeRiskConfirmationContentProps { + detailItems: string[] + onCancel: () => void + onConfirm: () => void +} - const handleConfirm = () => { - if (!canConfirm) return - setCheckedItems(new Set()) - setConfirmText('') - onConfirm() - } +export function StatusCodeRiskConfirmationContent( + props: StatusCodeRiskConfirmationContentProps +): ReactElement { + const { t } = useTranslation() - const handleCancel = () => { - setCheckedItems(new Set()) - setConfirmText('') - onOpenChange(false) - } + return ( + <> +
+ {props.detailItems.length > 0 && ( +
+

+ {t('Detected high-risk status code redirect rules')} +

+
    + {props.detailItems.map((item) => ( +
  • + {item} +
    • + ))} +
+
+ )} + +
    + {RISK_NOTICE_KEYS.map((key) => ( +
  • {t(key)}
    • + ))} +
+
+ + + + + + + ) +} - const toggleCheck = (idx: number) => { - setCheckedItems((prev) => { - const next = new Set(prev) - if (next.has(idx)) next.delete(idx) - else next.add(idx) - return next - }) - } +export function StatusCodeRiskDialog( + props: StatusCodeRiskDialogProps +): ReactElement { + const { t } = useTranslation() return ( - + - + @@ -96,72 +115,11 @@ export function StatusCodeRiskDialog({ -
- {detailItems.length > 0 && ( -
-

- {t('Detected high-risk status code redirect rules')} -

-
    - {detailItems.map((item) => ( -
  • - {item} -
    • - ))} -
-
- )} - -
- {CHECKLIST_KEYS.map((key, idx) => ( -
- toggleCheck(idx)} - /> - -
- ))} -
- -
- - setConfirmText(e.target.value)} - placeholder={t('High-risk status code retry input placeholder')} - /> - {confirmText && !textMatches && ( -

- {t('High-risk status code retry input mismatch')} -

- )} -
-
- - - - - + props.onOpenChange(false)} + onConfirm={props.onConfirm} + /> ) diff --git a/web/default/src/lib/oauth.ts b/web/default/src/lib/oauth.ts index a5a6bcaa..3c2de139 100644 --- a/web/default/src/lib/oauth.ts +++ b/web/default/src/lib/oauth.ts @@ -21,22 +21,33 @@ import { api } from './api' const oauthBindingStatePrefix = 'oauth:binding:state:' const oauthBindingStateTtlMs = 10 * 60 * 1000 -function getOAuthBindingStateKey(state: string): string { +function normalizeOAuthProvider(provider: string): string { + return provider.trim().toLowerCase() +} + +function getOAuthBindingStateKey(provider: string, state: string): string { + return `${oauthBindingStatePrefix}${normalizeOAuthProvider(provider)}:${state}` +} + +function getLegacyOAuthBindingStateKey(state: string): string { return `${oauthBindingStatePrefix}${state}` } -function markOAuthBindingState(state: string): void { +function markOAuthBindingState(provider: string, state: string): void { try { - localStorage.setItem(getOAuthBindingStateKey(state), String(Date.now())) + localStorage.setItem( + getOAuthBindingStateKey(provider, state), + String(Date.now()) + ) } catch { - // The callback retains the legacy opener check when storage is unavailable. + // Without this marker the callback safely falls back to login mode. } } -export function isOAuthBindingState(state?: string): boolean { - if (!state) return false +export function isOAuthBindingState(provider?: string, state?: string): boolean { + if (!provider || !state) return false try { - const key = getOAuthBindingStateKey(state) + const key = getOAuthBindingStateKey(provider, state) const createdAt = Number(localStorage.getItem(key)) if ( !Number.isFinite(createdAt) || @@ -51,17 +62,27 @@ export function isOAuthBindingState(state?: string): boolean { } } -export function clearOAuthBindingState(state?: string): void { +export function clearOAuthBindingState( + provider?: string, + state?: string +): void { if (!state) return try { - localStorage.removeItem(getOAuthBindingStateKey(state)) + if (provider) { + localStorage.removeItem(getOAuthBindingStateKey(provider, state)) + } + localStorage.removeItem(getLegacyOAuthBindingStateKey(state)) } catch { // Ignore storage cleanup failures; entries are bounded by their TTL. } } -function openOAuthBindingWindow(url: string, state: string): void { - markOAuthBindingState(state) +function openOAuthBindingWindow( + url: string, + provider: string, + state: string +): void { + markOAuthBindingState(provider, state) window.open(url, '_blank', 'noopener,noreferrer') } @@ -151,7 +172,7 @@ export async function handleGitHubOAuth(clientId: string): Promise { if (!state) return const url = buildGitHubOAuthUrl(clientId, state) - openOAuthBindingWindow(url, state) + openOAuthBindingWindow(url, 'github', state) } /** @@ -162,7 +183,7 @@ export async function handleDiscordOAuth(clientId: string): Promise { if (!state) return const url = buildDiscordOAuthUrl(clientId, state) - openOAuthBindingWindow(url, state) + openOAuthBindingWindow(url, 'discord', state) } /** @@ -176,7 +197,7 @@ export async function handleOIDCOAuth( if (!state) return const url = buildOIDCOAuthUrl(authUrl, clientId, state) - openOAuthBindingWindow(url, state) + openOAuthBindingWindow(url, 'oidc', state) } /** @@ -187,5 +208,5 @@ export async function handleLinuxDOOAuth(clientId: string): Promise { if (!state) return const url = buildLinuxDOOAuthUrl(clientId, state) - openOAuthBindingWindow(url, state) + openOAuthBindingWindow(url, 'linuxdo', state) } diff --git a/web/default/src/routes/oauth/$provider.tsx b/web/default/src/routes/oauth/$provider.tsx index 7e0ebf3f..71f0b973 100644 --- a/web/default/src/routes/oauth/$provider.tsx +++ b/web/default/src/routes/oauth/$provider.tsx @@ -16,49 +16,37 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useEffect, useState } from 'react' +import { useEffect, useRef } from 'react' import type { AxiosRequestConfig } from 'axios' -import { - createFileRoute, - useNavigate, - useParams, - useSearch, -} from '@tanstack/react-router' +import { createFileRoute, useNavigate } from '@tanstack/react-router' import i18next from 'i18next' import { toast } from 'sonner' import { useAuthStore, type AuthUser } from '@/stores/auth-store' import { api, getSelf } from '@/lib/api' -import { clearOAuthBindingState, isOAuthBindingState } from '@/lib/oauth' +import { clearOAuthBindingState } from '@/lib/oauth' import { normalizeInternalRedirect } from '@/lib/safe-redirect' import { OAuthCallbackScreen } from '@/features/auth/components/oauth-callback-screen' import { OAUTH_BIND_STORAGE_KEY } from '@/features/auth/constants' +import { + oauthCallbackSearchSchema, + resolveOAuthCallbackMode, +} from '@/features/auth/lib/oauth-callback-mode' type OAuthRequestConfig = AxiosRequestConfig & { skipBusinessError?: boolean } +export const Route = createFileRoute('/oauth/$provider')({ + component: OAuthCallback, + validateSearch: oauthCallbackSearchSchema, +}) + function OAuthCallback() { const navigate = useNavigate() - const { provider } = useParams({ from: '/oauth/$provider' }) as { - provider: string - } - const search = useSearch({ from: '/oauth/$provider' }) as { - code?: string - state?: string - redirect?: string - } - const [mode, setMode] = useState<'login' | 'bind'>(() => { - if (typeof window === 'undefined') return 'login' - return window.opener || isOAuthBindingState(search.state) ? 'bind' : 'login' - }) - - useEffect(() => { - if (typeof window === 'undefined') return - // eslint-disable-next-line react-hooks/set-state-in-effect - setMode( - window.opener || isOAuthBindingState(search.state) ? 'bind' : 'login' - ) - }, [search.state]) + const { provider } = Route.useParams() + const search = Route.useSearch() + const mode = resolveOAuthCallbackMode(provider, search.state, search.bind) + const exchangeAttemptKeysRef = useRef(new Set()) useEffect(() => { ;(async () => { @@ -85,18 +73,20 @@ function OAuthCallback() { safeNavigate('/sign-in') return } + + const exchangeKey = JSON.stringify([ + provider, + search.code, + search.state ?? null, + ]) + if (exchangeAttemptKeysRef.current.has(exchangeKey)) return + exchangeAttemptKeysRef.current.add(exchangeKey) + const isBindingFlow = - typeof window !== 'undefined' - ? Boolean(window.opener) || isOAuthBindingState(search.state) - : mode === 'bind' - if (isBindingFlow && mode !== 'bind') { - setMode('bind') - } else if (!isBindingFlow && mode !== 'login') { - setMode('login') - } + resolveOAuthCallbackMode(provider, search.state, search.bind) === 'bind' const notifyBindingResult = (status: 'success' | 'error') => { if (typeof window === 'undefined') return - clearOAuthBindingState(search.state) + clearOAuthBindingState(provider, search.state) try { window.localStorage.setItem( OAUTH_BIND_STORAGE_KEY, @@ -243,11 +233,14 @@ function OAuthCallback() { return } })() - }, [mode, navigate, provider, search]) + }, [ + navigate, + provider, + search.bind, + search.code, + search.redirect, + search.state, + ]) return } - -export const Route = createFileRoute('/oauth/$provider')({ - component: OAuthCallback, -})