From 7abd87fba0b1656e7923cf624f167f514af14efa Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sun, 12 Jul 2026 18:16:55 +0800 Subject: [PATCH 1/9] v1.0.4 --- common/quota_math.go | 138 ++++++- common/quota_math_test.go | 129 +++++- dto/billing_usage.go | 215 ++++++++++ dto/billing_usage_test.go | 89 ++++ dto/claude.go | 1 + dto/gemini.go | 38 +- dto/openai_response.go | 25 +- model/ability.go | 6 +- pkg/billingexpr/round.go | 22 +- pkg/billingexpr/settle.go | 5 +- pkg/billingexpr/types.go | 5 + relay/channel/claude/relay-claude.go | 6 + relay/channel/gemini/relay-gemini-native.go | 4 +- relay/channel/gemini/relay-gemini.go | 97 ++++- relay/channel/gemini/relay_responses.go | 4 +- relay/channel/openai/relay-openai.go | 1 + relay/channel/openai/relay_image.go | 1 + relay/channel/openai/relay_responses.go | 2 + .../channel/openai/relay_responses_compact.go | 1 + relay/common/relay_info.go | 4 + relay/helper/price.go | 29 +- relay/relay_task.go | 17 +- service/billing.go | 17 + service/billing_usage.go | 205 +++++++++ service/convert.go | 48 ++- service/log_info_generate.go | 24 ++ .../chat_to_responses_response.go | 5 + service/openaicompat/responses_to_chat.go | 5 + service/pre_consume_quota.go | 6 + service/quota.go | 17 +- service/task_billing.go | 10 +- service/text_quota.go | 49 ++- service/text_quota_test.go | 166 ++++++++ service/tiered_settle.go | 4 +- .../billing/section-registry.tsx | 2 +- .../models/model-ratio-form.tsx | 45 +- .../models/model-ratio-visual-editor.tsx | 390 ++++++++++++------ .../models/ratio-settings-card.tsx | 25 +- web/default/src/i18n/locales/en.json | 3 + web/default/src/i18n/locales/fr.json | 3 + web/default/src/i18n/locales/ja.json | 3 + web/default/src/i18n/locales/ru.json | 3 + web/default/src/i18n/locales/vi.json | 3 + web/default/src/i18n/locales/zh.json | 3 + 44 files changed, 1651 insertions(+), 224 deletions(-) create mode 100644 dto/billing_usage.go create mode 100644 dto/billing_usage_test.go create mode 100644 service/billing_usage.go diff --git a/common/quota_math.go b/common/quota_math.go index 9d71d8d..c7c05ec 100644 --- a/common/quota_math.go +++ b/common/quota_math.go @@ -1,21 +1,131 @@ package common -import "math" +import ( + "fmt" + "math" -// QuotaFromFloat converts a computed quota value to int with saturation. -// Quota products can include user-controlled multipliers such as image count, -// video seconds, or resolution ratios; oversized products must never wrap into -// a negative charge. The bound is int32 because quota columns are int fields -// used as 32-bit database integers in supported deployments. -func QuotaFromFloat(value float64) int { - if math.IsNaN(value) { - return 0 + "github.com/shopspring/decimal" +) + +// Quota conversions are centralized here so every billing path shares one +// saturation + logging policy. Quota columns are 32-bit integers in supported +// deployments, so an oversized product must clamp to the int32 range instead +// of wrapping around and turning a charge into a credit. +const ( + MaxQuota = math.MaxInt32 + MinQuota = math.MinInt32 +) + +// QuotaClampKind identifies why a quota conversion had to be saturated. +type QuotaClampKind string + +const ( + QuotaClampOverflow QuotaClampKind = "overflow" + QuotaClampUnderflow QuotaClampKind = "underflow" + QuotaClampNaN QuotaClampKind = "nan" +) + +// QuotaClamp describes a single saturation event and is stored under +// admin_info.quota_saturation for audit when a billing path has to clamp. +type QuotaClamp struct { + Op string `json:"op"` + Kind QuotaClampKind `json:"kind"` + Original float64 `json:"original"` + Clamped int `json:"clamped"` +} + +func (c *QuotaClamp) Error() string { + if c == nil { + return "" } - if value >= math.MaxInt32 { - return math.MaxInt32 + return fmt.Sprintf("quota conversion (%s) %s: original=%g, clamped=%d", c.Op, c.Kind, c.Original, c.Clamped) +} + +func (c *QuotaClamp) AuditMap() map[string]interface{} { + if c == nil { + return nil } - if value <= math.MinInt32 { - return math.MinInt32 + original := interface{}(c.Original) + switch { + case math.IsNaN(c.Original): + original = "NaN" + case math.IsInf(c.Original, 1): + original = "+Inf" + case math.IsInf(c.Original, -1): + original = "-Inf" } - return int(value) + return map[string]interface{}{ + "op": c.Op, + "kind": c.Kind, + "original": original, + "clamped": c.Clamped, + } +} + +func saturateQuota(value float64, op string) (int, *QuotaClamp) { + var clamp *QuotaClamp + switch { + case math.IsNaN(value): + clamp = &QuotaClamp{Op: op, Kind: QuotaClampNaN, Original: value, Clamped: 0} + case value >= MaxQuota: + clamp = &QuotaClamp{Op: op, Kind: QuotaClampOverflow, Original: value, Clamped: MaxQuota} + case value <= MinQuota: + clamp = &QuotaClamp{Op: op, Kind: QuotaClampUnderflow, Original: value, Clamped: MinQuota} + default: + return int(value), nil + } + SysError(clamp.Error()) + return clamp.Clamped, clamp +} + +func strictQuota(quota int, clamp *QuotaClamp) (int, error) { + if clamp != nil { + return 0, clamp + } + return quota, nil +} + +// QuotaFromFloat converts a computed quota value to int, truncating toward +// zero, with int32 saturation. +func QuotaFromFloat(value float64) int { + quota, _ := QuotaFromFloatChecked(value) + return quota +} + +// QuotaFromFloatChecked is QuotaFromFloat plus a non-nil clamp descriptor +// when saturation or NaN fallback happened. +func QuotaFromFloatChecked(value float64) (int, *QuotaClamp) { + return saturateQuota(value, "QuotaFromFloat") +} + +// QuotaFromFloatStrict rejects unrepresentable billing estimates instead of +// silently saturating them. +func QuotaFromFloatStrict(value float64) (int, error) { + return strictQuota(QuotaFromFloatChecked(value)) +} + +// QuotaRound converts a float64 quota value to int using half-away-from-zero +// rounding, with the same saturation policy. +func QuotaRound(value float64) int { + quota, _ := QuotaRoundChecked(value) + return quota +} + +func QuotaRoundChecked(value float64) (int, *QuotaClamp) { + return saturateQuota(math.Round(value), "QuotaRound") +} + +func QuotaRoundStrict(value float64) (int, error) { + return strictQuota(QuotaRoundChecked(value)) +} + +// QuotaFromDecimal rounds a computed quota decimal before conversion. +func QuotaFromDecimal(d decimal.Decimal) int { + quota, _ := QuotaFromDecimalChecked(d) + return quota +} + +func QuotaFromDecimalChecked(d decimal.Decimal) (int, *QuotaClamp) { + f, _ := d.Round(0).Float64() + return saturateQuota(f, "QuotaFromDecimal") } diff --git a/common/quota_math_test.go b/common/quota_math_test.go index 11ecdf5..0e5f538 100644 --- a/common/quota_math_test.go +++ b/common/quota_math_test.go @@ -4,15 +4,134 @@ import ( "math" "testing" + "github.com/shopspring/decimal" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) +const overflowingProduct = 2000 * 1.8446744073686647e19 + func TestQuotaFromFloat(t *testing.T) { assert.Equal(t, 42, QuotaFromFloat(42.4)) - assert.Equal(t, -42, QuotaFromFloat(-42.4)) - assert.Equal(t, math.MaxInt32, QuotaFromFloat(2000*1.8446744073686647e19)) - assert.Equal(t, math.MinInt32, QuotaFromFloat(-2000*1.8446744073686647e19)) - assert.Equal(t, math.MaxInt32, QuotaFromFloat(math.Inf(1))) - assert.Equal(t, math.MinInt32, QuotaFromFloat(math.Inf(-1))) + assert.Equal(t, 42, QuotaFromFloat(42.9)) + assert.Equal(t, -42, QuotaFromFloat(-42.9)) + assert.Equal(t, MaxQuota, QuotaFromFloat(overflowingProduct)) + assert.Equal(t, MinQuota, QuotaFromFloat(-overflowingProduct)) + assert.Equal(t, MaxQuota, QuotaFromFloat(math.Inf(1))) + assert.Equal(t, MinQuota, QuotaFromFloat(math.Inf(-1))) assert.Equal(t, 0, QuotaFromFloat(math.NaN())) } + +func TestQuotaRound(t *testing.T) { + assert.Equal(t, 42, QuotaRound(41.5)) + assert.Equal(t, 43, QuotaRound(42.5)) + assert.Equal(t, -43, QuotaRound(-42.5)) + assert.Equal(t, MaxQuota, QuotaRound(overflowingProduct)) + assert.Equal(t, MinQuota, QuotaRound(-overflowingProduct)) + assert.Equal(t, 0, QuotaRound(math.NaN())) +} + +func TestQuotaFromDecimal(t *testing.T) { + assert.Equal(t, 43, QuotaFromDecimal(decimal.NewFromFloat(42.5))) + assert.Equal(t, 42, QuotaFromDecimal(decimal.NewFromFloat(41.7))) + assert.Equal(t, MaxQuota, QuotaFromDecimal(decimal.NewFromInt(2000).Mul(decimal.NewFromFloat(1.8446744073686647e19)))) + assert.Equal(t, MinQuota, QuotaFromDecimal(decimal.NewFromInt(-2000).Mul(decimal.NewFromFloat(1.8446744073686647e19)))) +} + +func TestQuotaFromFloatChecked(t *testing.T) { + quota, clamp := QuotaFromFloatChecked(42.9) + assert.Equal(t, 42, quota) + assert.Nil(t, clamp) + + quota, clamp = QuotaFromFloatChecked(overflowingProduct) + assert.Equal(t, MaxQuota, quota) + if assert.NotNil(t, clamp) { + assert.Equal(t, "QuotaFromFloat", clamp.Op) + assert.Equal(t, QuotaClampOverflow, clamp.Kind) + assert.Equal(t, MaxQuota, clamp.Clamped) + } + + quota, clamp = QuotaFromFloatChecked(-overflowingProduct) + assert.Equal(t, MinQuota, quota) + if assert.NotNil(t, clamp) { + assert.Equal(t, QuotaClampUnderflow, clamp.Kind) + assert.Equal(t, MinQuota, clamp.Clamped) + } + + quota, clamp = QuotaFromFloatChecked(math.NaN()) + assert.Equal(t, 0, quota) + if assert.NotNil(t, clamp) { + assert.Equal(t, QuotaClampNaN, clamp.Kind) + assert.Equal(t, 0, clamp.Clamped) + } +} + +func TestQuotaFromFloatStrictReturnsTypedClampError(t *testing.T) { + quota, err := QuotaFromFloatStrict(42.9) + require.NoError(t, err) + assert.Equal(t, 42, quota) + + quota, err = QuotaFromFloatStrict(overflowingProduct) + assert.Zero(t, quota) + var clamp *QuotaClamp + require.ErrorAs(t, err, &clamp) + assert.Equal(t, QuotaClampOverflow, clamp.Kind) + assert.Equal(t, MaxQuota, clamp.Clamped) + assert.ErrorContains(t, err, "QuotaFromFloat") + assert.ErrorContains(t, err, "overflow") + assert.ErrorContains(t, err, "original=") + assert.ErrorContains(t, err, "clamped=2147483647") +} + +func TestQuotaClampAuditMapIsJSONSafe(t *testing.T) { + tests := []struct { + name string + value float64 + wantOriginal interface{} + }{ + {name: "nan", value: math.NaN(), wantOriginal: "NaN"}, + {name: "positive infinity", value: math.Inf(1), wantOriginal: "+Inf"}, + {name: "negative infinity", value: math.Inf(-1), wantOriginal: "-Inf"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, clamp := QuotaFromFloatChecked(tt.value) + require.NotNil(t, clamp) + + auditMap := clamp.AuditMap() + assert.Equal(t, tt.wantOriginal, auditMap["original"]) + assert.NotEmpty(t, MapToJsonStr(map[string]interface{}{ + "admin_info": map[string]interface{}{ + "quota_saturation": auditMap, + }, + })) + }) + } +} + +func TestQuotaRoundChecked(t *testing.T) { + quota, clamp := QuotaRoundChecked(42.5) + assert.Equal(t, 43, quota) + assert.Nil(t, clamp) + + quota, clamp = QuotaRoundChecked(overflowingProduct) + assert.Equal(t, MaxQuota, quota) + if assert.NotNil(t, clamp) { + assert.Equal(t, "QuotaRound", clamp.Op) + assert.Equal(t, QuotaClampOverflow, clamp.Kind) + } +} + +func TestQuotaFromDecimalChecked(t *testing.T) { + quota, clamp := QuotaFromDecimalChecked(decimal.NewFromFloat(41.7)) + assert.Equal(t, 42, quota) + assert.Nil(t, clamp) + + quota, clamp = QuotaFromDecimalChecked(decimal.NewFromInt(2000).Mul(decimal.NewFromFloat(1.8446744073686647e19))) + assert.Equal(t, MaxQuota, quota) + if assert.NotNil(t, clamp) { + assert.Equal(t, "QuotaFromDecimal", clamp.Op) + assert.Equal(t, QuotaClampOverflow, clamp.Kind) + } +} diff --git a/dto/billing_usage.go b/dto/billing_usage.go new file mode 100644 index 0000000..6773d4d --- /dev/null +++ b/dto/billing_usage.go @@ -0,0 +1,215 @@ +package dto + +const ( + BillingUsageSourceClaudeMessages = "claude_messages" + BillingUsageSourceGeminiChat = "gemini_chat" + BillingUsageSourceOAIChat = "oai_chat" + BillingUsageSourceOAIResponses = "oai_responses" + + BillingUsageSemanticAnthropic = "anthropic" + BillingUsageSemanticGemini = "gemini" + BillingUsageSemanticOpenAI = "openai" +) + +type BillingUsage struct { + Source string `json:"source,omitempty"` + Semantic string `json:"semantic,omitempty"` + Estimated bool `json:"estimated,omitempty"` + OpenAIUsage *Usage `json:"openai_usage,omitempty"` + ClaudeUsage *ClaudeUsage `json:"claude_usage,omitempty"` + GeminiUsageMetadata *GeminiUsageMetadata `json:"gemini_usage_metadata,omitempty"` +} + +func NewClaudeMessagesBillingUsage(usage *ClaudeUsage) *BillingUsage { + if !HasClaudeUsageTokens(usage) { + return nil + } + return &BillingUsage{ + Source: BillingUsageSourceClaudeMessages, + Semantic: BillingUsageSemanticAnthropic, + ClaudeUsage: cloneClaudeUsage(usage), + } +} + +func HasClaudeUsageTokens(usage *ClaudeUsage) bool { + if usage == nil { + return false + } + if usage.InputTokens != 0 || + usage.OutputTokens != 0 || + usage.CacheCreationInputTokens != 0 || + usage.CacheReadInputTokens != 0 || + usage.ClaudeCacheCreation5mTokens != 0 || + usage.ClaudeCacheCreation1hTokens != 0 { + return true + } + if usage.CacheCreation != nil && + (usage.CacheCreation.Ephemeral5mInputTokens != 0 || usage.CacheCreation.Ephemeral1hInputTokens != 0) { + return true + } + return false +} + +func NewOpenAIChatBillingUsage(usage *Usage) *BillingUsage { + return newOpenAIBillingUsage(BillingUsageSourceOAIChat, usage) +} + +func NewOpenAIResponsesBillingUsage(usage *Usage) *BillingUsage { + return newOpenAIBillingUsage(BillingUsageSourceOAIResponses, usage) +} + +func newOpenAIBillingUsage(source string, usage *Usage) *BillingUsage { + if !HasOpenAIUsageTokens(usage) { + return nil + } + return &BillingUsage{ + Source: source, + Semantic: BillingUsageSemanticOpenAI, + OpenAIUsage: cloneOpenAIUsage(usage), + } +} + +func HasOpenAIUsageTokens(usage *Usage) bool { + if usage == nil { + return false + } + if usage.PromptTokens != 0 || + usage.CompletionTokens != 0 || + usage.TotalTokens != 0 || + usage.InputTokens != 0 || + usage.OutputTokens != 0 || + usage.PromptCacheHitTokens != 0 || + usage.ClaudeCacheCreation5mTokens != 0 || + usage.ClaudeCacheCreation1hTokens != 0 { + return true + } + if usage.PromptTokensDetails.CachedTokens != 0 || + usage.PromptTokensDetails.CachedCreationTokens != 0 || + usage.PromptTokensDetails.CacheWriteTokens != 0 || + usage.PromptTokensDetails.TextTokens != 0 || + usage.PromptTokensDetails.ImageTokens != 0 || + usage.PromptTokensDetails.AudioTokens != 0 { + return true + } + if usage.CompletionTokenDetails.ReasoningTokens != 0 || + usage.CompletionTokenDetails.TextTokens != 0 || + usage.CompletionTokenDetails.ImageTokens != 0 || + usage.CompletionTokenDetails.AudioTokens != 0 { + return true + } + return usage.InputTokensDetails != nil +} + +func NewGeminiChatBillingUsage(metadata *GeminiUsageMetadata) *BillingUsage { + return newGeminiChatBillingUsage(metadata, false) +} + +func NewEstimatedGeminiChatBillingUsage(usage *Usage) *BillingUsage { + if usage == nil { + return nil + } + totalTokens := usage.TotalTokens + if totalTokens == 0 { + totalTokens = usage.PromptTokens + usage.CompletionTokens + } + return newGeminiChatBillingUsage(&GeminiUsageMetadata{ + PromptTokenCount: usage.PromptTokens, + CandidatesTokenCount: usage.CompletionTokens, + TotalTokenCount: totalTokens, + }, true) +} + +func newGeminiChatBillingUsage(metadata *GeminiUsageMetadata, estimated bool) *BillingUsage { + if !HasGeminiUsageMetadataTokens(metadata) { + return nil + } + usageMetadata := cloneGeminiUsageMetadata(*metadata) + return &BillingUsage{ + Source: BillingUsageSourceGeminiChat, + Semantic: BillingUsageSemanticGemini, + Estimated: estimated, + GeminiUsageMetadata: &usageMetadata, + } +} + +func CloneBillingUsage(usage *BillingUsage) *BillingUsage { + if usage == nil { + return nil + } + clone := *usage + clone.OpenAIUsage = cloneOpenAIUsage(usage.OpenAIUsage) + clone.ClaudeUsage = cloneClaudeUsage(usage.ClaudeUsage) + if usage.GeminiUsageMetadata != nil { + metadata := cloneGeminiUsageMetadata(*usage.GeminiUsageMetadata) + clone.GeminiUsageMetadata = &metadata + } + return &clone +} + +func cloneOpenAIUsage(usage *Usage) *Usage { + if usage == nil { + return nil + } + clone := *usage + clone.BillingUsage = nil + if usage.InputTokensDetails != nil { + inputTokensDetails := *usage.InputTokensDetails + clone.InputTokensDetails = &inputTokensDetails + } + return &clone +} + +func cloneClaudeUsage(usage *ClaudeUsage) *ClaudeUsage { + if usage == nil { + return nil + } + clone := *usage + clone.BillingUsage = nil + if usage.CacheCreation != nil { + cacheCreation := *usage.CacheCreation + clone.CacheCreation = &cacheCreation + } + if usage.ServerToolUse != nil { + serverToolUse := *usage.ServerToolUse + clone.ServerToolUse = &serverToolUse + } + return &clone +} + +func cloneGeminiUsageMetadata(metadata GeminiUsageMetadata) GeminiUsageMetadata { + metadata.PromptTokensDetails = append([]GeminiPromptTokensDetails{}, metadata.PromptTokensDetails...) + metadata.ToolUsePromptTokensDetails = append([]GeminiPromptTokensDetails{}, metadata.ToolUsePromptTokensDetails...) + metadata.CandidatesTokensDetails = append([]GeminiPromptTokensDetails{}, metadata.CandidatesTokensDetails...) + metadata.BillingUsage = nil + return metadata +} + +func HasGeminiUsageMetadataTokens(metadata *GeminiUsageMetadata) bool { + if metadata == nil { + return false + } + if metadata.PromptTokenCount != 0 || + metadata.ToolUsePromptTokenCount != 0 || + metadata.CandidatesTokenCount != 0 || + metadata.TotalTokenCount != 0 || + metadata.ThoughtsTokenCount != 0 || + metadata.CachedContentTokenCount != 0 { + return true + } + for _, detail := range metadata.PromptTokensDetails { + if detail.TokenCount != 0 { + return true + } + } + for _, detail := range metadata.ToolUsePromptTokensDetails { + if detail.TokenCount != 0 { + return true + } + } + for _, detail := range metadata.CandidatesTokensDetails { + if detail.TokenCount != 0 { + return true + } + } + return false +} diff --git a/dto/billing_usage_test.go b/dto/billing_usage_test.go new file mode 100644 index 0000000..b9bc8b3 --- /dev/null +++ b/dto/billing_usage_test.go @@ -0,0 +1,89 @@ +package dto + +import ( + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNewGeminiChatBillingUsageRequiresTokenContent(t *testing.T) { + require.Nil(t, NewGeminiChatBillingUsage(nil)) + require.Nil(t, NewGeminiChatBillingUsage(&GeminiUsageMetadata{})) + + billingUsage := NewGeminiChatBillingUsage(&GeminiUsageMetadata{PromptTokenCount: 1}) + require.NotNil(t, billingUsage) + require.NotNil(t, billingUsage.GeminiUsageMetadata) + assert.Equal(t, BillingUsageSourceGeminiChat, billingUsage.Source) + assert.Equal(t, BillingUsageSemanticGemini, billingUsage.Semantic) + assert.False(t, billingUsage.Estimated) +} + +func TestNewClaudeMessagesBillingUsageRequiresTokenContent(t *testing.T) { + require.Nil(t, NewClaudeMessagesBillingUsage(nil)) + require.Nil(t, NewClaudeMessagesBillingUsage(&ClaudeUsage{})) + require.Nil(t, NewClaudeMessagesBillingUsage(&ClaudeUsage{CacheCreation: &ClaudeCacheCreationUsage{}})) + + billingUsage := NewClaudeMessagesBillingUsage(&ClaudeUsage{InputTokens: 1}) + require.NotNil(t, billingUsage) + require.NotNil(t, billingUsage.ClaudeUsage) + assert.Equal(t, BillingUsageSourceClaudeMessages, billingUsage.Source) + assert.Equal(t, BillingUsageSemanticAnthropic, billingUsage.Semantic) + + cacheOnly := NewClaudeMessagesBillingUsage(&ClaudeUsage{ + CacheCreation: &ClaudeCacheCreationUsage{Ephemeral5mInputTokens: 4}, + }) + require.NotNil(t, cacheOnly) +} + +func TestNewOpenAIChatBillingUsageRequiresTokenContent(t *testing.T) { + require.Nil(t, NewOpenAIChatBillingUsage(nil)) + require.Nil(t, NewOpenAIChatBillingUsage(&Usage{})) + + billingUsage := NewOpenAIChatBillingUsage(&Usage{PromptTokens: 1}) + require.NotNil(t, billingUsage) + require.NotNil(t, billingUsage.OpenAIUsage) + assert.Equal(t, BillingUsageSourceOAIChat, billingUsage.Source) + assert.Equal(t, BillingUsageSemanticOpenAI, billingUsage.Semantic) + assert.Equal(t, 1, billingUsage.OpenAIUsage.PromptTokens) +} + +func TestNewEstimatedGeminiChatBillingUsage(t *testing.T) { + billingUsage := NewEstimatedGeminiChatBillingUsage(&Usage{ + PromptTokens: 11, + CompletionTokens: 7, + }) + + require.NotNil(t, billingUsage) + require.NotNil(t, billingUsage.GeminiUsageMetadata) + assert.True(t, billingUsage.Estimated) + assert.Equal(t, 11, billingUsage.GeminiUsageMetadata.PromptTokenCount) + assert.Equal(t, 7, billingUsage.GeminiUsageMetadata.CandidatesTokenCount) + assert.Equal(t, 18, billingUsage.GeminiUsageMetadata.TotalTokenCount) +} + +func TestBillingUsageJSONUsesProtocolNamedFields(t *testing.T) { + billingUsage := &BillingUsage{ + OpenAIUsage: &Usage{PromptTokens: 1, BillingUsage: NewClaudeMessagesBillingUsage(&ClaudeUsage{InputTokens: 9})}, + ClaudeUsage: &ClaudeUsage{InputTokens: 2, BillingUsage: NewOpenAIChatBillingUsage(&Usage{PromptTokens: 8})}, + GeminiUsageMetadata: &GeminiUsageMetadata{PromptTokenCount: 3, BillingUsage: NewOpenAIChatBillingUsage(&Usage{PromptTokens: 7})}, + } + + data, err := common.Marshal(billingUsage) + require.NoError(t, err) + + assert.Contains(t, string(data), `"openai_usage"`) + assert.Contains(t, string(data), `"claude_usage"`) + assert.Contains(t, string(data), `"gemini_usage_metadata"`) + assert.NotContains(t, string(data), `"usage":`) + assert.NotContains(t, string(data), `"usage_metadata"`) + + clone := CloneBillingUsage(billingUsage) + require.NotNil(t, clone.OpenAIUsage) + require.NotNil(t, clone.ClaudeUsage) + require.NotNil(t, clone.GeminiUsageMetadata) + assert.Nil(t, clone.OpenAIUsage.BillingUsage) + assert.Nil(t, clone.ClaudeUsage.BillingUsage) + assert.Nil(t, clone.GeminiUsageMetadata.BillingUsage) +} diff --git a/dto/claude.go b/dto/claude.go index 9d0b5ff..7c9863f 100644 --- a/dto/claude.go +++ b/dto/claude.go @@ -564,6 +564,7 @@ type ClaudeUsage struct { ClaudeCacheCreation5mTokens int `json:"claude_cache_creation_5_m_tokens"` ClaudeCacheCreation1hTokens int `json:"claude_cache_creation_1_h_tokens"` ServerToolUse *ClaudeServerToolUse `json:"server_tool_use,omitempty"` + BillingUsage *BillingUsage `json:"billing_usage,omitempty"` } type ClaudeCacheCreationUsage struct { diff --git a/dto/gemini.go b/dto/gemini.go index 3095144..6a445d6 100644 --- a/dto/gemini.go +++ b/dto/gemini.go @@ -455,9 +455,40 @@ type GeminiChatPromptFeedback struct { } type GeminiChatResponse struct { - Candidates []GeminiChatCandidate `json:"candidates"` - PromptFeedback *GeminiChatPromptFeedback `json:"promptFeedback,omitempty"` - UsageMetadata GeminiUsageMetadata `json:"usageMetadata"` + Candidates []GeminiChatCandidate `json:"candidates"` + PromptFeedback *GeminiChatPromptFeedback `json:"promptFeedback,omitempty"` + UsageMetadata GeminiUsageMetadata `json:"usageMetadata"` + HasUsageMetadata bool `json:"-"` +} + +func (r *GeminiChatResponse) UnmarshalJSON(data []byte) error { + var aux struct { + Candidates []GeminiChatCandidate `json:"candidates"` + PromptFeedback *GeminiChatPromptFeedback `json:"promptFeedback,omitempty"` + UsageMetadata *GeminiUsageMetadata `json:"usageMetadata"` + } + if err := common.Unmarshal(data, &aux); err != nil { + return err + } + r.Candidates = aux.Candidates + r.PromptFeedback = aux.PromptFeedback + r.HasUsageMetadata = aux.UsageMetadata != nil + if aux.UsageMetadata != nil { + r.UsageMetadata = *aux.UsageMetadata + } else { + r.UsageMetadata = GeminiUsageMetadata{} + } + return nil +} + +func (r *GeminiChatResponse) GetUsageMetadata() *GeminiUsageMetadata { + if r == nil { + return nil + } + if r.HasUsageMetadata || HasGeminiUsageMetadataTokens(&r.UsageMetadata) { + return &r.UsageMetadata + } + return nil } type GeminiUsageMetadata struct { @@ -470,6 +501,7 @@ type GeminiUsageMetadata struct { PromptTokensDetails []GeminiPromptTokensDetails `json:"promptTokensDetails"` ToolUsePromptTokensDetails []GeminiPromptTokensDetails `json:"toolUsePromptTokensDetails"` CandidatesTokensDetails []GeminiPromptTokensDetails `json:"candidatesTokensDetails"` + BillingUsage *BillingUsage `json:"billing_usage,omitempty"` } type GeminiPromptTokensDetails struct { diff --git a/dto/openai_response.go b/dto/openai_response.go index 00069d1..8242d14 100644 --- a/dto/openai_response.go +++ b/dto/openai_response.go @@ -221,12 +221,13 @@ type CompletionsStreamResponse struct { } type Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` - PromptCacheHitTokens int `json:"prompt_cache_hit_tokens,omitempty"` - UsageSemantic string `json:"usage_semantic,omitempty"` - UsageSource string `json:"usage_source,omitempty"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + PromptCacheHitTokens int `json:"prompt_cache_hit_tokens,omitempty"` + UsageSemantic string `json:"usage_semantic,omitempty"` + UsageSource string `json:"usage_source,omitempty"` + BillingUsage *BillingUsage `json:"billing_usage,omitempty"` PromptTokensDetails InputTokenDetails `json:"prompt_tokens_details"` CompletionTokenDetails OutputTokenDetails `json:"completion_tokens_details"` @@ -255,11 +256,23 @@ type OpenAIVideoResponse struct { type InputTokenDetails struct { CachedTokens int `json:"cached_tokens"` CachedCreationTokens int `json:"cached_creation_tokens,omitempty"` + CacheWriteTokens int `json:"cache_write_tokens,omitempty"` TextTokens int `json:"text_tokens"` AudioTokens int `json:"audio_tokens"` ImageTokens int `json:"image_tokens"` } +func (d InputTokenDetails) CacheCreationTokensTotal() int { + total := d.CachedCreationTokens + if d.CacheWriteTokens > total { + total = d.CacheWriteTokens + } + if total < 0 { + return 0 + } + return total +} + type OutputTokenDetails struct { TextTokens int `json:"text_tokens"` AudioTokens int `json:"audio_tokens"` diff --git a/model/ability.go b/model/ability.go index 520c1b7..289f99c 100644 --- a/model/ability.go +++ b/model/ability.go @@ -51,7 +51,11 @@ func GetGroupEnabledModels(group string) []string { func GetEnabledModels() []string { var models []string // Find distinct models - DB.Table("abilities").Where("enabled = ?", true).Distinct("model").Pluck("model", &models) + DB.Table("abilities"). + Joins("JOIN channels ON abilities.channel_id = channels.id"). + Where("abilities.enabled = ? AND channels.status = ?", true, common.ChannelStatusEnabled). + Distinct("abilities.model"). + Pluck("abilities.model", &models) return models } diff --git a/pkg/billingexpr/round.go b/pkg/billingexpr/round.go index fbdc64d..e3bf4a9 100644 --- a/pkg/billingexpr/round.go +++ b/pkg/billingexpr/round.go @@ -1,23 +1,15 @@ package billingexpr -import "math" +import "github.com/MAX-API-Next/MAX-API/common" // QuotaRound converts a float64 quota value to int using half-away-from-zero // rounding. Every tiered billing path (pre-consume, settlement, breakdown // validation, log fields) MUST use this function to avoid +-1 discrepancies. -// -// The result saturates at int32 bounds: quota columns are 32-bit integers in -// the database, and oversized expression results must never wrap around. func QuotaRound(f float64) int { - r := math.Round(f) - if math.IsNaN(r) { - return 0 - } - if r >= math.MaxInt32 { - return math.MaxInt32 - } - if r <= math.MinInt32 { - return math.MinInt32 - } - return int(r) + return common.QuotaRound(f) +} + +// QuotaRoundStrict rejects an unrepresentable pre-consume estimate. +func QuotaRoundStrict(f float64) (int, error) { + return common.QuotaRoundStrict(f) } diff --git a/pkg/billingexpr/settle.go b/pkg/billingexpr/settle.go index 7a6ca44..ec86adc 100644 --- a/pkg/billingexpr/settle.go +++ b/pkg/billingexpr/settle.go @@ -1,5 +1,7 @@ package billingexpr +import "github.com/MAX-API-Next/MAX-API/common" + // quotaConversion converts raw expression output to quota based on the // expression version. This is the central dispatch point for future versions // that may use a different conversion formula. @@ -23,7 +25,7 @@ func ComputeTieredQuotaWithRequest(snap *BillingSnapshot, params TokenParams, re } quotaBeforeGroup := quotaConversion(cost, snap) - afterGroup := QuotaRound(quotaBeforeGroup * snap.GroupRatio) + afterGroup, clamp := common.QuotaRoundChecked(quotaBeforeGroup * snap.GroupRatio) crossed := trace.MatchedTier != snap.EstimatedTier return TieredResult{ @@ -31,5 +33,6 @@ func ComputeTieredQuotaWithRequest(snap *BillingSnapshot, params TokenParams, re ActualQuotaAfterGroup: afterGroup, MatchedTier: trace.MatchedTier, CrossedTier: crossed, + Clamp: clamp, }, nil } diff --git a/pkg/billingexpr/types.go b/pkg/billingexpr/types.go index 12e0d3c..8e363c4 100644 --- a/pkg/billingexpr/types.go +++ b/pkg/billingexpr/types.go @@ -3,6 +3,8 @@ package billingexpr import ( "crypto/sha256" "fmt" + + "github.com/MAX-API-Next/MAX-API/common" ) type RequestInput struct { @@ -57,6 +59,9 @@ type TieredResult struct { ActualQuotaAfterGroup int `json:"actual_quota_after_group"` MatchedTier string `json:"matched_tier"` CrossedTier bool `json:"crossed_tier"` + // Clamp records an int32 saturation event during quota conversion so the + // caller can surface it on the consume log for admin auditing. + Clamp *common.QuotaClamp `json:"-"` } // ExprHashString returns the SHA-256 hex digest of an expression string. diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 1a544fc..aeb89f8 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -665,12 +665,14 @@ func buildOpenAIStyleUsageFromClaudeUsage(usage *dto.Usage) dto.Usage { return dto.Usage{} } clone := *usage + clone.BillingUsage = dto.CloneBillingUsage(usage.BillingUsage) clone.ClaudeCacheCreation5mTokens, clone.ClaudeCacheCreation1hTokens = service.NormalizeCacheCreationSplit( usage.PromptTokensDetails.CachedCreationTokens, usage.ClaudeCacheCreation5mTokens, usage.ClaudeCacheCreation1hTokens, ) cacheCreationTokens := cacheCreationTokensForOpenAIUsage(usage) + clone.PromptTokensDetails.CacheWriteTokens = cacheCreationTokens totalInputTokens := usage.PromptTokens + usage.PromptTokensDetails.CachedTokens + cacheCreationTokens clone.PromptTokens = totalInputTokens clone.InputTokens = totalInputTokens @@ -908,6 +910,9 @@ func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, clau if claudeInfo.Usage != nil { claudeInfo.Usage.UsageSemantic = "anthropic" } + if claudeInfo.Usage != nil && claudeInfo.Usage.BillingUsage == nil { + claudeInfo.Usage.BillingUsage = dto.NewClaudeMessagesBillingUsage(buildMessageDeltaPatchUsage(nil, claudeInfo)) + } if info.RelayFormat == types.RelayFormatClaude { // @@ -968,6 +973,7 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud claudeInfo.Usage.CompletionTokens = claudeResponse.Usage.OutputTokens claudeInfo.Usage.TotalTokens = claudeResponse.Usage.InputTokens + claudeResponse.Usage.OutputTokens claudeInfo.Usage.UsageSemantic = "anthropic" + claudeInfo.Usage.BillingUsage = dto.NewClaudeMessagesBillingUsage(claudeResponse.Usage) claudeInfo.Usage.PromptTokensDetails.CachedTokens = claudeResponse.Usage.CacheReadInputTokens claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens = claudeResponse.Usage.CacheCreationInputTokens claudeInfo.Usage.ClaudeCacheCreation5mTokens = claudeResponse.Usage.GetCacheCreation5mTokens() diff --git a/relay/channel/gemini/relay-gemini-native.go b/relay/channel/gemini/relay-gemini-native.go index f2350d3..cd9c4c5 100644 --- a/relay/channel/gemini/relay-gemini-native.go +++ b/relay/channel/gemini/relay-gemini-native.go @@ -39,8 +39,8 @@ func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, re common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason)) } - // 计算使用量(基于 UsageMetadata) - usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) + // 计算使用量(优先上游 UsageMetadata,缺失时本地估算并保留 Gemini 计费语义) + usage := buildUsageFromGeminiResponse(c, info, &geminiResponse) if service.ResponseAuditEnabled() { service.SetRelayResponseAuditContent(info, service.BuildGeminiResponseAuditContent(&geminiResponse)) } diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 0bdbe63..85d4b8f 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1058,9 +1058,91 @@ func buildUsageFromGeminiMetadata(metadata dto.GeminiUsageMetadata, fallbackProm usage.PromptTokensDetails.TextTokens = usage.PromptTokens } + usage.BillingUsage = dto.NewGeminiChatBillingUsage(&metadata) return usage } +func attachEstimatedGeminiBillingUsage(usage *dto.Usage) *dto.Usage { + if usage != nil && usage.BillingUsage == nil { + usage.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(usage) + } + return usage +} + +// patchGeminiZeroCompletionUsage estimates completion tokens locally when +// upstream usageMetadata was billable but reported zero completion tokens even +// though output content was received. Without replacing BillingUsage, settlement +// would still prefer the prompt-only metadata and bill zero completion. +func patchGeminiZeroCompletionUsage(c *gin.Context, info *relaycommon.RelayInfo, usage *dto.Usage, responseText string, imageCount int) { + if usage == nil || usage.CompletionTokens > 0 { + return + } + if responseText == "" && imageCount == 0 { + return + } + promptTokens := usage.PromptTokens + if promptTokens <= 0 { + promptTokens = info.GetEstimatePromptTokens() + } + estimated := service.ResponseText2Usage(c, responseText, info.UpstreamModelName, promptTokens) + if usage.PromptTokens == 0 { + usage.PromptTokens = estimated.PromptTokens + } + usage.CompletionTokens = estimated.CompletionTokens + if imageCount != 0 && usage.CompletionTokens == 0 { + usage.CompletionTokens = imageCount * 1400 + } + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + usage.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(usage) +} + +func geminiResponseUsageText(response *dto.GeminiChatResponse) string { + if response == nil { + return "" + } + var text strings.Builder + for _, candidate := range response.Candidates { + for _, part := range candidate.Content.Parts { + if part.Text != "" { + text.WriteString(part.Text) + } + } + } + return text.String() +} + +func geminiResponseInlineImageCount(response *dto.GeminiChatResponse) int { + if response == nil { + return 0 + } + count := 0 + for _, candidate := range response.Candidates { + for _, part := range candidate.Content.Parts { + if part.InlineData != nil && strings.HasPrefix(part.InlineData.MimeType, "image") { + count++ + } + } + } + return count +} + +func buildUsageFromGeminiResponse(c *gin.Context, info *relaycommon.RelayInfo, response *dto.GeminiChatResponse) dto.Usage { + if response == nil { + return dto.Usage{} + } + if metadata := response.GetUsageMetadata(); dto.HasGeminiUsageMetadataTokens(metadata) { + usage := buildUsageFromGeminiMetadata(*metadata, info.GetEstimatePromptTokens()) + patchGeminiZeroCompletionUsage(c, info, &usage, geminiResponseUsageText(response), geminiResponseInlineImageCount(response)) + return usage + } + usage := service.ResponseText2Usage(c, geminiResponseUsageText(response), info.UpstreamModelName, info.GetEstimatePromptTokens()) + attachEstimatedGeminiBillingUsage(usage) + if usage == nil { + return dto.Usage{} + } + return *usage +} + func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse) *dto.OpenAITextResponse { fullTextResponse := dto.OpenAITextResponse{ Id: helper.GetResponseID(c), @@ -1380,8 +1462,8 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } // 更新使用量统计 - if geminiResponse.UsageMetadata.TotalTokenCount != 0 { - mappedUsage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) + if metadata := geminiResponse.GetUsageMetadata(); dto.HasGeminiUsageMetadataTokens(metadata) { + mappedUsage := buildUsageFromGeminiMetadata(*metadata, info.GetEstimatePromptTokens()) *usage = mappedUsage } @@ -1390,15 +1472,12 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } }) - if imageCount != 0 { - if usage.CompletionTokens == 0 { - usage.CompletionTokens = imageCount * 1400 - } - } + patchGeminiZeroCompletionUsage(c, info, usage, responseText.String(), imageCount) if usage.CompletionTokens <= 0 { if info.ReceivedResponseCount > 0 { usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + attachEstimatedGeminiBillingUsage(usage) } else { usage = &dto.Usage{} } @@ -1535,7 +1614,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } if len(geminiResponse.Candidates) == 0 { - usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) + usage := buildUsageFromGeminiResponse(c, info, &geminiResponse) var maxAPIError *types.MaxAPIError if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil { @@ -1571,7 +1650,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R } fullTextResponse := responseGeminiChat2OpenAI(c, &geminiResponse) fullTextResponse.Model = info.UpstreamModelName - usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) + usage := buildUsageFromGeminiResponse(c, info, &geminiResponse) if service.ResponseAuditEnabled() { service.SetRelayResponseAuditContent(info, service.BuildGeminiResponseAuditContent(&geminiResponse)) } diff --git a/relay/channel/gemini/relay_responses.go b/relay/channel/gemini/relay_responses.go index 6780ed4..a5762e4 100644 --- a/relay/channel/gemini/relay_responses.go +++ b/relay/channel/gemini/relay_responses.go @@ -32,7 +32,7 @@ func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } if len(geminiResponse.Candidates) == 0 { - usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) + usage := buildUsageFromGeminiResponse(c, info, &geminiResponse) if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil { common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason)) return &usage, types.NewOpenAIError( @@ -51,7 +51,7 @@ func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h chatResp := responseGeminiChat2OpenAI(c, &geminiResponse) chatResp.Model = info.UpstreamModelName - usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) + usage := buildUsageFromGeminiResponse(c, info, &geminiResponse) chatResp.Usage = usage responsesResp, responsesUsage, err := service.ChatCompletionsResponseToResponsesResponse(chatResp, helper.GetResponseID(c)) diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 69bdc79..4ba3c13 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -594,6 +594,7 @@ func OpenaiHandlerWithUsage(c *gin.Context, info *relaycommon.RelayInfo, resp *h if usageResp.InputTokensDetails != nil { usageResp.PromptTokensDetails.ImageTokens += usageResp.InputTokensDetails.ImageTokens usageResp.PromptTokensDetails.TextTokens += usageResp.InputTokensDetails.TextTokens + usageResp.PromptTokensDetails.CacheWriteTokens += usageResp.InputTokensDetails.CacheWriteTokens } applyUsagePostProcessing(info, &usageResp.Usage, responseBody) return &usageResp.Usage, nil diff --git a/relay/channel/openai/relay_image.go b/relay/channel/openai/relay_image.go index 91a08da..ab12cad 100644 --- a/relay/channel/openai/relay_image.go +++ b/relay/channel/openai/relay_image.go @@ -63,6 +63,7 @@ func normalizeOpenAIImageUsage(usage *dto.Usage) { if usage.InputTokensDetails != nil { usage.PromptTokensDetails.CachedTokens = usage.InputTokensDetails.CachedTokens usage.PromptTokensDetails.CachedCreationTokens = usage.InputTokensDetails.CachedCreationTokens + usage.PromptTokensDetails.CacheWriteTokens = usage.InputTokensDetails.CacheWriteTokens usage.PromptTokensDetails.ImageTokens = usage.InputTokensDetails.ImageTokens usage.PromptTokensDetails.TextTokens = usage.InputTokensDetails.TextTokens usage.PromptTokensDetails.AudioTokens = usage.InputTokensDetails.AudioTokens diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 2b160e3..ab5ea4f 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -53,6 +53,7 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http usage.TotalTokens = responsesResponse.Usage.TotalTokens if responsesResponse.Usage.InputTokensDetails != nil { usage.PromptTokensDetails.CachedTokens = responsesResponse.Usage.InputTokensDetails.CachedTokens + usage.PromptTokensDetails.CacheWriteTokens = responsesResponse.Usage.InputTokensDetails.CacheWriteTokens } } emptyCompletion := isEmptyResponsesCompletion(&responsesResponse) @@ -146,6 +147,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } if streamResponse.Response.Usage.InputTokensDetails != nil { usage.PromptTokensDetails.CachedTokens = streamResponse.Response.Usage.InputTokensDetails.CachedTokens + usage.PromptTokensDetails.CacheWriteTokens = streamResponse.Response.Usage.InputTokensDetails.CacheWriteTokens } } if streamResponse.Response.HasImageGenerationCall() { diff --git a/relay/channel/openai/relay_responses_compact.go b/relay/channel/openai/relay_responses_compact.go index 4a1ed1a..b24a19c 100644 --- a/relay/channel/openai/relay_responses_compact.go +++ b/relay/channel/openai/relay_responses_compact.go @@ -41,6 +41,7 @@ func OaiResponsesCompactionHandler(c *gin.Context, info *relaycommon.RelayInfo, usage.TotalTokens = compactResp.Usage.TotalTokens if compactResp.Usage.InputTokensDetails != nil { usage.PromptTokensDetails.CachedTokens = compactResp.Usage.InputTokensDetails.CachedTokens + usage.PromptTokensDetails.CacheWriteTokens = compactResp.Usage.InputTokensDetails.CacheWriteTokens } } diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index f2aaeea..7760284 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -166,6 +166,10 @@ type RelayInfo struct { PriceData types.PriceData TaskBilling *types.TaskBillingResult + // QuotaClamp is set when a quota conversion saturated at the int32 bound + // or fell back from NaN while computing this request's charge. + QuotaClamp *common.QuotaClamp + // TieredBillingSnapshot is a frozen snapshot of tiered billing rules // captured at pre-consume time. Non-nil only when billing mode is "tiered_expr". TieredBillingSnapshot *billingexpr.BillingSnapshot diff --git a/relay/helper/price.go b/relay/helper/price.go index 10585bf..73a35eb 100644 --- a/relay/helper/price.go +++ b/relay/helper/price.go @@ -113,12 +113,20 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens audioRatio = ratio_setting.GetAudioRatio(info.OriginModelName) audioCompletionRatio = ratio_setting.GetAudioCompletionRatio(info.OriginModelName) ratio := modelRatio * groupRatioInfo.GroupRatio - preConsumedQuota = common.QuotaFromFloat(float64(preConsumedTokens) * ratio) + quota, err := common.QuotaFromFloatStrict(float64(preConsumedTokens) * ratio) + if err != nil { + return types.PriceData{}, err + } + preConsumedQuota = quota } else { if meta.ImagePriceRatio != 0 { modelPrice = modelPrice * meta.ImagePriceRatio } - preConsumedQuota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + quota, err := common.QuotaFromFloatStrict(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + if err != nil { + return types.PriceData{}, err + } + preConsumedQuota = quota } // check if free model pre-consume is disabled @@ -199,7 +207,11 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types freeModel := false if usePrice { - quota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + var err error + quota, err = common.QuotaFromFloatStrict(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + if err != nil { + return types.PriceData{}, err + } if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { if groupRatioInfo.GroupRatio == 0 || (modelPrice == 0 && !rateCardPriced) { quota = 0 @@ -208,7 +220,11 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types } } else { // 按量计费:以模型倍率的一半作为预扣额度 - quota = common.QuotaFromFloat(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + var err error + quota, err = common.QuotaFromFloatStrict(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + if err != nil { + return types.PriceData{}, err + } modelPrice = -1 if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { if groupRatioInfo.GroupRatio == 0 || modelRatio == 0 { @@ -270,7 +286,10 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT // Expression coefficients are $/1M tokens prices; convert to quota the same way per-call billing does. quotaBeforeGroup := rawCost / 1_000_000 * common.QuotaPerUnit - preConsumedQuota := billingexpr.QuotaRound(quotaBeforeGroup * groupRatioInfo.GroupRatio) + preConsumedQuota, err := billingexpr.QuotaRoundStrict(quotaBeforeGroup * groupRatioInfo.GroupRatio) + if err != nil { + return types.PriceData{}, err + } freeModel := false if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { diff --git a/relay/relay_task.go b/relay/relay_task.go index 6e892b5..34e015b 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -222,7 +222,9 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe quotaWithRatios *= ra } } - info.PriceData.Quota = common.QuotaFromFloat(quotaWithRatios) + quota, clamp := common.QuotaFromFloatChecked(quotaWithRatios) + info.PriceData.Quota = quota + noteTaskQuotaClamp(info, clamp) } } @@ -373,7 +375,18 @@ func recalcQuotaFromRatios(info *relaycommon.RelayInfo, ratios map[string]float6 result *= ra } } - return common.QuotaFromFloat(result) + quota, clamp := common.QuotaFromFloatChecked(result) + noteTaskQuotaClamp(info, clamp) + return quota +} + +func noteTaskQuotaClamp(info *relaycommon.RelayInfo, clamp *common.QuotaClamp) { + if clamp == nil || info == nil { + return + } + if info.QuotaClamp == nil { + info.QuotaClamp = clamp + } } var fetchRespBuilders = map[int]func(c *gin.Context) (respBody []byte, taskResp *dto.TaskError){ diff --git a/service/billing.go b/service/billing.go index 80c42c6..67d1156 100644 --- a/service/billing.go +++ b/service/billing.go @@ -2,6 +2,7 @@ package service import ( "fmt" + "net/http" "github.com/MAX-API-Next/MAX-API/logger" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" @@ -17,6 +18,22 @@ const ( // PreConsumeBilling 根据用户计费偏好创建 BillingSession 并执行预扣费。 // 会话存储在 relayInfo.Billing 上,供后续 Settle / Refund 使用。 func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { + if relayInfo != nil && relayInfo.QuotaClamp != nil { + return types.NewErrorWithStatusCode( + relayInfo.QuotaClamp, + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + if preConsumedQuota < 0 { + return types.NewErrorWithStatusCode( + fmt.Errorf("pre-consume quota cannot be negative: %d", preConsumedQuota), + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } session, apiErr := NewBillingSession(c, relayInfo, preConsumedQuota) if apiErr != nil { return apiErr diff --git a/service/billing_usage.go b/service/billing_usage.go new file mode 100644 index 0000000..a83d5e0 --- /dev/null +++ b/service/billing_usage.go @@ -0,0 +1,205 @@ +package service + +import ( + "strings" + + "github.com/MAX-API-Next/MAX-API/dto" +) + +const ( + usageBillingPathLocal = "local" + usageBillingPathUpstream = "upstream" + usageBillingPathOpenAI = "billing-usage-openai" + usageBillingPathOpenAIEstimated = "billing-usage-openai-estimated" + usageBillingPathAnthropic = "billing-usage-anthropic" + usageBillingPathAnthropicEstimated = "billing-usage-anthropic-estimated" + usageBillingPathGemini = "billing-usage-gemini" + usageBillingPathGeminiEstimated = "billing-usage-gemini-estimated" +) + +func effectiveBillingUsage(usage *dto.Usage) *dto.Usage { + if billingUsage, ok := usageFromBillingUsage(usage); ok { + return billingUsage + } + return usage +} + +func usageBillingPathForLog(isLocalCountTokens bool, usage *dto.Usage) string { + if isLocalCountTokens { + return usageBillingPathLocal + } + if usage == nil || usage.BillingUsage == nil { + 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) { + if usage.BillingUsage.Estimated { + return usageBillingPathOpenAIEstimated + } + return usageBillingPathOpenAI + } + if strings.EqualFold(source, dto.BillingUsageSourceClaudeMessages) || + strings.EqualFold(semantic, dto.BillingUsageSemanticAnthropic) { + if usage.BillingUsage.Estimated { + return usageBillingPathAnthropicEstimated + } + return usageBillingPathAnthropic + } + if strings.EqualFold(source, dto.BillingUsageSourceGeminiChat) || + strings.EqualFold(semantic, dto.BillingUsageSemanticGemini) { + if usage.BillingUsage.Estimated { + return usageBillingPathGeminiEstimated + } + return usageBillingPathGemini + } + return usageBillingPathUpstream +} + +func appendUsageBillingPathForLog(other map[string]interface{}, isLocalCountTokens bool, usage *dto.Usage) { + if other == nil { + return + } + adminInfo, ok := other["admin_info"].(map[string]interface{}) + if !ok || adminInfo == nil { + adminInfo = make(map[string]interface{}) + other["admin_info"] = adminInfo + } + adminInfo["usage_billing_path"] = usageBillingPathForLog(isLocalCountTokens, usage) +} + +func usageFromBillingUsage(usage *dto.Usage) (*dto.Usage, bool) { + if usage == nil || usage.BillingUsage == nil { + return nil, false + } + billingUsage := usage.BillingUsage + source := strings.TrimSpace(billingUsage.Source) + semantic := strings.TrimSpace(billingUsage.Semantic) + + if billingUsage.OpenAIUsage != nil && + (strings.EqualFold(source, dto.BillingUsageSourceOAIChat) || + strings.EqualFold(source, dto.BillingUsageSourceOAIResponses) || + strings.EqualFold(semantic, dto.BillingUsageSemanticOpenAI)) { + return usageFromOpenAIBillingUsage(billingUsage), true + } + + if billingUsage.ClaudeUsage != nil && + (strings.EqualFold(source, dto.BillingUsageSourceClaudeMessages) || + strings.EqualFold(semantic, dto.BillingUsageSemanticAnthropic)) { + return usageFromClaudeBillingUsage(billingUsage), true + } + + if billingUsage.GeminiUsageMetadata != nil && + (strings.EqualFold(source, dto.BillingUsageSourceGeminiChat) || + strings.EqualFold(semantic, dto.BillingUsageSemanticGemini)) { + return usageFromGeminiBillingUsage(billingUsage), true + } + + return nil, false +} + +func usageFromOpenAIBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage { + usage := *billingUsage.OpenAIUsage + if usage.PromptTokens == 0 && usage.InputTokens > 0 { + usage.PromptTokens = usage.InputTokens + } + if usage.CompletionTokens == 0 && usage.OutputTokens > 0 { + usage.CompletionTokens = usage.OutputTokens + } + if usage.InputTokens == 0 && usage.PromptTokens > 0 { + usage.InputTokens = usage.PromptTokens + } + if usage.OutputTokens == 0 && usage.CompletionTokens > 0 { + usage.OutputTokens = usage.CompletionTokens + } + if usage.TotalTokens == 0 { + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + } + usage.UsageSemantic = dto.BillingUsageSemanticOpenAI + usage.UsageSource = billingUsage.Source + usage.BillingUsage = dto.CloneBillingUsage(billingUsage) + return &usage +} + +func usageFromClaudeBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage { + claudeUsage := billingUsage.ClaudeUsage + cacheCreation5m := claudeUsage.GetCacheCreation5mTokens() + if cacheCreation5m == 0 { + cacheCreation5m = claudeUsage.ClaudeCacheCreation5mTokens + } + cacheCreation1h := claudeUsage.GetCacheCreation1hTokens() + if cacheCreation1h == 0 { + cacheCreation1h = claudeUsage.ClaudeCacheCreation1hTokens + } + + usage := &dto.Usage{ + PromptTokens: claudeUsage.InputTokens, + CompletionTokens: claudeUsage.OutputTokens, + TotalTokens: claudeUsage.InputTokens + claudeUsage.OutputTokens, + InputTokens: claudeUsage.InputTokens + claudeUsage.CacheReadInputTokens + claudeUsage.CacheCreationInputTokens, + OutputTokens: claudeUsage.OutputTokens, + UsageSemantic: dto.BillingUsageSemanticAnthropic, + UsageSource: dto.BillingUsageSourceClaudeMessages, + BillingUsage: dto.CloneBillingUsage(billingUsage), + ClaudeCacheCreation5mTokens: cacheCreation5m, + ClaudeCacheCreation1hTokens: cacheCreation1h, + } + usage.PromptTokensDetails.CachedTokens = claudeUsage.CacheReadInputTokens + usage.PromptTokensDetails.CachedCreationTokens = claudeUsage.CacheCreationInputTokens + return usage +} + +func usageFromGeminiBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage { + metadata := *billingUsage.GeminiUsageMetadata + promptTokens := metadata.PromptTokenCount + metadata.ToolUsePromptTokenCount + usage := &dto.Usage{ + PromptTokens: promptTokens, + CompletionTokens: metadata.CandidatesTokenCount + metadata.ThoughtsTokenCount, + TotalTokens: metadata.TotalTokenCount, + UsageSemantic: dto.BillingUsageSemanticGemini, + UsageSource: dto.BillingUsageSourceGeminiChat, + BillingUsage: dto.CloneBillingUsage(billingUsage), + } + usage.CompletionTokenDetails.ReasoningTokens = metadata.ThoughtsTokenCount + usage.PromptTokensDetails.CachedTokens = metadata.CachedContentTokenCount + + for _, detail := range metadata.PromptTokensDetails { + addGeminiInputTokenDetail(&usage.PromptTokensDetails, detail) + } + for _, detail := range metadata.ToolUsePromptTokensDetails { + addGeminiInputTokenDetail(&usage.PromptTokensDetails, detail) + } + for _, detail := range metadata.CandidatesTokensDetails { + switch detail.Modality { + case "IMAGE": + usage.CompletionTokenDetails.ImageTokens += detail.TokenCount + case "AUDIO": + usage.CompletionTokenDetails.AudioTokens += detail.TokenCount + case "TEXT": + usage.CompletionTokenDetails.TextTokens += detail.TokenCount + } + } + + if usage.TotalTokens == 0 { + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + } else if usage.CompletionTokens <= 0 { + usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens + } + if usage.PromptTokens > 0 && usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 { + usage.PromptTokensDetails.TextTokens = usage.PromptTokens + } + return usage +} + +func addGeminiInputTokenDetail(details *dto.InputTokenDetails, detail dto.GeminiPromptTokensDetails) { + switch detail.Modality { + case "AUDIO": + details.AudioTokens += detail.TokenCount + case "IMAGE": + details.ImageTokens += detail.TokenCount + case "TEXT": + details.TextTokens += detail.TokenCount + } +} diff --git a/service/convert.go b/service/convert.go index cadf93e..7c8c43f 100644 --- a/service/convert.go +++ b/service/convert.go @@ -833,14 +833,23 @@ func extractTextFromGeminiParts(parts []dto.GeminiPart) string { // ResponseOpenAI2Gemini 将 OpenAI 响应转换为 Gemini 格式 func ResponseOpenAI2Gemini(openAIResponse *dto.OpenAITextResponse, info *relaycommon.RelayInfo) *dto.GeminiChatResponse { + totalTokens := openAIResponse.TotalTokens + if totalTokens == 0 { + totalTokens = openAIResponse.PromptTokens + openAIResponse.CompletionTokens + } geminiResponse := &dto.GeminiChatResponse{ - Candidates: make([]dto.GeminiChatCandidate, 0, len(openAIResponse.Choices)), + Candidates: make([]dto.GeminiChatCandidate, 0, len(openAIResponse.Choices)), + HasUsageMetadata: true, UsageMetadata: dto.GeminiUsageMetadata{ PromptTokenCount: openAIResponse.PromptTokens, CandidatesTokenCount: openAIResponse.CompletionTokens, - TotalTokenCount: openAIResponse.PromptTokens + openAIResponse.CompletionTokens, + TotalTokenCount: totalTokens, + BillingUsage: openAIBillingUsageFromUsage(&openAIResponse.Usage), }, } + if metadata, ok := geminiBillingMetadataFromOpenAIUsage(&openAIResponse.Usage); ok { + geminiResponse.UsageMetadata = metadata + } for _, choice := range openAIResponse.Choices { candidate := dto.GeminiChatCandidate{ @@ -918,7 +927,8 @@ func StreamResponseOpenAI2Gemini(openAIResponse *dto.ChatCompletionsStreamRespon } geminiResponse := &dto.GeminiChatResponse{ - Candidates: make([]dto.GeminiChatCandidate, 0, len(openAIResponse.Choices)), + Candidates: make([]dto.GeminiChatCandidate, 0, len(openAIResponse.Choices)), + HasUsageMetadata: true, UsageMetadata: dto.GeminiUsageMetadata{ PromptTokenCount: info.GetEstimatePromptTokens(), CandidatesTokenCount: 0, // 流式响应中可能没有完整的 usage 信息 @@ -930,6 +940,10 @@ func StreamResponseOpenAI2Gemini(openAIResponse *dto.ChatCompletionsStreamRespon geminiResponse.UsageMetadata.PromptTokenCount = openAIResponse.Usage.PromptTokens geminiResponse.UsageMetadata.CandidatesTokenCount = openAIResponse.Usage.CompletionTokens geminiResponse.UsageMetadata.TotalTokenCount = openAIResponse.Usage.TotalTokens + geminiResponse.UsageMetadata.BillingUsage = openAIBillingUsageFromUsage(openAIResponse.Usage) + if metadata, ok := geminiBillingMetadataFromOpenAIUsage(openAIResponse.Usage); ok { + geminiResponse.UsageMetadata = metadata + } } for _, choice := range openAIResponse.Choices { @@ -991,3 +1005,31 @@ func StreamResponseOpenAI2Gemini(openAIResponse *dto.ChatCompletionsStreamRespon return geminiResponse } + +func geminiBillingMetadataFromOpenAIUsage(usage *dto.Usage) (dto.GeminiUsageMetadata, bool) { + if usage == nil || usage.BillingUsage == nil || usage.BillingUsage.GeminiUsageMetadata == nil { + return dto.GeminiUsageMetadata{}, false + } + if usage.BillingUsage.Source != dto.BillingUsageSourceGeminiChat && usage.BillingUsage.Semantic != dto.BillingUsageSemanticGemini { + return dto.GeminiUsageMetadata{}, false + } + billingUsage := dto.CloneBillingUsage(usage.BillingUsage) + if billingUsage == nil || billingUsage.GeminiUsageMetadata == nil { + return dto.GeminiUsageMetadata{}, false + } + return *billingUsage.GeminiUsageMetadata, true +} + +func openAIBillingUsageFromUsage(usage *dto.Usage) *dto.BillingUsage { + if usage == nil { + return nil + } + if existingBillingUsage := dto.CloneBillingUsage(usage.BillingUsage); existingBillingUsage != nil && existingBillingUsage.OpenAIUsage != nil { + if existingBillingUsage.Source == dto.BillingUsageSourceOAIChat || + existingBillingUsage.Source == dto.BillingUsageSourceOAIResponses || + existingBillingUsage.Semantic == dto.BillingUsageSemanticOpenAI { + return existingBillingUsage + } + } + return dto.NewOpenAIChatBillingUsage(usage) +} diff --git a/service/log_info_generate.go b/service/log_info_generate.go index cebecc6..8260e71 100644 --- a/service/log_info_generate.go +++ b/service/log_info_generate.go @@ -2,11 +2,13 @@ package service import ( "encoding/base64" + "fmt" "strings" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/constant" "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/logger" "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" @@ -14,6 +16,28 @@ import ( "github.com/gin-gonic/gin" ) +func attachQuotaSaturationToOther(other map[string]interface{}, clamp *common.QuotaClamp) { + if clamp == nil || other == nil { + return + } + adminInfo, ok := other["admin_info"].(map[string]interface{}) + if !ok || adminInfo == nil { + adminInfo = map[string]interface{}{} + other["admin_info"] = adminInfo + } + adminInfo["quota_saturation"] = clamp.AuditMap() +} + +func attachQuotaSaturation(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, other map[string]interface{}) { + if relayInfo == nil || relayInfo.QuotaClamp == nil { + return + } + clamp := relayInfo.QuotaClamp + attachQuotaSaturationToOther(other, clamp) + logger.LogWarn(ctx, fmt.Sprintf("quota saturation on consume log: op=%s kind=%s original=%g clamped=%d user=%d model=%s", + clamp.Op, clamp.Kind, clamp.Original, clamp.Clamped, relayInfo.UserId, relayInfo.OriginModelName)) +} + func appendRequestPath(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, other map[string]interface{}) { if other == nil { return diff --git a/service/openaicompat/chat_to_responses_response.go b/service/openaicompat/chat_to_responses_response.go index b461c25..bc2cc78 100644 --- a/service/openaicompat/chat_to_responses_response.go +++ b/service/openaicompat/chat_to_responses_response.go @@ -148,6 +148,7 @@ func UsageFromChatUsage(src *dto.Usage) *dto.Usage { src.PromptTokensDetails.ImageTokens != 0 || src.PromptTokensDetails.AudioTokens != 0 || src.PromptTokensDetails.CachedCreationTokens != 0 || + src.PromptTokensDetails.CacheWriteTokens != 0 || src.PromptTokensDetails.TextTokens != 0 { details := src.PromptTokensDetails usage.InputTokensDetails = &details @@ -158,6 +159,10 @@ func UsageFromChatUsage(src *dto.Usage) *dto.Usage { src.CompletionTokenDetails.ImageTokens != 0 { usage.CompletionTokenDetails = src.CompletionTokenDetails } + usage.BillingUsage = dto.CloneBillingUsage(src.BillingUsage) + if usage.BillingUsage == nil { + usage.BillingUsage = dto.NewOpenAIChatBillingUsage(src) + } return usage } diff --git a/service/openaicompat/responses_to_chat.go b/service/openaicompat/responses_to_chat.go index 8d18f66..7620fc0 100644 --- a/service/openaicompat/responses_to_chat.go +++ b/service/openaicompat/responses_to_chat.go @@ -140,9 +140,14 @@ func UsageFromResponsesUsage(src *dto.Usage) *dto.Usage { usage.PromptTokensDetails.ImageTokens = src.InputTokensDetails.ImageTokens usage.PromptTokensDetails.AudioTokens = src.InputTokensDetails.AudioTokens usage.PromptTokensDetails.CachedCreationTokens = src.InputTokensDetails.CachedCreationTokens + usage.PromptTokensDetails.CacheWriteTokens = src.InputTokensDetails.CacheWriteTokens usage.PromptTokensDetails.TextTokens = src.InputTokensDetails.TextTokens } usage.CompletionTokenDetails = src.CompletionTokenDetails + usage.BillingUsage = dto.CloneBillingUsage(src.BillingUsage) + if usage.BillingUsage == nil { + usage.BillingUsage = dto.NewOpenAIResponsesBillingUsage(src) + } return usage } diff --git a/service/pre_consume_quota.go b/service/pre_consume_quota.go index 7928da0..cee7730 100644 --- a/service/pre_consume_quota.go +++ b/service/pre_consume_quota.go @@ -31,6 +31,12 @@ func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo) { // PreConsumeQuota checks if the user has enough quota to pre-consume. // It returns the pre-consumed quota if successful, or an error if not. func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { + if relayInfo != nil && relayInfo.QuotaClamp != nil { + return types.NewErrorWithStatusCode(relayInfo.QuotaClamp, types.ErrorCodeModelPriceError, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + if preConsumedQuota < 0 { + return types.NewErrorWithStatusCode(fmt.Errorf("pre-consume quota cannot be negative: %d", preConsumedQuota), types.ErrorCodeModelPriceError, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } userQuota, err := model.GetUserQuota(relayInfo.UserId, false) if err != nil { return types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry()) diff --git a/service/quota.go b/service/quota.go index 8b1412f..bfc89c2 100644 --- a/service/quota.go +++ b/service/quota.go @@ -47,14 +47,14 @@ func hasCustomModelRatio(modelName string, currentRatio float64) bool { return currentRatio != defaultRatio } -func calculateAudioQuota(info QuotaInfo) int { +func calculateAudioQuota(info QuotaInfo) (int, *common.QuotaClamp) { if info.UsePrice { modelPrice := decimal.NewFromFloat(info.ModelPrice) quotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) groupRatio := decimal.NewFromFloat(info.GroupRatio) quota := modelPrice.Mul(quotaPerUnit).Mul(groupRatio) - return decimalToQuota(quota) + return common.QuotaFromDecimalChecked(quota) } completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(info.ModelName)) @@ -83,7 +83,7 @@ func calculateAudioQuota(info QuotaInfo) int { quota = decimal.NewFromInt(1) } - return decimalToQuota(quota) + return common.QuotaFromDecimalChecked(quota) } func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.RealtimeUsage) error { @@ -136,7 +136,8 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag GroupRatio: actualGroupRatio, } - quota := calculateAudioQuota(quotaInfo) + quota, clamp := calculateAudioQuota(quotaInfo) + noteQuotaClamp(relayInfo, clamp) if userQuota < int64(quota) { return fmt.Errorf("user quota is not enough, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(quota)) @@ -199,7 +200,8 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod GroupRatio: groupRatio, } - quota := calculateAudioQuota(quotaInfo) + quota, clamp := calculateAudioQuota(quotaInfo) + noteQuotaClamp(relayInfo, clamp) if tieredOk { quota = tieredQuota } @@ -247,6 +249,7 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod if tieredResult != nil { InjectTieredBillingInfo(other, relayInfo, tieredResult) } + attachQuotaSaturation(ctx, relayInfo, other) model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, PromptTokens: usage.InputTokens, @@ -328,7 +331,8 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u GroupRatio: groupRatio, } - quota := calculateAudioQuota(quotaInfo) + quota, clamp := calculateAudioQuota(quotaInfo) + noteQuotaClamp(relayInfo, clamp) if tieredOk { quota = tieredQuota } @@ -376,6 +380,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u if tieredResult != nil { InjectTieredBillingInfo(other, relayInfo, tieredResult) } + attachQuotaSaturation(ctx, relayInfo, other) model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, PromptTokens: usage.PromptTokens, diff --git a/service/task_billing.go b/service/task_billing.go index 9ba0749..c5462bd 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -55,6 +55,7 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) { other["is_model_mapped"] = true other["upstream_model_name"] = info.UpstreamModelName } + attachQuotaSaturation(c, info, other) model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{ ChannelId: info.ChannelId, ModelName: info.OriginModelName, @@ -194,7 +195,7 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) { // RecalculateTaskQuota 通用的异步差额结算。 // actualQuota 是任务完成后的实际应扣额度,与预扣额度 (task.Quota) 做差额结算。 // reason 用于日志记录(例如 "token重算" 或 "adaptor调整")。 -func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int, reason string) { +func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int, reason string, clamps ...*common.QuotaClamp) { if actualQuota <= 0 { return } @@ -244,6 +245,9 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int other["task_id"] = task.TaskID other["pre_consumed_quota"] = preConsumedQuota other["actual_quota"] = actualQuota + for _, clamp := range clamps { + attachQuotaSaturationToOther(other, clamp) + } model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{ UserId: task.UserId, LogType: logType, @@ -308,8 +312,8 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo } // 计算实际应扣费额度: totalTokens * modelRatio * groupRatio * otherMultiplier - actualQuota := common.QuotaFromFloat(float64(totalTokens) * modelRatio * finalGroupRatio * otherMultiplier) + actualQuota, clamp := common.QuotaFromFloatChecked(float64(totalTokens) * modelRatio * finalGroupRatio * otherMultiplier) reason := fmt.Sprintf("token重算:tokens=%d, modelRatio=%.2f, groupRatio=%.2f, otherMultiplier=%.4f", totalTokens, modelRatio, finalGroupRatio, otherMultiplier) - RecalculateTaskQuota(ctx, task, actualQuota, reason) + RecalculateTaskQuota(ctx, task, actualQuota, reason, clamp) } diff --git a/service/text_quota.go b/service/text_quota.go index 3eced8e..4b33fe2 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -138,6 +138,17 @@ func calculateTextToolCallSurcharge(ctx *gin.Context, relayInfo *relaycommon.Rel return surcharge } +// noteQuotaClamp records the first quota saturation event onto relayInfo so it +// can later be attached to the consume/task log for admin auditing. +func noteQuotaClamp(relayInfo *relaycommon.RelayInfo, clamp *common.QuotaClamp) { + if clamp == nil || relayInfo == nil { + return + } + if relayInfo.QuotaClamp == nil { + relayInfo.QuotaClamp = clamp + } +} + func composeTieredTextQuota(relayInfo *relaycommon.RelayInfo, summary textQuotaSummary, tieredQuota int, tieredResult *billingexpr.TieredResult) int { if summary.ToolCallSurchargeQuota.IsZero() { return tieredQuota @@ -145,13 +156,17 @@ func composeTieredTextQuota(relayInfo *relaycommon.RelayInfo, summary textQuotaS if tieredResult != nil { if snap := relayInfo.TieredBillingSnapshot; snap != nil { - return decimalToQuota(decimal.NewFromFloat(tieredResult.ActualQuotaBeforeGroup). + quota, clamp := common.QuotaFromDecimalChecked(decimal.NewFromFloat(tieredResult.ActualQuotaBeforeGroup). Mul(decimal.NewFromFloat(snap.GroupRatio)). Add(summary.ToolCallSurchargeQuota)) + noteQuotaClamp(relayInfo, clamp) + return quota } } - return decimalToQuota(decimal.NewFromInt(int64(tieredQuota)).Add(summary.ToolCallSurchargeQuota)) + quota, clamp := common.QuotaFromDecimalChecked(decimal.NewFromInt(int64(tieredQuota)).Add(summary.ToolCallSurchargeQuota)) + noteQuotaClamp(relayInfo, clamp) + return quota } func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage) textQuotaSummary { @@ -184,7 +199,7 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf summary.CompletionTokens = usage.CompletionTokens summary.TotalTokens = usage.PromptTokens + usage.CompletionTokens summary.CacheTokens = usage.PromptTokensDetails.CachedTokens - summary.CacheCreationTokens = usage.PromptTokensDetails.CachedCreationTokens + summary.CacheCreationTokens = usage.PromptTokensDetails.CacheCreationTokensTotal() summary.CacheCreationTokens5m = usage.ClaudeCacheCreation5mTokens summary.CacheCreationTokens1h = usage.ClaudeCacheCreation1hTokens summary.ImageTokens = usage.PromptTokensDetails.ImageTokens @@ -270,6 +285,10 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf } } + if baseTokens.IsNegative() { + baseTokens = decimal.Zero + } + promptQuota := baseTokens.Add(cachedTokensWithRatio).Add(imageTokensWithRatio).Add(cachedCreationTokensWithRatio) completionQuota := dCompletionTokens.Mul(dCompletionRatio) quotaCalculateDecimal := promptQuota.Add(completionQuota).Mul(ratio) @@ -285,7 +304,9 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf if !ratio.IsZero() && quotaCalculateDecimal.LessThanOrEqual(decimal.Zero) { quotaCalculateDecimal = decimal.NewFromInt(1) } - summary.Quota = decimalToQuota(quotaCalculateDecimal) + quota, clamp := common.QuotaFromDecimalChecked(quotaCalculateDecimal) + summary.Quota = quota + noteQuotaClamp(relayInfo, clamp) } else { quotaCalculateDecimal := dModelPrice.Mul(dQuotaPerUnit).Mul(dGroupRatio) quotaCalculateDecimal = quotaCalculateDecimal.Add(summary.ToolCallSurchargeQuota) @@ -295,7 +316,9 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf quotaCalculateDecimal = quotaCalculateDecimal.Mul(decimal.NewFromFloat(otherRatio)) } } - summary.Quota = decimalToQuota(quotaCalculateDecimal) + quota, clamp := common.QuotaFromDecimalChecked(quotaCalculateDecimal) + summary.Quota = quota + noteQuotaClamp(relayInfo, clamp) } if summary.TotalTokens == 0 { @@ -308,8 +331,7 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf } func decimalToQuota(d decimal.Decimal) int { - f, _ := d.Round(0).Float64() - return common.QuotaFromFloat(f) + return common.QuotaFromDecimal(d) } func usageSemanticFromUsage(relayInfo *relaycommon.RelayInfo, usage *dto.Usage) string { @@ -347,15 +369,16 @@ func streamFallbackQuota(relayInfo *relaycommon.RelayInfo, quota int) (int, bool func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent []string) { originUsage := usage + billingUsage := effectiveBillingUsage(usage) if usage == nil { extraContent = append(extraContent, "上游无计费信息") } if originUsage != nil { - ObserveChannelAffinityUsageCacheByRelayFormat(ctx, usage, relayInfo.GetFinalRequestRelayFormat()) + ObserveChannelAffinityUsageCacheByRelayFormat(ctx, billingUsage, relayInfo.GetFinalRequestRelayFormat()) } adminRejectReason := common.GetContextKeyString(ctx, constant.ContextKeyAdminRejectReason) - summary := calculateTextQuotaSummary(ctx, relayInfo, usage) + summary := calculateTextQuotaSummary(ctx, relayInfo, billingUsage) var tieredResult *billingexpr.TieredResult tieredBillingApplied := false @@ -364,7 +387,7 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us if snap := relayInfo.TieredBillingSnapshot; snap != nil { tieredUsedVars = billingexpr.UsedVars(snap.ExprString) } - tieredOk, tieredQuota, tieredRes := TryTieredSettle(relayInfo, BuildTieredTokenParams(usage, summary.IsClaudeUsageSemantic, tieredUsedVars)) + tieredOk, tieredQuota, tieredRes := TryTieredSettle(relayInfo, BuildTieredTokenParams(billingUsage, summary.IsClaudeUsageSemantic, tieredUsedVars)) if tieredOk { tieredBillingApplied = true tieredResult = tieredRes @@ -436,6 +459,7 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us } else { other = GenerateTextOtherInfo(ctx, relayInfo, summary.ModelRatio, summary.GroupRatio, summary.CompletionRatio, summary.CacheTokens, summary.CacheRatio, summary.ModelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio) } + appendUsageBillingPathForLog(other, common.GetContextKeyBool(ctx, constant.ContextKeyLocalCountTokens), originUsage) if adminRejectReason != "" { other["reject_reason"] = adminRejectReason } @@ -486,16 +510,17 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us // to cache_creation_tokens. other["cache_write_tokens"] = cacheWriteTokens } - if relayInfo.GetFinalRequestRelayFormat() != types.RelayFormatClaude && usage != nil && usage.UsageSource != "" && usage.InputTokens > 0 { + if relayInfo.GetFinalRequestRelayFormat() != types.RelayFormatClaude && billingUsage != nil && billingUsage.UsageSource != "" && billingUsage.InputTokens > 0 { // input_tokens_total: explicit normalized total input used by the usage log UI. // Only write this field when upstream/current conversion has already provided a // reliable total input value and tagged the usage source. Do not infer it from // prompt/cache fields here, otherwise old upstream payloads may be double-counted. - other["input_tokens_total"] = usage.InputTokens + other["input_tokens_total"] = billingUsage.InputTokens } if tieredBillingApplied { InjectTieredBillingInfo(other, relayInfo, tieredResult) } + attachQuotaSaturation(ctx, relayInfo, other) model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, diff --git a/service/text_quota_test.go b/service/text_quota_test.go index c0305af..94f638e 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -318,6 +318,172 @@ func TestCalculateTextQuotaSummaryUsesAnthropicUsageSemanticFromUpstreamUsage(t require.Equal(t, 1488, summary.Quota) } +func TestCalculateTextQuotaSummaryUsesClaudeBillingUsageBeforeTopLevelUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(w) + + relayInfo := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + OriginModelName: "claude-3-7-sonnet", + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 2, + CacheRatio: 0.1, + CacheCreationRatio: 1.25, + CacheCreation5mRatio: 1.25, + CacheCreation1hRatio: 2, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + StartTime: time.Now(), + } + + usage := &dto.Usage{ + PromptTokens: 999, + CompletionTokens: 999, + TotalTokens: 1998, + BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{ + InputTokens: 70, + CacheReadInputTokens: 30, + CacheCreationInputTokens: 20, + OutputTokens: 7, + CacheCreation: &dto.ClaudeCacheCreationUsage{ + Ephemeral5mInputTokens: 12, + Ephemeral1hInputTokens: 8, + }, + }), + } + + summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveBillingUsage(usage)) + + require.True(t, summary.IsClaudeUsageSemantic) + require.Equal(t, dto.BillingUsageSemanticAnthropic, summary.UsageSemantic) + require.Equal(t, 70, summary.PromptTokens) + require.Equal(t, 7, summary.CompletionTokens) + require.Equal(t, 30, summary.CacheTokens) + require.Equal(t, 20, summary.CacheCreationTokens) + require.Equal(t, 12, summary.CacheCreationTokens5m) + require.Equal(t, 8, summary.CacheCreationTokens1h) + require.Equal(t, 118, summary.Quota) +} + +func TestCalculateTextQuotaSummaryUsesGeminiBillingUsageBeforeTopLevelUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(w) + + relayInfo := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + OriginModelName: "gemini-2.5-flash", + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 2, + CacheRatio: 0.1, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + StartTime: time.Now(), + } + + usage := &dto.Usage{ + PromptTokens: 999, + CompletionTokens: 999, + TotalTokens: 1998, + BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{ + PromptTokenCount: 100, + ToolUsePromptTokenCount: 5, + CandidatesTokenCount: 20, + ThoughtsTokenCount: 3, + TotalTokenCount: 128, + CachedContentTokenCount: 7, + }), + } + + summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveBillingUsage(usage)) + + require.False(t, summary.IsClaudeUsageSemantic) + require.Equal(t, dto.BillingUsageSemanticGemini, summary.UsageSemantic) + require.Equal(t, 105, summary.PromptTokens) + require.Equal(t, 23, summary.CompletionTokens) + require.Equal(t, 7, summary.CacheTokens) + require.Equal(t, 128, summary.TotalTokens) + require.Equal(t, 145, summary.Quota) +} + +func TestCalculateTextQuotaSummaryUsesOpenAIBillingUsageBeforeTopLevelUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(w) + + relayInfo := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatClaude, + OriginModelName: "gpt-4o", + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 2, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + StartTime: time.Now(), + } + + usage := &dto.Usage{ + PromptTokens: 999, + CompletionTokens: 999, + TotalTokens: 1998, + BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{ + PromptTokens: 80, + CompletionTokens: 9, + TotalTokens: 89, + }), + } + + summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveBillingUsage(usage)) + + require.False(t, summary.IsClaudeUsageSemantic) + require.Equal(t, dto.BillingUsageSemanticOpenAI, summary.UsageSemantic) + require.Equal(t, 80, summary.PromptTokens) + require.Equal(t, 9, summary.CompletionTokens) + require.Equal(t, 89, summary.TotalTokens) + require.Equal(t, 98, summary.Quota) +} + +func TestUsageBillingPathForLog(t *testing.T) { + require.Equal(t, usageBillingPathLocal, usageBillingPathForLog(true, &dto.Usage{ + BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), + })) + require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, &dto.Usage{})) + require.Equal(t, usageBillingPathOpenAI, usageBillingPathForLog(false, &dto.Usage{ + BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{PromptTokens: 1}), + })) + require.Equal(t, usageBillingPathAnthropic, usageBillingPathForLog(false, &dto.Usage{ + BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), + })) + require.Equal(t, usageBillingPathGemini, usageBillingPathForLog(false, &dto.Usage{ + BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{PromptTokenCount: 1}), + })) + require.Equal(t, usageBillingPathGeminiEstimated, usageBillingPathForLog(false, &dto.Usage{ + BillingUsage: dto.NewEstimatedGeminiChatBillingUsage(&dto.Usage{PromptTokens: 1}), + })) +} + +func TestAppendUsageBillingPathForLogWritesAdminInfo(t *testing.T) { + other := map[string]interface{}{ + "admin_info": map[string]interface{}{}, + } + appendUsageBillingPathForLog(other, false, &dto.Usage{ + BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), + }) + + adminInfo, ok := other["admin_info"].(map[string]interface{}) + require.True(t, ok) + require.Equal(t, usageBillingPathAnthropic, adminInfo["usage_billing_path"]) + + other = map[string]interface{}{} + appendUsageBillingPathForLog(other, true, nil) + adminInfo, ok = other["admin_info"].(map[string]interface{}) + require.True(t, ok) + require.Equal(t, usageBillingPathLocal, adminInfo["usage_billing_path"]) +} + func TestCacheWriteTokensTotal(t *testing.T) { t.Run("split cache creation", func(t *testing.T) { summary := textQuotaSummary{ diff --git a/service/tiered_settle.go b/service/tiered_settle.go index 28b9fae..6077952 100644 --- a/service/tiered_settle.go +++ b/service/tiered_settle.go @@ -22,7 +22,7 @@ func BuildTieredTokenParams(usage *dto.Usage, isClaudeUsageSemantic bool, usedVa p := float64(usage.PromptTokens) c := float64(usage.CompletionTokens) cr := float64(usage.PromptTokensDetails.CachedTokens) - cc5m := float64(usage.PromptTokensDetails.CachedCreationTokens) + cc5m := float64(usage.PromptTokensDetails.CacheCreationTokensTotal()) cc1h := float64(0) if usage.UsageSemantic == "anthropic" { @@ -112,5 +112,7 @@ func TryTieredSettle(relayInfo *relaycommon.RelayInfo, params billingexpr.TokenP return true, quota, nil } + noteQuotaClamp(relayInfo, tr.Clamp) + return true, tr.ActualQuotaAfterGroup, &tr } diff --git a/web/default/src/features/system-settings/billing/section-registry.tsx b/web/default/src/features/system-settings/billing/section-registry.tsx index c9e4dcf..d97ff3e 100644 --- a/web/default/src/features/system-settings/billing/section-registry.tsx +++ b/web/default/src/features/system-settings/billing/section-registry.tsx @@ -112,7 +112,7 @@ const BILLING_SECTIONS = [ modelDefaults={getModelDefaults(settings)} groupDefaults={getGroupDefaults(settings)} toolPricesDefault={settings['tool_price_setting.prices']} - visibleTabs={['models', 'tool-prices', 'upstream-sync']} + visibleTabs={['models', 'unset-models', 'tool-prices', 'upstream-sync']} /> ), }, diff --git a/web/default/src/features/system-settings/models/model-ratio-form.tsx b/web/default/src/features/system-settings/models/model-ratio-form.tsx index e76ecbf..dd6976e 100644 --- a/web/default/src/features/system-settings/models/model-ratio-form.tsx +++ b/web/default/src/features/system-settings/models/model-ratio-form.tsx @@ -16,10 +16,12 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { memo, useCallback, useState } from 'react' +import { memo, useCallback, useEffect, useState } from 'react' import { type UseFormReturn } from 'react-hook-form' +import { useQuery } from '@tanstack/react-query' import { Code2, Eye } from 'lucide-react' import { useTranslation } from 'react-i18next' +import { toast } from 'sonner' import { Button } from '@/components/ui/button' import { Form, @@ -32,6 +34,7 @@ import { } from '@/components/ui/form' import { Switch } from '@/components/ui/switch' import { Textarea } from '@/components/ui/textarea' +import { getEnabledModels } from '@/features/channels/api' import { SettingsForm, SettingsSwitchContent, @@ -56,22 +59,45 @@ type ModelFormValues = { type ModelRatioFormProps = { form: UseFormReturn + savedValues: ModelFormValues onSave: (values: ModelFormValues) => Promise onReset: () => void isSaving: boolean isResetting: boolean + variant?: 'default' | 'unset' } export const ModelRatioForm = memo(function ModelRatioForm({ form, + savedValues, onSave, onReset, isSaving, isResetting, + variant = 'default', }: ModelRatioFormProps) { const { t } = useTranslation() + const isUnsetVariant = variant === 'unset' const [editMode, setEditMode] = useState<'visual' | 'json'>('visual') + const enabledModelsQuery = useQuery({ + queryKey: ['enabled-models'], + queryFn: getEnabledModels, + enabled: isUnsetVariant, + }) + + const enabledModelsError = isUnsetVariant + ? enabledModelsQuery.isError || + (enabledModelsQuery.data !== undefined && + !enabledModelsQuery.data.success) + : false + const enabledModelsErrorMessage = enabledModelsQuery.data?.message + + useEffect(() => { + if (!enabledModelsError) return + toast.error(enabledModelsErrorMessage || t('Failed to load enabled models')) + }, [enabledModelsError, enabledModelsErrorMessage, t]) + const handleFieldChange = useCallback( (field: keyof ModelFormValues, value: string) => { form.setValue(field, value, { @@ -127,6 +153,16 @@ export const ModelRatioForm = memo(function ModelRatioForm({ {editMode === 'visual' ? (
{ const fieldMap: Record = { 'billing_setting.billing_mode': 'BillingMode', diff --git a/web/default/src/features/system-settings/models/model-ratio-visual-editor.tsx b/web/default/src/features/system-settings/models/model-ratio-visual-editor.tsx index 31d0470..9e90a83 100644 --- a/web/default/src/features/system-settings/models/model-ratio-visual-editor.tsx +++ b/web/default/src/features/system-settings/models/model-ratio-visual-editor.tsx @@ -68,6 +68,16 @@ import { import { formatPricingNumber } from './pricing-format' type ModelRatioVisualEditorProps = { + savedModelPrice: string + savedModelRatio: string + savedCacheRatio: string + savedCreateCacheRatio: string + savedCompletionRatio: string + savedImageRatio: string + savedAudioRatio: string + savedAudioCompletionRatio: string + savedBillingMode: string + savedBillingExpr: string modelPrice: string modelRatio: string cacheRatio: string @@ -78,9 +88,25 @@ type ModelRatioVisualEditorProps = { audioCompletionRatio: string billingMode: string billingExpr: string + candidateModelNames?: string[] + candidateModelsLoading?: boolean + filterMode?: 'all' | 'unset' onChange: (field: string, value: string) => void } +type ModelRatioSettingsInput = { + modelPrice: string + modelRatio: string + cacheRatio: string + createCacheRatio: string + completionRatio: string + imageRatio: string + audioRatio: string + audioCompletionRatio: string + billingMode: string + billingExpr: string +} + type ModelRow = { name: string price?: string @@ -192,8 +218,146 @@ const getPriceDetail = (row: ModelRow, t: (key: string) => string) => { return details.length > 0 ? details.join(' · ') : t('Base input price only') } +const isBasePricingUnset = (row?: ModelRow) => + !row || + (row.billingMode !== 'tiered_expr' && + !hasValue(row.price) && + !hasValue(row.ratio)) + +const buildModelRows = ({ + modelPrice, + modelRatio, + cacheRatio, + createCacheRatio, + completionRatio, + imageRatio, + audioRatio, + audioCompletionRatio, + billingMode, + billingExpr, +}: ModelRatioSettingsInput): ModelRow[] => { + const priceMap = safeJsonParse>(modelPrice, { + fallback: {}, + context: 'model prices', + }) + const ratioMap = safeJsonParse>(modelRatio, { + fallback: {}, + context: 'model ratios', + }) + const cacheMap = safeJsonParse>(cacheRatio, { + fallback: {}, + context: 'cache ratios', + }) + const createCacheMap = safeJsonParse>( + createCacheRatio, + { fallback: {}, context: 'create cache ratios' } + ) + const completionMap = safeJsonParse>(completionRatio, { + fallback: {}, + context: 'completion ratios', + }) + const imageMap = safeJsonParse>(imageRatio, { + fallback: {}, + context: 'image ratios', + }) + const audioMap = safeJsonParse>(audioRatio, { + fallback: {}, + context: 'audio ratios', + }) + const audioCompletionMap = safeJsonParse>( + audioCompletionRatio, + { fallback: {}, context: 'audio completion ratios' } + ) + const billingModeMap = safeJsonParse>(billingMode, { + fallback: {}, + context: 'billing mode', + }) + const billingExprMap = safeJsonParse>(billingExpr, { + fallback: {}, + context: 'billing expression', + }) + + const modelNames = new Set([ + ...Object.keys(priceMap), + ...Object.keys(ratioMap), + ...Object.keys(cacheMap), + ...Object.keys(createCacheMap), + ...Object.keys(completionMap), + ...Object.keys(imageMap), + ...Object.keys(audioMap), + ...Object.keys(audioCompletionMap), + ...Object.keys(billingModeMap), + ...Object.keys(billingExprMap), + ]) + + return Array.from(modelNames).map((name) => { + const price = priceMap[name]?.toString() || '' + const ratio = ratioMap[name]?.toString() || '' + const cache = cacheMap[name]?.toString() || '' + const createCache = createCacheMap[name]?.toString() || '' + const completion = completionMap[name]?.toString() || '' + const image = imageMap[name]?.toString() || '' + const audio = audioMap[name]?.toString() || '' + const audioCompletion = audioCompletionMap[name]?.toString() || '' + + const modeForModel = billingModeMap[name] + if (modeForModel === 'tiered_expr') { + const fullExpr = billingExprMap[name] || '' + const { billingExpr: pureExpr, requestRuleExpr } = + splitBillingExprAndRequestRules(fullExpr) + return { + name, + billingMode: 'tiered_expr', + billingExpr: pureExpr, + requestRuleExpr, + price, + ratio, + cacheRatio: cache, + createCacheRatio: createCache, + completionRatio: completion, + imageRatio: image, + audioRatio: audio, + audioCompletionRatio: audioCompletion, + hasConflict: false, + } + } + + return { + name, + price, + ratio, + cacheRatio: cache, + createCacheRatio: createCache, + completionRatio: completion, + imageRatio: image, + audioRatio: audio, + audioCompletionRatio: audioCompletion, + billingMode: price !== '' ? 'per-request' : 'per-token', + hasConflict: + price !== '' && + (ratio !== '' || + completion !== '' || + cache !== '' || + createCache !== '' || + image !== '' || + audio !== '' || + audioCompletion !== ''), + } + }) +} + export const ModelRatioVisualEditor = memo( function ModelRatioVisualEditor({ + savedModelPrice, + savedModelRatio, + savedCacheRatio, + savedCreateCacheRatio, + savedCompletionRatio, + savedImageRatio, + savedAudioRatio, + savedAudioCompletionRatio, + savedBillingMode, + savedBillingExpr, modelPrice, modelRatio, cacheRatio, @@ -204,6 +368,9 @@ export const ModelRatioVisualEditor = memo( audioCompletionRatio, billingMode, billingExpr, + candidateModelNames, + candidateModelsLoading, + filterMode = 'all', onChange, }: ModelRatioVisualEditorProps) { const { t } = useTranslation() @@ -259,126 +426,69 @@ export const ModelRatioVisualEditor = memo( }, [columnVisibility]) const models = useMemo(() => { - const priceMap = safeJsonParse>(modelPrice, { - fallback: {}, - context: 'model prices', - }) - const ratioMap = safeJsonParse>(modelRatio, { - fallback: {}, - context: 'model ratios', - }) - const cacheMap = safeJsonParse>(cacheRatio, { - fallback: {}, - context: 'cache ratios', + const savedRows = buildModelRows({ + modelPrice: savedModelPrice, + modelRatio: savedModelRatio, + cacheRatio: savedCacheRatio, + createCacheRatio: savedCreateCacheRatio, + completionRatio: savedCompletionRatio, + imageRatio: savedImageRatio, + audioRatio: savedAudioRatio, + audioCompletionRatio: savedAudioCompletionRatio, + billingMode: savedBillingMode, + billingExpr: savedBillingExpr, }) - const createCacheMap = safeJsonParse>( + const draftRows = buildModelRows({ + modelPrice, + modelRatio, + cacheRatio, createCacheRatio, - { fallback: {}, context: 'create cache ratios' } - ) - const completionMap = safeJsonParse>( completionRatio, - { fallback: {}, context: 'completion ratios' } - ) - const imageMap = safeJsonParse>(imageRatio, { - fallback: {}, - context: 'image ratios', - }) - const audioMap = safeJsonParse>(audioRatio, { - fallback: {}, - context: 'audio ratios', - }) - const audioCompletionMap = safeJsonParse>( + imageRatio, + audioRatio, audioCompletionRatio, - { fallback: {}, context: 'audio completion ratios' } - ) - const billingModeMap = safeJsonParse>( billingMode, - { - fallback: {}, - context: 'billing mode', - } - ) - const billingExprMap = safeJsonParse>( billingExpr, - { - fallback: {}, - context: 'billing expression', - } - ) - - const modelNames = new Set([ - ...Object.keys(priceMap), - ...Object.keys(ratioMap), - ...Object.keys(cacheMap), - ...Object.keys(createCacheMap), - ...Object.keys(completionMap), - ...Object.keys(imageMap), - ...Object.keys(audioMap), - ...Object.keys(audioCompletionMap), - ...Object.keys(billingModeMap), - ...Object.keys(billingExprMap), - ]) - - const modelData: ModelRow[] = Array.from(modelNames).map((name) => { - const price = priceMap[name]?.toString() || '' - const ratio = ratioMap[name]?.toString() || '' - const cache = cacheMap[name]?.toString() || '' - const createCache = createCacheMap[name]?.toString() || '' - const completion = completionMap[name]?.toString() || '' - const image = imageMap[name]?.toString() || '' - const audio = audioMap[name]?.toString() || '' - const audioCompletion = audioCompletionMap[name]?.toString() || '' - - const modeForModel = billingModeMap[name] - if (modeForModel === 'tiered_expr') { - // Tiered_expr models may also retain ratio/price values as fallback - // during multi-instance sync delays. We preserve them in the row so - // the edit dialog round-trip and the next save don't drop them. - const fullExpr = billingExprMap[name] || '' - const { billingExpr: pureExpr, requestRuleExpr } = - splitBillingExprAndRequestRules(fullExpr) - return { - name, - billingMode: 'tiered_expr', - billingExpr: pureExpr, - requestRuleExpr, - price, - ratio, - cacheRatio: cache, - createCacheRatio: createCache, - completionRatio: completion, - imageRatio: image, - audioRatio: audio, - audioCompletionRatio: audioCompletion, - hasConflict: false, - } - } - - return { - name, - price, - ratio, - cacheRatio: cache, - createCacheRatio: createCache, - completionRatio: completion, - imageRatio: image, - audioRatio: audio, - audioCompletionRatio: audioCompletion, - billingMode: price !== '' ? 'per-request' : 'per-token', - hasConflict: - price !== '' && - (ratio !== '' || - completion !== '' || - cache !== '' || - createCache !== '' || - image !== '' || - audio !== '' || - audioCompletion !== ''), - } }) + const savedByName = new Map(savedRows.map((row) => [row.name, row])) + const draftByName = new Map(draftRows.map((row) => [row.name, row])) + const modelNames = + filterMode === 'unset' + ? new Set(candidateModelNames ?? []) + : new Set([...savedByName.keys(), ...draftByName.keys()]) - return modelData.sort((a, b) => a.name.localeCompare(b.name)) + return Array.from(modelNames) + .map((name) => { + const draft = draftByName.get(name) + const saved = savedByName.get(name) + return ( + draft ?? + saved ?? { + name, + billingMode: 'per-token', + hasConflict: false, + } + ) + }) + .filter( + (row) => + filterMode !== 'unset' || + isBasePricingUnset(savedByName.get(row.name)) + ) + .sort((a, b) => a.name.localeCompare(b.name)) }, [ + candidateModelNames, + filterMode, + savedModelPrice, + savedModelRatio, + savedCacheRatio, + savedCreateCacheRatio, + savedCompletionRatio, + savedImageRatio, + savedAudioRatio, + savedAudioCompletionRatio, + savedBillingMode, + savedBillingExpr, modelPrice, modelRatio, cacheRatio, @@ -664,19 +774,21 @@ export const ModelRatioVisualEditor = memo( > - + {filterMode !== 'unset' && ( + + )}
), enableHiding: false, }, ] - }, [handleEdit, handleDelete, t]) + }, [filterMode, handleEdit, handleDelete, t]) const table = useReactTable({ data: models, @@ -915,10 +1027,12 @@ export const ModelRatioVisualEditor = memo( }, ]} preActions={ - + filterMode === 'unset' ? undefined : ( + + ) } /> @@ -926,7 +1040,11 @@ export const ModelRatioVisualEditor = memo(
{table.getState().globalFilter ? t('No models match your search') - : t('No models configured. Use Add model to get started.')} + : filterMode === 'unset' + ? candidateModelsLoading + ? t('Loading...') + : t('No models with unset prices') + : t('No models configured. Use Add model to get started.')}
) : (
@@ -1005,10 +1123,12 @@ export const ModelRatioVisualEditor = memo( 'Use the full-width table to scan prices, then select a row to edit it here.' )}

- + {filterMode !== 'unset' && ( + + )}
)} @@ -1039,6 +1159,17 @@ export const ModelRatioVisualEditor = memo( // Custom equality check - only re-render if JSON props actually changed (prevProps, nextProps) => { return ( + prevProps.savedModelPrice === nextProps.savedModelPrice && + prevProps.savedModelRatio === nextProps.savedModelRatio && + prevProps.savedCacheRatio === nextProps.savedCacheRatio && + prevProps.savedCreateCacheRatio === nextProps.savedCreateCacheRatio && + prevProps.savedCompletionRatio === nextProps.savedCompletionRatio && + prevProps.savedImageRatio === nextProps.savedImageRatio && + prevProps.savedAudioRatio === nextProps.savedAudioRatio && + prevProps.savedAudioCompletionRatio === + nextProps.savedAudioCompletionRatio && + prevProps.savedBillingMode === nextProps.savedBillingMode && + prevProps.savedBillingExpr === nextProps.savedBillingExpr && prevProps.modelPrice === nextProps.modelPrice && prevProps.modelRatio === nextProps.modelRatio && prevProps.cacheRatio === nextProps.cacheRatio && @@ -1049,6 +1180,9 @@ export const ModelRatioVisualEditor = memo( prevProps.audioCompletionRatio === nextProps.audioCompletionRatio && prevProps.billingMode === nextProps.billingMode && prevProps.billingExpr === nextProps.billingExpr && + prevProps.candidateModelNames === nextProps.candidateModelNames && + prevProps.candidateModelsLoading === nextProps.candidateModelsLoading && + prevProps.filterMode === nextProps.filterMode && prevProps.onChange === nextProps.onChange ) } diff --git a/web/default/src/features/system-settings/models/ratio-settings-card.tsx b/web/default/src/features/system-settings/models/ratio-settings-card.tsx index 6f48857..7377f00 100644 --- a/web/default/src/features/system-settings/models/ratio-settings-card.tsx +++ b/web/default/src/features/system-settings/models/ratio-settings-card.tsx @@ -214,7 +214,12 @@ const groupSchema = z.object({ type ModelFormValues = z.infer type GroupFormValues = z.infer -type RatioTabId = 'models' | 'groups' | 'tool-prices' | 'upstream-sync' +type RatioTabId = + | 'models' + | 'unset-models' + | 'groups' + | 'tool-prices' + | 'upstream-sync' function normalizeAutoGroupRoutesString( autoGroupRoutes: string, @@ -271,7 +276,7 @@ export function RatioSettingsCard({ }, }) - const modelNormalizedDefaults = useRef({ + const initialModelNormalizedDefaults = { ModelPrice: normalizeJsonString(modelDefaults.ModelPrice), ModelRatio: normalizeJsonString(modelDefaults.ModelRatio), CacheRatio: normalizeJsonString(modelDefaults.CacheRatio), @@ -285,7 +290,11 @@ export function RatioSettingsCard({ ExposeRatioEnabled: modelDefaults.ExposeRatioEnabled, BillingMode: normalizeJsonString(modelDefaults.BillingMode), BillingExpr: normalizeJsonString(modelDefaults.BillingExpr), - }) + } + const modelNormalizedDefaults = useRef(initialModelNormalizedDefaults) + const [savedModelValues, setSavedModelValues] = useState( + initialModelNormalizedDefaults + ) const groupNormalizedDefaults = useRef({ GroupRatio: normalizeJsonString(groupDefaults.GroupRatio), @@ -359,6 +368,7 @@ export function RatioSettingsCard({ BillingMode: normalizeJsonString(modelDefaults.BillingMode), BillingExpr: normalizeJsonString(modelDefaults.BillingExpr), } + setSavedModelValues(modelNormalizedDefaults.current) modelForm.reset({ ...modelDefaults, @@ -447,6 +457,9 @@ export function RatioSettingsCard({ const apiKey = apiKeyMap[key as string] || (key as string) await updateOption.mutateAsync({ key: apiKey, value: normalized[key] }) } + + modelNormalizedDefaults.current = normalized + setSavedModelValues(normalized) }, [t, updateOption] ) @@ -510,6 +523,7 @@ export function RatioSettingsCard({ const tabLabels: Record = { models: 'Model prices', + 'unset-models': 'Unset price models', groups: 'Group ratios', 'tool-prices': 'Tool prices', 'upstream-sync': 'Upstream price sync', @@ -520,18 +534,21 @@ export function RatioSettingsCard({ 2: 'grid-cols-2', 3: 'grid-cols-3', 4: 'grid-cols-4', + 5: 'grid-cols-5', }[visibleTabs.length] ?? 'grid-cols-4' const defaultTab = visibleTabs[0] ?? 'models' const renderTabContent = (tab: RatioTabId) => { - if (tab === 'models') { + if (tab === 'models' || tab === 'unset-models') { return ( ) } diff --git a/web/default/src/i18n/locales/en.json b/web/default/src/i18n/locales/en.json index 591590c..829150d 100644 --- a/web/default/src/i18n/locales/en.json +++ b/web/default/src/i18n/locales/en.json @@ -1857,6 +1857,7 @@ "Failed to load": "Failed to load", "Failed to load API keys": "Failed to load API keys", "Failed to load billing history": "Failed to load billing history", + "Failed to load enabled models": "Failed to load enabled models", "Failed to load home page content": "Failed to load home page content", "Failed to load image": "Failed to load image", "Failed to load key status": "Failed to load key status", @@ -2962,6 +2963,7 @@ "No models to add": "No models to add", "No models to copy": "No models to copy", "No models to remove": "No models to remove", + "No models with unset prices": "No models with unset prices", "No new models to add": "No new models to add", "No new models yet": "No new models yet", "No notable climbers right now": "No notable climbers right now", @@ -4810,6 +4812,7 @@ "Unlimited Quota": "Unlimited Quota", "Unsaved changes": "Unsaved changes", "Unset price": "Unset price", + "Unset price models": "Unset price models", "unsupported auto group routes config version: {{version}}": "unsupported auto group routes config version: {{version}}", "Until": "Until", "Untitled": "Untitled", diff --git a/web/default/src/i18n/locales/fr.json b/web/default/src/i18n/locales/fr.json index 0a7d14b..f40c214 100644 --- a/web/default/src/i18n/locales/fr.json +++ b/web/default/src/i18n/locales/fr.json @@ -1857,6 +1857,7 @@ "Failed to load": "Échec du chargement", "Failed to load API keys": "Échec du chargement des Clés API", "Failed to load billing history": "Échec du chargement de l'historique de facturation", + "Failed to load enabled models": "Échec du chargement des modèles activés", "Failed to load home page content": "Échec du chargement du contenu de la page d'accueil", "Failed to load image": "Échec du chargement de l'image", "Failed to load key status": "Échec du chargement du statut des clés", @@ -2962,6 +2963,7 @@ "No models to add": "Aucun modèle à ajouter", "No models to copy": "Aucun modèle à copier", "No models to remove": "Aucun modèle à supprimer", + "No models with unset prices": "Aucun modèle avec prix non défini", "No new models to add": "Aucun nouveau modèle à ajouter", "No new models yet": "Pas encore de nouveaux modèles", "No notable climbers right now": "Aucune progression notable pour le moment", @@ -4810,6 +4812,7 @@ "Unlimited Quota": "Quota illimité", "Unsaved changes": "Modifications non enregistrées", "Unset price": "Prix non défini", + "Unset price models": "Modèles sans prix défini", "unsupported auto group routes config version: {{version}}": "Version de configuration des routes auto non prise en charge : {{version}}", "Until": "Jusqu'au", "Untitled": "Sans titre", diff --git a/web/default/src/i18n/locales/ja.json b/web/default/src/i18n/locales/ja.json index c2d714e..91e7087 100644 --- a/web/default/src/i18n/locales/ja.json +++ b/web/default/src/i18n/locales/ja.json @@ -1857,6 +1857,7 @@ "Failed to load": "読み込みに失敗しました", "Failed to load API keys": "APIキーの読み込みに失敗しました", "Failed to load billing history": "請求履歴の読み込みに失敗しました", + "Failed to load enabled models": "有効なモデルの読み込みに失敗しました", "Failed to load home page content": "ホームページの内容の読み込みに失敗しました", "Failed to load image": "画像の読み込みに失敗しました", "Failed to load key status": "キー状態の読み込みに失敗しました", @@ -2962,6 +2963,7 @@ "No models to add": "追加するモデルがありません", "No models to copy": "コピーするモデルがありません", "No models to remove": "削除するモデルがありません", + "No models with unset prices": "価格未設定のモデルはありません", "No new models to add": "追加する新しいモデルはありません", "No new models yet": "新しいモデルはまだありません", "No notable climbers right now": "現在、目立った上昇はありません", @@ -4810,6 +4812,7 @@ "Unlimited Quota": "無制限のクォータ", "Unsaved changes": "未保存の変更", "Unset price": "価格未設定", + "Unset price models": "価格未設定モデル", "unsupported auto group routes config version: {{version}}": "サポートされていない自動ルート設定バージョンです: {{version}}", "Until": "まで", "Untitled": "無題", diff --git a/web/default/src/i18n/locales/ru.json b/web/default/src/i18n/locales/ru.json index 9aed47b..53b9524 100644 --- a/web/default/src/i18n/locales/ru.json +++ b/web/default/src/i18n/locales/ru.json @@ -1857,6 +1857,7 @@ "Failed to load": "Не удалось загрузить", "Failed to load API keys": "Не удалось загрузить API ключи", "Failed to load billing history": "Не удалось загрузить историю платежей", + "Failed to load enabled models": "Не удалось загрузить включенные модели", "Failed to load home page content": "Не удалось загрузить содержимое главной страницы", "Failed to load image": "Не удалось загрузить изображение", "Failed to load key status": "Не удалось загрузить статус ключей", @@ -2962,6 +2963,7 @@ "No models to add": "Нет моделей для добавления", "No models to copy": "Нет моделей для копирования", "No models to remove": "Нет моделей для удаления", + "No models with unset prices": "Нет моделей без заданной цены", "No new models to add": "Нет новых моделей для добавления", "No new models yet": "Новых моделей пока нет", "No notable climbers right now": "Сейчас нет значимых поднявшихся моделей", @@ -4810,6 +4812,7 @@ "Unlimited Quota": "Неограниченная квота", "Unsaved changes": "Несохранённые изменения", "Unset price": "Цена не задана", + "Unset price models": "Модели без заданной цены", "unsupported auto group routes config version: {{version}}": "Неподдерживаемая версия конфигурации автомаршрутов: {{version}}", "Until": "До", "Untitled": "Без названия", diff --git a/web/default/src/i18n/locales/vi.json b/web/default/src/i18n/locales/vi.json index fe86a05..09fe50e 100644 --- a/web/default/src/i18n/locales/vi.json +++ b/web/default/src/i18n/locales/vi.json @@ -1857,6 +1857,7 @@ "Failed to load": "Tải thất bại", "Failed to load API keys": "Không thể tải khóa API", "Failed to load billing history": "Không thể tải lịch sử thanh toán", + "Failed to load enabled models": "Không tải được các mô hình đã bật", "Failed to load home page content": "Không thể tải nội dung trang chủ", "Failed to load image": "Không thể tải ảnh", "Failed to load key status": "Không thể tải trạng thái khóa", @@ -2962,6 +2963,7 @@ "No models to add": "Không có mô hình để thêm", "No models to copy": "Không có mô hình nào để sao chép", "No models to remove": "Không có mô hình để xóa", + "No models with unset prices": "Không có mô hình nào chưa đặt giá", "No new models to add": "Không có mô hình mới để thêm", "No new models yet": "Chưa có mô hình mới", "No notable climbers right now": "Hiện chưa có mô hình nào tăng đáng kể", @@ -4810,6 +4812,7 @@ "Unlimited Quota": "Hạn mức không giới hạn", "Unsaved changes": "Thay đổi chưa được lưu", "Unset price": "Chưa đặt giá", + "Unset price models": "Mô hình chưa đặt giá", "unsupported auto group routes config version: {{version}}": "Phiên bản cấu hình tuyến tự động không được hỗ trợ: {{version}}", "Until": "Cho đến", "Untitled": "Không có tiêu đề", diff --git a/web/default/src/i18n/locales/zh.json b/web/default/src/i18n/locales/zh.json index cef2175..d56c176 100644 --- a/web/default/src/i18n/locales/zh.json +++ b/web/default/src/i18n/locales/zh.json @@ -1857,6 +1857,7 @@ "Failed to load": "加载失败", "Failed to load API keys": "加载 API 密钥失败", "Failed to load billing history": "加载计费历史失败", + "Failed to load enabled models": "加载已启用模型失败", "Failed to load home page content": "加载首页内容失败", "Failed to load image": "无法加载图像", "Failed to load key status": "加载密钥状态失败", @@ -2962,6 +2963,7 @@ "No models to add": "无待新增模型", "No models to copy": "没有模型可复制", "No models to remove": "无待删除模型", + "No models with unset prices": "没有未设置价格的模型", "No new models to add": "没有新模型可添加", "No new models yet": "暂无新模型", "No notable climbers right now": "当前没有显著上升的模型", @@ -4810,6 +4812,7 @@ "Unlimited Quota": "无限配额", "Unsaved changes": "未保存的更改", "Unset price": "未设置价格", + "Unset price models": "未设置价格模型", "unsupported auto group routes config version: {{version}}": "不支持的自动链路配置版本:{{version}}", "Until": "至", "Untitled": "未命名", From 5c9f603a4411313ed1c3a2dba8e62242cdbe2c08 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sun, 12 Jul 2026 18:54:33 +0800 Subject: [PATCH 2/9] v1.0.4 --- common/quota_math.go | 4 +- common/quota_math_test.go | 22 ++++++ service/billing.go | 13 +++- service/billing_usage.go | 25 +++++++ service/log_info_generate.go | 10 ++- service/log_info_generate_test.go | 73 +++++++++++++++++++ service/pre_consume_quota.go | 7 +- service/pre_consume_validation_test.go | 61 ++++++++++++++++ service/text_quota_test.go | 60 +++++++++++++++ .../models/model-ratio-visual-editor.tsx | 20 +++-- 10 files changed, 275 insertions(+), 20 deletions(-) create mode 100644 service/pre_consume_validation_test.go diff --git a/common/quota_math.go b/common/quota_math.go index c7c05ec..77dc17e 100644 --- a/common/quota_math.go +++ b/common/quota_math.go @@ -67,9 +67,9 @@ func saturateQuota(value float64, op string) (int, *QuotaClamp) { switch { case math.IsNaN(value): clamp = &QuotaClamp{Op: op, Kind: QuotaClampNaN, Original: value, Clamped: 0} - case value >= MaxQuota: + case value > MaxQuota: clamp = &QuotaClamp{Op: op, Kind: QuotaClampOverflow, Original: value, Clamped: MaxQuota} - case value <= MinQuota: + case value < MinQuota: clamp = &QuotaClamp{Op: op, Kind: QuotaClampUnderflow, Original: value, Clamped: MinQuota} default: return int(value), nil diff --git a/common/quota_math_test.go b/common/quota_math_test.go index 0e5f538..18eee6f 100644 --- a/common/quota_math_test.go +++ b/common/quota_math_test.go @@ -83,6 +83,28 @@ func TestQuotaFromFloatStrictReturnsTypedClampError(t *testing.T) { assert.ErrorContains(t, err, "clamped=2147483647") } +func TestQuotaFromFloatStrictAcceptsExactInt32Boundaries(t *testing.T) { + quota, err := QuotaFromFloatStrict(float64(MaxQuota)) + require.NoError(t, err) + assert.Equal(t, MaxQuota, quota) + + quota, err = QuotaFromFloatStrict(float64(MinQuota)) + require.NoError(t, err) + assert.Equal(t, MinQuota, quota) + + quota, err = QuotaFromFloatStrict(float64(MaxQuota) + 1) + assert.Zero(t, quota) + var overflow *QuotaClamp + require.ErrorAs(t, err, &overflow) + assert.Equal(t, QuotaClampOverflow, overflow.Kind) + + quota, err = QuotaFromFloatStrict(float64(MinQuota) - 1) + assert.Zero(t, quota) + var underflow *QuotaClamp + require.ErrorAs(t, err, &underflow) + assert.Equal(t, QuotaClampUnderflow, underflow.Kind) +} + func TestQuotaClampAuditMapIsJSONSafe(t *testing.T) { tests := []struct { name string diff --git a/service/billing.go b/service/billing.go index 67d1156..6c62de5 100644 --- a/service/billing.go +++ b/service/billing.go @@ -15,9 +15,7 @@ const ( BillingSourceSubscription = "subscription" ) -// PreConsumeBilling 根据用户计费偏好创建 BillingSession 并执行预扣费。 -// 会话存储在 relayInfo.Billing 上,供后续 Settle / Refund 使用。 -func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { +func validatePreConsumedQuota(preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { if relayInfo != nil && relayInfo.QuotaClamp != nil { return types.NewErrorWithStatusCode( relayInfo.QuotaClamp, @@ -34,6 +32,15 @@ func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycom types.ErrOptionWithSkipRetry(), ) } + return nil +} + +// PreConsumeBilling 根据用户计费偏好创建 BillingSession 并执行预扣费。 +// 会话存储在 relayInfo.Billing 上,供后续 Settle / Refund 使用。 +func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { + if apiErr := validatePreConsumedQuota(preConsumedQuota, relayInfo); apiErr != nil { + return apiErr + } session, apiErr := NewBillingSession(c, relayInfo, preConsumedQuota) if apiErr != nil { return apiErr diff --git a/service/billing_usage.go b/service/billing_usage.go index a83d5e0..cfb5ce3 100644 --- a/service/billing_usage.go +++ b/service/billing_usage.go @@ -117,12 +117,37 @@ func usageFromOpenAIBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage { if usage.TotalTokens == 0 { usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } + fillPromptTokenDetailsFromInputDetails(&usage.PromptTokensDetails, usage.InputTokensDetails) usage.UsageSemantic = dto.BillingUsageSemanticOpenAI usage.UsageSource = billingUsage.Source usage.BillingUsage = dto.CloneBillingUsage(billingUsage) return &usage } +func fillPromptTokenDetailsFromInputDetails(promptDetails *dto.InputTokenDetails, inputDetails *dto.InputTokenDetails) { + if inputDetails == nil { + return + } + if promptDetails.CachedTokens == 0 { + promptDetails.CachedTokens = inputDetails.CachedTokens + } + if promptDetails.CachedCreationTokens == 0 { + promptDetails.CachedCreationTokens = inputDetails.CachedCreationTokens + } + if promptDetails.CacheWriteTokens == 0 { + promptDetails.CacheWriteTokens = inputDetails.CacheWriteTokens + } + if promptDetails.TextTokens == 0 { + promptDetails.TextTokens = inputDetails.TextTokens + } + if promptDetails.ImageTokens == 0 { + promptDetails.ImageTokens = inputDetails.ImageTokens + } + if promptDetails.AudioTokens == 0 { + promptDetails.AudioTokens = inputDetails.AudioTokens + } +} + func usageFromClaudeBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage { claudeUsage := billingUsage.ClaudeUsage cacheCreation5m := claudeUsage.GetCacheCreation5mTokens() diff --git a/service/log_info_generate.go b/service/log_info_generate.go index 8260e71..bc9d73d 100644 --- a/service/log_info_generate.go +++ b/service/log_info_generate.go @@ -20,8 +20,12 @@ func attachQuotaSaturationToOther(other map[string]interface{}, clamp *common.Qu if clamp == nil || other == nil { return } - adminInfo, ok := other["admin_info"].(map[string]interface{}) - if !ok || adminInfo == nil { + rawAdminInfo, exists := other["admin_info"] + adminInfo, ok := rawAdminInfo.(map[string]interface{}) + if exists && !ok { + return + } + if !exists || adminInfo == nil { adminInfo = map[string]interface{}{} other["admin_info"] = adminInfo } @@ -29,7 +33,7 @@ func attachQuotaSaturationToOther(other map[string]interface{}, clamp *common.Qu } func attachQuotaSaturation(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, other map[string]interface{}) { - if relayInfo == nil || relayInfo.QuotaClamp == nil { + if relayInfo == nil || relayInfo.QuotaClamp == nil || other == nil { return } clamp := relayInfo.QuotaClamp diff --git a/service/log_info_generate_test.go b/service/log_info_generate_test.go index 1641c8e..38baff8 100644 --- a/service/log_info_generate_test.go +++ b/service/log_info_generate_test.go @@ -1,10 +1,12 @@ package service import ( + "bytes" "net/http/httptest" "testing" "time" + "github.com/MAX-API-Next/MAX-API/common" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -48,3 +50,74 @@ func TestGenerateTextOtherInfoDoesNotMarkSingleChannelAsRetry(t *testing.T) { require.NotContains(t, other, "retry_log") } + +func TestAttachQuotaSaturationToOtherPreservesMalformedAdminInfo(t *testing.T) { + clamp := &common.QuotaClamp{ + Op: "test", + Kind: common.QuotaClampOverflow, + Original: float64(common.MaxQuota) + 1, + Clamped: common.MaxQuota, + } + other := map[string]interface{}{ + "admin_info": "malformed", + } + + attachQuotaSaturationToOther(other, clamp) + + require.Equal(t, "malformed", other["admin_info"]) + require.NotContains(t, other, "quota_saturation") +} + +func TestAttachQuotaSaturationToOtherWritesValidAdminInfo(t *testing.T) { + clamp := &common.QuotaClamp{ + Op: "test", + Kind: common.QuotaClampOverflow, + Original: float64(common.MaxQuota) + 1, + Clamped: common.MaxQuota, + } + + other := map[string]interface{}{} + attachQuotaSaturationToOther(other, clamp) + adminInfo, ok := other["admin_info"].(map[string]interface{}) + require.True(t, ok) + require.Equal(t, clamp.AuditMap(), adminInfo["quota_saturation"]) + + adminInfo = map[string]interface{}{ + "existing": "value", + } + other = map[string]interface{}{ + "admin_info": adminInfo, + } + attachQuotaSaturationToOther(other, clamp) + require.Equal(t, "value", adminInfo["existing"]) + require.Equal(t, clamp.AuditMap(), adminInfo["quota_saturation"]) +} + +func TestAttachQuotaSaturationSkipsNilOtherWithoutWarning(t *testing.T) { + gin.SetMode(gin.TestMode) + var logBuffer bytes.Buffer + common.LogWriterMu.Lock() + oldWriter := gin.DefaultErrorWriter + gin.DefaultErrorWriter = &logBuffer + common.LogWriterMu.Unlock() + t.Cleanup(func() { + common.LogWriterMu.Lock() + gin.DefaultErrorWriter = oldWriter + common.LogWriterMu.Unlock() + }) + + relayInfo := &relaycommon.RelayInfo{ + QuotaClamp: &common.QuotaClamp{ + Op: "test", + Kind: common.QuotaClampOverflow, + Original: float64(common.MaxQuota) + 1, + Clamped: common.MaxQuota, + }, + } + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + + attachQuotaSaturation(ctx, relayInfo, nil) + + require.NotContains(t, logBuffer.String(), "quota saturation on consume log") +} diff --git a/service/pre_consume_quota.go b/service/pre_consume_quota.go index cee7730..448fac0 100644 --- a/service/pre_consume_quota.go +++ b/service/pre_consume_quota.go @@ -31,11 +31,8 @@ func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo) { // PreConsumeQuota checks if the user has enough quota to pre-consume. // It returns the pre-consumed quota if successful, or an error if not. func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { - if relayInfo != nil && relayInfo.QuotaClamp != nil { - return types.NewErrorWithStatusCode(relayInfo.QuotaClamp, types.ErrorCodeModelPriceError, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) - } - if preConsumedQuota < 0 { - return types.NewErrorWithStatusCode(fmt.Errorf("pre-consume quota cannot be negative: %d", preConsumedQuota), types.ErrorCodeModelPriceError, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + if apiErr := validatePreConsumedQuota(preConsumedQuota, relayInfo); apiErr != nil { + return apiErr } userQuota, err := model.GetUserQuota(relayInfo.UserId, false) if err != nil { diff --git a/service/pre_consume_validation_test.go b/service/pre_consume_validation_test.go new file mode 100644 index 0000000..082bc8b --- /dev/null +++ b/service/pre_consume_validation_test.go @@ -0,0 +1,61 @@ +package service + +import ( + "net/http" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + + "github.com/stretchr/testify/require" +) + +func TestPreConsumeValidationSharedByBillingAndQuota(t *testing.T) { + entries := []struct { + name string + call func(int, *relaycommon.RelayInfo) *types.MaxAPIError + }{ + { + name: "billing session", + call: func(preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { + return PreConsumeBilling(nil, preConsumedQuota, relayInfo) + }, + }, + { + name: "legacy quota", + call: func(preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { + return PreConsumeQuota(nil, preConsumedQuota, relayInfo) + }, + }, + } + + for _, entry := range entries { + t.Run(entry.name+"/quota_clamp", func(t *testing.T) { + clamp := &common.QuotaClamp{ + Op: "test", + Kind: common.QuotaClampOverflow, + Original: float64(common.MaxQuota) + 1, + Clamped: common.MaxQuota, + } + + apiErr := entry.call(1, &relaycommon.RelayInfo{QuotaClamp: clamp}) + + require.NotNil(t, apiErr) + require.Equal(t, clamp, apiErr.Err) + require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode()) + require.Equal(t, http.StatusBadRequest, apiErr.StatusCode) + require.True(t, types.IsSkipRetryError(apiErr)) + }) + + t.Run(entry.name+"/negative_quota", func(t *testing.T) { + apiErr := entry.call(-1, &relaycommon.RelayInfo{}) + + require.NotNil(t, apiErr) + require.ErrorContains(t, apiErr.Err, "pre-consume quota cannot be negative: -1") + require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode()) + require.Equal(t, http.StatusBadRequest, apiErr.StatusCode) + require.True(t, types.IsSkipRetryError(apiErr)) + }) + } +} diff --git a/service/text_quota_test.go b/service/text_quota_test.go index 94f638e..5890a4e 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -11,6 +11,7 @@ import ( "github.com/MAX-API-Next/MAX-API/model" "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/service/openaicompat" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" @@ -446,6 +447,65 @@ func TestCalculateTextQuotaSummaryUsesOpenAIBillingUsageBeforeTopLevelUsage(t *t require.Equal(t, 98, summary.Quota) } +func TestCalculateTextQuotaSummaryPreservesResponsesInputDetailsThroughBillingUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(w) + + chatResp, usage, err := openaicompat.ResponsesResponseToChatCompletionsResponse(&dto.OpenAIResponsesResponse{ + ID: "resp_1", + Model: "gpt-test", + CreatedAt: 456, + Status: []byte(`"completed"`), + Usage: &dto.Usage{ + InputTokens: 100, + OutputTokens: 10, + TotalTokens: 110, + InputTokensDetails: &dto.InputTokenDetails{ + CachedTokens: 5, + CacheWriteTokens: 11, + ImageTokens: 7, + AudioTokens: 13, + }, + }, + Output: []dto.ResponsesOutput{{ + Type: "message", + Role: "assistant", + Content: []dto.ResponsesOutputContent{{ + Type: "output_text", + Text: "done", + }}, + }}, + }, "chatcmpl_1") + require.NoError(t, err) + require.Equal(t, usage.PromptTokensDetails, chatResp.Usage.PromptTokensDetails) + + relayInfo := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + OriginModelName: "gpt-test", + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 2, + CacheRatio: 0, + CacheCreationRatio: 4, + ImageRatio: 3, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + StartTime: time.Now(), + } + + summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveBillingUsage(usage)) + + require.Equal(t, dto.BillingUsageSemanticOpenAI, summary.UsageSemantic) + require.Equal(t, 100, summary.PromptTokens) + require.Equal(t, 10, summary.CompletionTokens) + require.Equal(t, 5, summary.CacheTokens) + require.Equal(t, 11, summary.CacheCreationTokens) + require.Equal(t, 7, summary.ImageTokens) + require.Equal(t, 13, summary.AudioTokens) + require.Equal(t, 162, summary.Quota) +} + func TestUsageBillingPathForLog(t *testing.T) { require.Equal(t, usageBillingPathLocal, usageBillingPathForLog(true, &dto.Usage{ BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), diff --git a/web/default/src/features/system-settings/models/model-ratio-visual-editor.tsx b/web/default/src/features/system-settings/models/model-ratio-visual-editor.tsx index 9e90a83..a78f4df 100644 --- a/web/default/src/features/system-settings/models/model-ratio-visual-editor.tsx +++ b/web/default/src/features/system-settings/models/model-ratio-visual-editor.tsx @@ -995,6 +995,18 @@ export const ModelRatioVisualEditor = memo( }, [editData, persistPricingData, t, table]) const selectedTargetCount = table.getFilteredSelectedRowModel().rows.length + const emptyStateText = (() => { + if (table.getState().globalFilter) { + return t('No models match your search') + } + if (filterMode === 'unset') { + if (candidateModelsLoading) { + return t('Loading...') + } + return t('No models with unset prices') + } + return t('No models configured. Use Add model to get started.') + })() return (
@@ -1038,13 +1050,7 @@ export const ModelRatioVisualEditor = memo( {table.getRowModel().rows.length === 0 ? (
- {table.getState().globalFilter - ? t('No models match your search') - : filterMode === 'unset' - ? candidateModelsLoading - ? t('Loading...') - : t('No models with unset prices') - : t('No models configured. Use Add model to get started.')} + {emptyStateText}
) : (
From 5291db23410b271e71ab7ff6b11db59d61969b1b Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sun, 12 Jul 2026 21:14:59 +0800 Subject: [PATCH 3/9] v1.0.4 --- common/verification.go | 46 ++++++++- common/verification_test.go | 94 ++++++++++++++++++- controller/misc.go | 8 +- controller/user.go | 17 +++- controller/user_setting_test.go | 79 ++++++++++++++++ dto/billing_usage_test.go | 40 ++++++++ dto/openai_response.go | 24 +++++ service/billing_usage.go | 22 +---- service/openaicompat/responses_compat_test.go | 2 + service/openaicompat/responses_to_chat.go | 9 +- .../users/components/user-quota-dialog.tsx | 24 ++++- .../features/users/lib/quota-safety.test.ts | 54 +++++++++++ .../src/features/users/lib/quota-safety.ts | 39 ++++++++ web/default/src/i18n/locales/en.json | 1 + web/default/src/i18n/locales/fr.json | 1 + web/default/src/i18n/locales/ja.json | 1 + web/default/src/i18n/locales/ru.json | 1 + web/default/src/i18n/locales/vi.json | 1 + web/default/src/i18n/locales/zh.json | 1 + 19 files changed, 421 insertions(+), 43 deletions(-) create mode 100644 web/default/src/features/users/lib/quota-safety.test.ts create mode 100644 web/default/src/features/users/lib/quota-safety.ts diff --git a/common/verification.go b/common/verification.go index 1a390fd..c20238a 100644 --- a/common/verification.go +++ b/common/verification.go @@ -9,8 +9,9 @@ import ( ) type verificationValue struct { - code string - time time.Time + code string + time time.Time + claimID string } const ( @@ -49,19 +50,56 @@ func VerifyCodeWithKey(key string, code string, purpose string) bool { defer verificationMutex.Unlock() value, okay := verificationMap[purpose+key] now := time.Now() - if !okay || int(now.Sub(value.time).Seconds()) >= VerificationValidMinutes*60 { + if !okay || value.claimID != "" || int(now.Sub(value.time).Seconds()) >= VerificationValidMinutes*60 { return false } return code == value.code } +// VerifyCodeWithKeyAndRun reserves a code while fn runs. A successful fn +// consumes the code; an error or panic releases the reservation for retry. +func VerifyCodeWithKeyAndRun(key string, code string, purpose string, fn func() error) (verified bool, err error) { + mapKey := purpose + key + claimID := uuid.NewString() + + verificationMutex.Lock() + value, okay := verificationMap[mapKey] + now := time.Now() + if !okay || value.claimID != "" || int(now.Sub(value.time).Seconds()) >= VerificationValidMinutes*60 || code != value.code { + verificationMutex.Unlock() + return false, nil + } + value.claimID = claimID + verificationMap[mapKey] = value + verificationMutex.Unlock() + + committed := false + defer func() { + verificationMutex.Lock() + current, exists := verificationMap[mapKey] + if exists && current.claimID == claimID { + if committed { + delete(verificationMap, mapKey) + } else { + current.claimID = "" + verificationMap[mapKey] = current + } + } + verificationMutex.Unlock() + }() + + err = fn() + committed = err == nil + return true, err +} + func VerifyAndDeleteCodeWithKey(key string, code string, purpose string) bool { verificationMutex.Lock() defer verificationMutex.Unlock() mapKey := purpose + key value, okay := verificationMap[mapKey] now := time.Now() - if !okay || int(now.Sub(value.time).Seconds()) >= VerificationValidMinutes*60 || code != value.code { + if !okay || value.claimID != "" || int(now.Sub(value.time).Seconds()) >= VerificationValidMinutes*60 || code != value.code { return false } delete(verificationMap, mapKey) diff --git a/common/verification_test.go b/common/verification_test.go index 2e03519..31e72f3 100644 --- a/common/verification_test.go +++ b/common/verification_test.go @@ -1,6 +1,10 @@ package common -import "testing" +import ( + "errors" + "sync/atomic" + "testing" +) func TestVerifyAndDeleteCodeWithKeyConsumesCodeOnce(t *testing.T) { key := "consume-once@example.com" @@ -17,3 +21,91 @@ func TestVerifyAndDeleteCodeWithKeyConsumesCodeOnce(t *testing.T) { t.Fatal("expected second verification to fail") } } + +func TestVerifyCodeWithKeyAndRunRestoresCodeAfterFailure(t *testing.T) { + key := "rollback@example.com" + code := "rollback-code" + wantErr := errors.New("password update failed") + RegisterVerificationCodeWithKey(key, code, PasswordResetPurpose) + t.Cleanup(func() { DeleteKey(key, PasswordResetPurpose) }) + + verified, err := VerifyCodeWithKeyAndRun(key, code, PasswordResetPurpose, func() error { + return wantErr + }) + + if !verified { + t.Fatal("expected code verification to succeed") + } + if !errors.Is(err, wantErr) { + t.Fatalf("expected callback error %v, got %v", wantErr, err) + } + if !VerifyCodeWithKey(key, code, PasswordResetPurpose) { + t.Fatal("expected code to remain valid after callback failure") + } +} + +func TestVerifyCodeWithKeyAndRunRestoresCodeAfterPanic(t *testing.T) { + key := "panic@example.com" + code := "panic-code" + RegisterVerificationCodeWithKey(key, code, PasswordResetPurpose) + t.Cleanup(func() { DeleteKey(key, PasswordResetPurpose) }) + + func() { + defer func() { + if recovered := recover(); recovered == nil { + t.Fatal("expected callback panic") + } + }() + _, _ = VerifyCodeWithKeyAndRun(key, code, PasswordResetPurpose, func() error { + panic("password update panic") + }) + }() + + if !VerifyCodeWithKey(key, code, PasswordResetPurpose) { + t.Fatal("expected code to remain valid after callback panic") + } +} + +func TestVerifyCodeWithKeyAndRunAllowsOnlyOneConcurrentClaim(t *testing.T) { + key := "concurrent@example.com" + code := "concurrent-code" + RegisterVerificationCodeWithKey(key, code, PasswordResetPurpose) + t.Cleanup(func() { DeleteKey(key, PasswordResetPurpose) }) + + started := make(chan struct{}) + release := make(chan struct{}) + firstDone := make(chan struct{}) + var callbackCount atomic.Int32 + var firstVerified bool + var firstErr error + go func() { + defer close(firstDone) + firstVerified, firstErr = VerifyCodeWithKeyAndRun(key, code, PasswordResetPurpose, func() error { + callbackCount.Add(1) + close(started) + <-release + return nil + }) + }() + + <-started + secondVerified, secondErr := VerifyCodeWithKeyAndRun(key, code, PasswordResetPurpose, func() error { + callbackCount.Add(1) + return nil + }) + close(release) + <-firstDone + + if !firstVerified || firstErr != nil { + t.Fatalf("expected first claim to succeed, verified=%t err=%v", firstVerified, firstErr) + } + if secondVerified || secondErr != nil { + t.Fatalf("expected concurrent claim to be rejected, verified=%t err=%v", secondVerified, secondErr) + } + if got := callbackCount.Load(); got != 1 { + t.Fatalf("expected one callback execution, got %d", got) + } + if VerifyCodeWithKey(key, code, PasswordResetPurpose) { + t.Fatal("expected successful claim to consume the code") + } +} diff --git a/controller/misc.go b/controller/misc.go index b2b5d4b..95248c9 100644 --- a/controller/misc.go +++ b/controller/misc.go @@ -348,12 +348,14 @@ func ResetPassword(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgInvalidParams) return } - if !common.VerifyAndDeleteCodeWithKey(req.Email, req.Token, common.PasswordResetPurpose) { + password := common.GenerateVerificationCode(12) + verified, err := common.VerifyCodeWithKeyAndRun(req.Email, req.Token, common.PasswordResetPurpose, func() error { + return model.ResetUserPasswordByEmail(req.Email, password) + }) + if !verified { common.ApiErrorI18n(c, i18n.MsgUserPasswordResetLinkInvalid) return } - password := common.GenerateVerificationCode(12) - err = model.ResetUserPasswordByEmail(req.Email, password) if err != nil { if errors.Is(err, model.ErrEmailNotFound) || errors.Is(err, model.ErrEmailAmbiguous) { common.ApiErrorI18n(c, i18n.MsgUserPasswordResetLinkInvalid) diff --git a/controller/user.go b/controller/user.go index b1cdde9..95014cd 100644 --- a/controller/user.go +++ b/controller/user.go @@ -35,7 +35,9 @@ var ( errOriginalPasswordFail = errors.New("original password is incorrect") ) -const maxUserQuotaValue = int64(^uint64(0) >> 1) +// User quota crosses the JSON boundary as a JavaScript number in the default +// frontend, so management operations must stay within its exact integer range. +const maxUserQuotaValue = int64(1<<53 - 1) func Login(c *gin.Context) { if !common.PasswordLoginEnabled { @@ -1000,12 +1002,19 @@ func isValidQuotaOverride(value int64) bool { } func isValidQuotaAddition(current int64, delta int64) bool { - if delta <= 0 || delta > maxUserQuotaValue { + if current < -maxUserQuotaValue || current > maxUserQuotaValue || delta <= 0 || delta > maxUserQuotaValue { return false } return current <= maxUserQuotaValue-delta } +func isValidQuotaSubtraction(current int64, delta int64) bool { + if current < -maxUserQuotaValue || current > maxUserQuotaValue || delta <= 0 || delta > maxUserQuotaValue { + return false + } + return current >= -maxUserQuotaValue+delta +} + // ManageUser Only admin user can do this func ManageUser(c *gin.Context) { var req ManageRequest @@ -1098,6 +1107,10 @@ func ManageUser(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgUserQuotaChangeZero) return } + if !isValidQuotaSubtraction(user.Quota, req.Value) { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) + return + } if err := model.DecreaseUserQuota(user.Id, req.Value, true); err != nil { common.ApiError(c, err) return diff --git a/controller/user_setting_test.go b/controller/user_setting_test.go index 5f2531f..52d20bf 100644 --- a/controller/user_setting_test.go +++ b/controller/user_setting_test.go @@ -2,6 +2,7 @@ package controller import ( "bytes" + "errors" "fmt" "net/http" "net/http/httptest" @@ -219,12 +220,15 @@ func TestRegisterConsumesEmailVerificationCode(t *testing.T) { func TestQuotaBoundsValidation(t *testing.T) { legacyInt32Max := int64(1<<31 - 1) aboveLegacyInt32Max := legacyInt32Max + 1 + javascriptSafeIntegerMax := int64(1<<53 - 1) maxQuota := maxUserQuotaValue overflowingHalfMax := maxQuota/2 + 1 + require.Equal(t, javascriptSafeIntegerMax, maxQuota) require.True(t, isValidQuotaOverride(0)) require.True(t, isValidQuotaOverride(aboveLegacyInt32Max)) require.True(t, isValidQuotaOverride(maxQuota)) + require.False(t, isValidQuotaOverride(maxQuota+1)) require.False(t, isValidQuotaOverride(-1)) require.True(t, isValidQuotaAddition(aboveLegacyInt32Max, 1)) @@ -236,6 +240,50 @@ func TestQuotaBoundsValidation(t *testing.T) { require.False(t, isValidQuotaAddition(overflowingHalfMax, overflowingHalfMax)) require.False(t, isValidQuotaAddition(maxQuota, maxQuota)) require.False(t, isValidQuotaAddition(0, 0)) + + require.True(t, isValidQuotaSubtraction(-maxQuota+1, 1)) + require.True(t, isValidQuotaSubtraction(0, maxQuota)) + require.False(t, isValidQuotaSubtraction(-maxQuota, 1)) + require.False(t, isValidQuotaSubtraction(0, maxQuota+1)) + require.False(t, isValidQuotaSubtraction(0, 0)) +} + +func TestResetPasswordKeepsVerificationCodeWhenPasswordUpdateFails(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + user := model.User{ + Username: "reset-failure-user", + Password: "ExistingPassword123", + Email: "reset-failure@example.com", + Status: common.UserStatusEnabled, + } + require.NoError(t, db.Create(&user).Error) + + wantErr := errors.New("forced password update failure") + const callbackName = "test:force_password_reset_failure" + require.NoError(t, db.Callback().Update().Before("gorm:update").Register(callbackName, func(tx *gorm.DB) { + tx.AddError(wantErr) + })) + t.Cleanup(func() { + require.NoError(t, db.Callback().Update().Remove(callbackName)) + }) + + code := "reset-code" + common.RegisterVerificationCodeWithKey(user.Email, code, common.PasswordResetPurpose) + t.Cleanup(func() { common.DeleteKey(user.Email, common.PasswordResetPurpose) }) + payload, err := common.Marshal(PasswordResetRequest{Email: user.Email, Token: code}) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/user/reset", bytes.NewReader(payload)) + ctx.Request.Header.Set("Content-Type", "application/json") + + ResetPassword(ctx) + + require.Contains(t, recorder.Body.String(), `"success":false`) + require.True(t, common.VerifyCodeWithKey(user.Email, code, common.PasswordResetPurpose)) + var stored model.User + require.NoError(t, db.First(&stored, user.Id).Error) + require.Equal(t, user.Password, stored.Password) } func TestManageUserRejectsQuotaValueAboveInt64Max(t *testing.T) { @@ -251,3 +299,34 @@ func TestManageUserRejectsQuotaValueAboveInt64Max(t *testing.T) { require.Equal(t, http.StatusOK, recorder.Code) require.Contains(t, recorder.Body.String(), `"success":false`) } + +func TestManageUserRejectsQuotaValueAboveJavaScriptSafeInteger(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + user := model.User{ + Username: "unsafe-quota-user", + Password: "password123", + Quota: 100, + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + } + require.NoError(t, db.Create(&user).Error) + payload, err := common.Marshal(ManageRequest{ + Id: user.Id, + Action: "add_quota", + Mode: "override", + Value: maxUserQuotaValue + 1, + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/user/manage", bytes.NewReader(payload)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("role", common.RoleRootUser) + + ManageUser(ctx) + + require.Contains(t, recorder.Body.String(), `"success":false`) + var stored model.User + require.NoError(t, db.First(&stored, user.Id).Error) + require.EqualValues(t, 100, stored.Quota) +} diff --git a/dto/billing_usage_test.go b/dto/billing_usage_test.go index b9bc8b3..26b1b49 100644 --- a/dto/billing_usage_test.go +++ b/dto/billing_usage_test.go @@ -63,6 +63,46 @@ func TestNewEstimatedGeminiChatBillingUsage(t *testing.T) { assert.Equal(t, 18, billingUsage.GeminiUsageMetadata.TotalTokenCount) } +func TestCopyInputTokenDetails(t *testing.T) { + src := &InputTokenDetails{ + CachedTokens: 1, + CachedCreationTokens: 2, + CacheWriteTokens: 3, + TextTokens: 4, + ImageTokens: 5, + AudioTokens: 6, + } + + overwritten := InputTokenDetails{ + CachedTokens: 10, + CachedCreationTokens: 20, + CacheWriteTokens: 30, + TextTokens: 40, + ImageTokens: 50, + AudioTokens: 60, + } + CopyInputTokenDetails(&overwritten, src, true) + assert.Equal(t, *src, overwritten) + + filled := InputTokenDetails{ + CachedTokens: 10, + CachedCreationTokens: 0, + CacheWriteTokens: 30, + TextTokens: 0, + ImageTokens: 50, + AudioTokens: 0, + } + CopyInputTokenDetails(&filled, src, false) + assert.Equal(t, InputTokenDetails{ + CachedTokens: 10, + CachedCreationTokens: 2, + CacheWriteTokens: 30, + TextTokens: 4, + ImageTokens: 50, + AudioTokens: 6, + }, filled) +} + func TestBillingUsageJSONUsesProtocolNamedFields(t *testing.T) { billingUsage := &BillingUsage{ OpenAIUsage: &Usage{PromptTokens: 1, BillingUsage: NewClaudeMessagesBillingUsage(&ClaudeUsage{InputTokens: 9})}, diff --git a/dto/openai_response.go b/dto/openai_response.go index 8242d14..6e536aa 100644 --- a/dto/openai_response.go +++ b/dto/openai_response.go @@ -262,6 +262,30 @@ type InputTokenDetails struct { ImageTokens int `json:"image_tokens"` } +func CopyInputTokenDetails(dst *InputTokenDetails, src *InputTokenDetails, overwriteExisting bool) { + if dst == nil || src == nil { + return + } + if overwriteExisting || dst.CachedTokens == 0 { + dst.CachedTokens = src.CachedTokens + } + if overwriteExisting || dst.CachedCreationTokens == 0 { + dst.CachedCreationTokens = src.CachedCreationTokens + } + if overwriteExisting || dst.CacheWriteTokens == 0 { + dst.CacheWriteTokens = src.CacheWriteTokens + } + if overwriteExisting || dst.TextTokens == 0 { + dst.TextTokens = src.TextTokens + } + if overwriteExisting || dst.ImageTokens == 0 { + dst.ImageTokens = src.ImageTokens + } + if overwriteExisting || dst.AudioTokens == 0 { + dst.AudioTokens = src.AudioTokens + } +} + func (d InputTokenDetails) CacheCreationTokensTotal() int { total := d.CachedCreationTokens if d.CacheWriteTokens > total { diff --git a/service/billing_usage.go b/service/billing_usage.go index cfb5ce3..f236120 100644 --- a/service/billing_usage.go +++ b/service/billing_usage.go @@ -125,27 +125,7 @@ func usageFromOpenAIBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage { } func fillPromptTokenDetailsFromInputDetails(promptDetails *dto.InputTokenDetails, inputDetails *dto.InputTokenDetails) { - if inputDetails == nil { - return - } - if promptDetails.CachedTokens == 0 { - promptDetails.CachedTokens = inputDetails.CachedTokens - } - if promptDetails.CachedCreationTokens == 0 { - promptDetails.CachedCreationTokens = inputDetails.CachedCreationTokens - } - if promptDetails.CacheWriteTokens == 0 { - promptDetails.CacheWriteTokens = inputDetails.CacheWriteTokens - } - if promptDetails.TextTokens == 0 { - promptDetails.TextTokens = inputDetails.TextTokens - } - if promptDetails.ImageTokens == 0 { - promptDetails.ImageTokens = inputDetails.ImageTokens - } - if promptDetails.AudioTokens == 0 { - promptDetails.AudioTokens = inputDetails.AudioTokens - } + dto.CopyInputTokenDetails(promptDetails, inputDetails, false) } func usageFromClaudeBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage { diff --git a/service/openaicompat/responses_compat_test.go b/service/openaicompat/responses_compat_test.go index 92043a5..b14a44c 100644 --- a/service/openaicompat/responses_compat_test.go +++ b/service/openaicompat/responses_compat_test.go @@ -289,6 +289,7 @@ func TestResponsesResponseToChatCompletionsResponsePreservesUsageDetails(t *test InputTokensDetails: &dto.InputTokenDetails{ CachedTokens: 1, CachedCreationTokens: 2, + CacheWriteTokens: 10, TextTokens: 3, AudioTokens: 4, ImageTokens: 5, @@ -317,6 +318,7 @@ func TestResponsesResponseToChatCompletionsResponsePreservesUsageDetails(t *test assert.Equal(t, 18, usage.TotalTokens) assert.Equal(t, 1, usage.PromptTokensDetails.CachedTokens) assert.Equal(t, 2, usage.PromptTokensDetails.CachedCreationTokens) + assert.Equal(t, 10, usage.PromptTokensDetails.CacheWriteTokens) assert.Equal(t, 3, usage.PromptTokensDetails.TextTokens) assert.Equal(t, 4, usage.PromptTokensDetails.AudioTokens) assert.Equal(t, 5, usage.PromptTokensDetails.ImageTokens) diff --git a/service/openaicompat/responses_to_chat.go b/service/openaicompat/responses_to_chat.go index 7620fc0..9f2986c 100644 --- a/service/openaicompat/responses_to_chat.go +++ b/service/openaicompat/responses_to_chat.go @@ -135,14 +135,7 @@ func UsageFromResponsesUsage(src *dto.Usage) *dto.Usage { } else { usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } - if src.InputTokensDetails != nil { - usage.PromptTokensDetails.CachedTokens = src.InputTokensDetails.CachedTokens - usage.PromptTokensDetails.ImageTokens = src.InputTokensDetails.ImageTokens - usage.PromptTokensDetails.AudioTokens = src.InputTokensDetails.AudioTokens - usage.PromptTokensDetails.CachedCreationTokens = src.InputTokensDetails.CachedCreationTokens - usage.PromptTokensDetails.CacheWriteTokens = src.InputTokensDetails.CacheWriteTokens - usage.PromptTokensDetails.TextTokens = src.InputTokensDetails.TextTokens - } + dto.CopyInputTokenDetails(&usage.PromptTokensDetails, src.InputTokensDetails, true) usage.CompletionTokenDetails = src.CompletionTokenDetails usage.BillingUsage = dto.CloneBillingUsage(src.BillingUsage) if usage.BillingUsage == nil { diff --git a/web/default/src/features/users/components/user-quota-dialog.tsx b/web/default/src/features/users/components/user-quota-dialog.tsx index 915a128..9c4990f 100644 --- a/web/default/src/features/users/components/user-quota-dialog.tsx +++ b/web/default/src/features/users/components/user-quota-dialog.tsx @@ -34,6 +34,7 @@ import { import { Input } from '@/components/ui/input' import { Label } from '@/components/ui/label' import { adjustUserQuota } from '../api' +import { isSafeQuotaAdjustment } from '../lib/quota-safety' import type { QuotaAdjustMode } from '../types' interface UserQuotaDialogProps { @@ -56,6 +57,13 @@ export function UserQuotaDialog(props: UserQuotaDialogProps) { const amountValue = parseFloat(amount) || 0 const quotaValue = parseQuotaFromDollars(Math.abs(amountValue)) + const adjustmentValue = + mode === 'override' ? parseQuotaFromDollars(amountValue) : quotaValue + const quotaIsSafe = isSafeQuotaAdjustment( + props.currentQuota, + mode, + adjustmentValue + ) const getPreviewText = () => { const current = props.currentQuota @@ -77,16 +85,19 @@ export function UserQuotaDialog(props: UserQuotaDialogProps) { const handleConfirm = async () => { if (!amount && mode !== 'override') return if (quotaValue <= 0 && mode !== 'override') return + if (!quotaIsSafe) { + toast.error(t('Quota must be a safe integer within the supported range')) + return + } setLoading(true) try { - const value = - mode === 'override' ? parseQuotaFromDollars(amountValue) : quotaValue const result = await adjustUserQuota({ id: props.userId, action: 'add_quota', mode, - value: mode === 'override' ? value : Math.abs(value), + value: + mode === 'override' ? adjustmentValue : Math.abs(adjustmentValue), }) if (result.success) { toast.success(t('Quota adjusted successfully')) @@ -171,13 +182,18 @@ export function UserQuotaDialog(props: UserQuotaDialogProps) { if (e.key === 'Enter') handleConfirm() }} /> + {amount && !quotaIsSafe && ( +

+ {t('Quota must be a safe integer within the supported range')} +

+ )}
- diff --git a/web/default/src/features/users/lib/quota-safety.test.ts b/web/default/src/features/users/lib/quota-safety.test.ts new file mode 100644 index 0000000..c50475d --- /dev/null +++ b/web/default/src/features/users/lib/quota-safety.test.ts @@ -0,0 +1,54 @@ +/* +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 assert from 'node:assert/strict' +import { describe, test } from 'node:test' +import { isSafeQuotaAdjustment, isSafeQuotaValue } from './quota-safety' + +describe('quota safety', () => { + test('accepts only JavaScript-safe integer quota values', () => { + assert.equal(isSafeQuotaValue(Number.MAX_SAFE_INTEGER), true) + assert.equal(isSafeQuotaValue(Number.MAX_SAFE_INTEGER + 1), false) + assert.equal(isSafeQuotaValue(1.5), false) + assert.equal(isSafeQuotaValue(Number.POSITIVE_INFINITY), false) + }) + + test('rejects adjustments whose result leaves the safe integer range', () => { + assert.equal( + isSafeQuotaAdjustment(Number.MAX_SAFE_INTEGER - 1, 'add', 1), + true + ) + assert.equal( + isSafeQuotaAdjustment(Number.MAX_SAFE_INTEGER, 'add', 1), + false + ) + assert.equal( + isSafeQuotaAdjustment(-Number.MAX_SAFE_INTEGER + 1, 'subtract', 1), + true + ) + assert.equal( + isSafeQuotaAdjustment(-Number.MAX_SAFE_INTEGER, 'subtract', 1), + false + ) + assert.equal(isSafeQuotaAdjustment(0, 'override', 0), true) + assert.equal( + isSafeQuotaAdjustment(Number.MAX_SAFE_INTEGER + 1, 'override', 100), + true + ) + }) +}) diff --git a/web/default/src/features/users/lib/quota-safety.ts b/web/default/src/features/users/lib/quota-safety.ts new file mode 100644 index 0000000..4754585 --- /dev/null +++ b/web/default/src/features/users/lib/quota-safety.ts @@ -0,0 +1,39 @@ +/* +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 type { QuotaAdjustMode } from '../types' + +export function isSafeQuotaValue(value: number) { + return ( + Number.isSafeInteger(value) && Math.abs(value) <= Number.MAX_SAFE_INTEGER + ) +} + +export function isSafeQuotaAdjustment( + currentQuota: number, + mode: QuotaAdjustMode, + value: number +) { + if (!isSafeQuotaValue(value)) return false + if (mode === 'override') return value >= 0 + if (!isSafeQuotaValue(currentQuota)) return false + if (value <= 0) return false + + const result = mode === 'add' ? currentQuota + value : currentQuota - value + return isSafeQuotaValue(result) +} diff --git a/web/default/src/i18n/locales/en.json b/web/default/src/i18n/locales/en.json index 829150d..786fb69 100644 --- a/web/default/src/i18n/locales/en.json +++ b/web/default/src/i18n/locales/en.json @@ -3602,6 +3602,7 @@ "Quota ({{currency}})": "Quota ({{currency}})", "quota + refund": "quota + refund", "Quota adjusted successfully": "Quota adjusted successfully", + "Quota must be a safe integer within the supported range": "Quota must be a safe integer within the supported range", "Quota consumed before charging users": "Quota consumed before charging users", "Quota Distribution": "Quota Distribution", "Quota given to invited users": "Quota given to invited users", diff --git a/web/default/src/i18n/locales/fr.json b/web/default/src/i18n/locales/fr.json index f40c214..8de60bc 100644 --- a/web/default/src/i18n/locales/fr.json +++ b/web/default/src/i18n/locales/fr.json @@ -3602,6 +3602,7 @@ "Quota ({{currency}})": "Quota ({{currency}})", "quota + refund": "quota + refund", "Quota adjusted successfully": "Quota ajusté avec succès", + "Quota must be a safe integer within the supported range": "Le quota doit être un entier sûr dans la plage prise en charge", "Quota consumed before charging users": "Quota consommé avant de facturer les utilisateurs", "Quota Distribution": "Distribution des quotas", "Quota given to invited users": "Quota attribué aux utilisateurs invités", diff --git a/web/default/src/i18n/locales/ja.json b/web/default/src/i18n/locales/ja.json index 91e7087..33870bd 100644 --- a/web/default/src/i18n/locales/ja.json +++ b/web/default/src/i18n/locales/ja.json @@ -3602,6 +3602,7 @@ "Quota ({{currency}})": "クォータ ({{currency}})", "quota + refund": "quota + refund", "Quota adjusted successfully": "クォータの調整に成功しました", + "Quota must be a safe integer within the supported range": "クォータはサポート範囲内の安全な整数である必要があります", "Quota consumed before charging users": "ユーザーに請求する前に消費されるクォータ", "Quota Distribution": "クォータの分配", "Quota given to invited users": "招待されたユーザーに付与されるクォータ", diff --git a/web/default/src/i18n/locales/ru.json b/web/default/src/i18n/locales/ru.json index 53b9524..f14eda6 100644 --- a/web/default/src/i18n/locales/ru.json +++ b/web/default/src/i18n/locales/ru.json @@ -3602,6 +3602,7 @@ "Quota ({{currency}})": "Квота ({{currency}})", "quota + refund": "quota + refund", "Quota adjusted successfully": "Квота успешно изменена", + "Quota must be a safe integer within the supported range": "Квота должна быть безопасным целым числом в поддерживаемом диапазоне", "Quota consumed before charging users": "Квота, потребляемая до взимания платы с пользователей", "Quota Distribution": "Распределение квоты", "Quota given to invited users": "Квота, предоставляемая приглашенным пользователям", diff --git a/web/default/src/i18n/locales/vi.json b/web/default/src/i18n/locales/vi.json index 09fe50e..39af4f6 100644 --- a/web/default/src/i18n/locales/vi.json +++ b/web/default/src/i18n/locales/vi.json @@ -3602,6 +3602,7 @@ "Quota ({{currency}})": "Hạn mức ({{currency}})", "quota + refund": "quota + refund", "Quota adjusted successfully": "Điều chỉnh hạn mức thành công", + "Quota must be a safe integer within the supported range": "Hạn mức phải là số nguyên an toàn trong phạm vi được hỗ trợ", "Quota consumed before charging users": "Hạn mức tiêu thụ trước khi tính phí người dùng", "Quota Distribution": "Phân bổ hạn ngạch", "Quota given to invited users": "Hạn mức được cấp cho người dùng được mời", diff --git a/web/default/src/i18n/locales/zh.json b/web/default/src/i18n/locales/zh.json index d56c176..2e0e5ad 100644 --- a/web/default/src/i18n/locales/zh.json +++ b/web/default/src/i18n/locales/zh.json @@ -3602,6 +3602,7 @@ "Quota ({{currency}})": "额度 ({{currency}})", "quota + refund": "额度 + 退款", "Quota adjusted successfully": "调整额度成功", + "Quota must be a safe integer within the supported range": "额度必须是支持范围内的安全整数", "Quota consumed before charging users": "向用户收费前消耗的配额", "Quota Distribution": "消耗分布", "Quota given to invited users": "授予被邀请用户的配额", From 444a05ebd9018388c616848f0f3e3623d8e24203 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sun, 12 Jul 2026 21:56:01 +0800 Subject: [PATCH 4/9] v1.0.4 --- web/default/src/features/users/lib/quota-safety.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/web/default/src/features/users/lib/quota-safety.ts b/web/default/src/features/users/lib/quota-safety.ts index 4754585..4839a18 100644 --- a/web/default/src/features/users/lib/quota-safety.ts +++ b/web/default/src/features/users/lib/quota-safety.ts @@ -19,9 +19,7 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API import type { QuotaAdjustMode } from '../types' export function isSafeQuotaValue(value: number) { - return ( - Number.isSafeInteger(value) && Math.abs(value) <= Number.MAX_SAFE_INTEGER - ) + return Number.isSafeInteger(value) } export function isSafeQuotaAdjustment( From bab6a4cba1877784db25f0b63852dfd30c5a44fc Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sun, 12 Jul 2026 22:20:04 +0800 Subject: [PATCH 5/9] v1.0.4 --- web/default/src/features/users/lib/quota-safety.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/web/default/src/features/users/lib/quota-safety.ts b/web/default/src/features/users/lib/quota-safety.ts index 4839a18..885a39e 100644 --- a/web/default/src/features/users/lib/quota-safety.ts +++ b/web/default/src/features/users/lib/quota-safety.ts @@ -18,7 +18,7 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API */ import type { QuotaAdjustMode } from '../types' -export function isSafeQuotaValue(value: number) { +export function isSafeQuotaValue(value: number): boolean { return Number.isSafeInteger(value) } @@ -26,7 +26,7 @@ export function isSafeQuotaAdjustment( currentQuota: number, mode: QuotaAdjustMode, value: number -) { +): boolean { if (!isSafeQuotaValue(value)) return false if (mode === 'override') return value >= 0 if (!isSafeQuotaValue(currentQuota)) return false From 9f4a745997842bd8a6e7dfcfb401fa0026a37491 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 13 Jul 2026 00:25:14 +0800 Subject: [PATCH 6/9] v1.0.4 --- controller/user.go | 17 +- controller/user_setting_test.go | 27 +++ router/api-router.go | 2 +- router/api_router_test.go | 35 ++++ web/default/src/features/profile/api.ts | 13 +- .../dialogs/access-token-dialog.tsx | 168 +++++++++++------- .../profile/hooks/use-access-token.ts | 57 +++--- 7 files changed, 229 insertions(+), 90 deletions(-) create mode 100644 router/api_router_test.go diff --git a/controller/user.go b/controller/user.go index 95014cd..29e2c93 100644 --- a/controller/user.go +++ b/controller/user.go @@ -23,6 +23,7 @@ import ( "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" + "gorm.io/gorm" ) type LoginRequest struct { @@ -1024,13 +1025,19 @@ func ManageUser(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgInvalidParams) return } - user := model.User{ - Id: req.Id, + if req.Id <= 0 { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) + return } + + var user model.User // Fill attributes - model.DB.Unscoped().Where(&user).First(&user) - if user.Id == 0 { - common.ApiErrorI18n(c, i18n.MsgUserNotExists) + if err := model.DB.Unscoped().Where("id = ?", req.Id).First(&user).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + common.ApiErrorI18n(c, i18n.MsgUserNotExists) + return + } + common.ApiError(c, err) return } myRole := c.GetInt("role") diff --git a/controller/user_setting_test.go b/controller/user_setting_test.go index 52d20bf..55a8259 100644 --- a/controller/user_setting_test.go +++ b/controller/user_setting_test.go @@ -330,3 +330,30 @@ func TestManageUserRejectsQuotaValueAboveJavaScriptSafeInteger(t *testing.T) { require.NoError(t, db.First(&stored, user.Id).Error) require.EqualValues(t, 100, stored.Quota) } + +func TestManageUserRejectsMissingIdWithoutTouchingFirstUser(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + user := model.User{ + Username: "manage-missing-id-first-user", + Password: "password123", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + } + require.NoError(t, db.Create(&user).Error) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/user/manage", strings.NewReader( + `{"action":"disable"}`, + )) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("role", common.RoleRootUser) + + ManageUser(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":false`) + var stored model.User + require.NoError(t, db.First(&stored, user.Id).Error) + require.Equal(t, common.UserStatusEnabled, stored.Status) +} diff --git a/router/api-router.go b/router/api-router.go index d8d58f0..1a16aab 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -85,7 +85,7 @@ func SetApiRouter(router *gin.Engine) { selfRoute.GET("/models", controller.GetUserModels) selfRoute.PUT("/self", middleware.CriticalRateLimit(), controller.UpdateSelf) selfRoute.DELETE("/self", controller.DeleteSelf) - selfRoute.GET("/token", controller.GenerateAccessToken) + selfRoute.POST("/token", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(), controller.GenerateAccessToken) selfRoute.GET("/passkey", controller.PasskeyStatus) selfRoute.POST("/passkey/register/begin", controller.PasskeyRegisterBegin) selfRoute.POST("/passkey/register/finish", controller.PasskeyRegisterFinish) diff --git a/router/api_router_test.go b/router/api_router_test.go new file mode 100644 index 0000000..332b25d --- /dev/null +++ b/router/api_router_test.go @@ -0,0 +1,35 @@ +package router + +import ( + "testing" + + "github.com/gin-gonic/gin" +) + +func TestUserTokenRouteUsesPostOnly(t *testing.T) { + gin.SetMode(gin.TestMode) + + engine := gin.New() + SetApiRouter(engine) + + var hasPost bool + var hasGet bool + for _, route := range engine.Routes() { + if route.Path != "/api/user/token" { + continue + } + switch route.Method { + case "POST": + hasPost = true + case "GET": + hasGet = true + } + } + + if !hasPost { + t.Fatal("expected POST /api/user/token route to be registered") + } + if hasGet { + t.Fatal("did not expect GET /api/user/token route to be registered") + } +} diff --git a/web/default/src/features/profile/api.ts b/web/default/src/features/profile/api.ts index 20e0f2b..e108849 100644 --- a/web/default/src/features/profile/api.ts +++ b/web/default/src/features/profile/api.ts @@ -16,7 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { api } from '@/lib/api' +import { api, type ApiRequestConfig } from '@/lib/api' import type { ApiResponse, UserProfile, @@ -27,6 +27,11 @@ import type { CheckinResponse, } from './types' +const sensitiveActionConfig: ApiRequestConfig = { + skipBusinessError: true, + skipErrorHandler: true, +} + // ============================================================================ // User Profile APIs // ============================================================================ @@ -83,7 +88,11 @@ export async function deleteUserAccount( * Generate/regenerate system access token */ export async function generateAccessToken(): Promise> { - const res = await api.get('/api/user/token') + const res = await api.post>( + '/api/user/token', + undefined, + sensitiveActionConfig + ) return res.data } diff --git a/web/default/src/features/profile/components/dialogs/access-token-dialog.tsx b/web/default/src/features/profile/components/dialogs/access-token-dialog.tsx index 665f394..c1c2f81 100644 --- a/web/default/src/features/profile/components/dialogs/access-token-dialog.tsx +++ b/web/default/src/features/profile/components/dialogs/access-token-dialog.tsx @@ -16,10 +16,14 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useEffect } from 'react' +import { useCallback, useEffect } from 'react' import { RefreshCw, Loader2 } from 'lucide-react' import { useTranslation } from 'react-i18next' import { Button } from '@/components/ui/button' +import { + SecureVerificationDialog, + useSecureVerification, +} from '@/features/auth/secure-verification' import { Dialog, DialogContent, @@ -48,76 +52,118 @@ export function AccessTokenDialog({ }: AccessTokenDialogProps) { const { t } = useTranslation() const { token, generating, generate } = useAccessToken() + const { + open: verificationOpen, + methods: verificationMethods, + state: verificationState, + executeVerification, + withVerification, + cancel: cancelVerification, + setCode: setVerificationCode, + switchMethod: switchVerificationMethod, + } = useSecureVerification() + + const handleGenerate = useCallback(async () => { + await withVerification( + () => generate({ rethrowVerificationRequired: true }), + { + preferredMethod: 'passkey', + title: t('Security verification'), + description: t( + "Your system access token for API authentication. Keep it secure and don't share it with others." + ), + } + ) + }, [generate, t, withVerification]) // Auto-generate token when dialog opens if no token exists useEffect(() => { if (open && !token) { - generate() + handleGenerate() } - }, [open, token, generate]) + }, [open, token, handleGenerate]) return ( - - - - {t('Access Token')} - - {t( - "Your system access token for API authentication. Keep it secure and don't share it with others." - )} - - + <> + + + + {t('Access Token')} + + {t( + "Your system access token for API authentication. Keep it secure and don't share it with others." + )} + + -
-
- -
- - +
+
+ +
+ + +
+

+ {t('Use this token for API authentication')} +

-

- {t('Use this token for API authentication')} -

-
- - - - - -
+ + + + +
+
+ + { + if (!nextOpen) { + cancelVerification() + } + }} + methods={verificationMethods} + state={verificationState} + onVerify={async (method, code) => { + await executeVerification(method, code) + }} + onCancel={cancelVerification} + onCodeChange={setVerificationCode} + onMethodChange={switchVerificationMethod} + /> + ) } diff --git a/web/default/src/features/profile/hooks/use-access-token.ts b/web/default/src/features/profile/hooks/use-access-token.ts index 04300da..675a27f 100644 --- a/web/default/src/features/profile/hooks/use-access-token.ts +++ b/web/default/src/features/profile/hooks/use-access-token.ts @@ -20,41 +20,56 @@ import { useState, useCallback } from 'react' import i18next from 'i18next' import { toast } from 'sonner' import { useCopyToClipboard } from '@/hooks/use-copy-to-clipboard' +import { isVerificationRequiredError } from '@/lib/secure-verification' import { generateAccessToken } from '../api' // ============================================================================ // Access Token Hook // ============================================================================ +interface GenerateAccessTokenOptions { + rethrowVerificationRequired?: boolean +} + export function useAccessToken() { const [token, setToken] = useState('') const [generating, setGenerating] = useState(false) const { copyToClipboard } = useCopyToClipboard({ notify: false }) // Generate new access token - const generate = useCallback(async (): Promise => { - try { - setGenerating(true) - const response = await generateAccessToken() + const generate = useCallback( + async (options: GenerateAccessTokenOptions = {}): Promise => { + try { + setGenerating(true) + const response = await generateAccessToken() - if (response.success && response.data) { - setToken(response.data) - copyToClipboard(response.data) - toast.success(i18next.t('Token regenerated and copied to clipboard')) - return true - } + if (response.success && response.data) { + setToken(response.data) + copyToClipboard(response.data) + toast.success(i18next.t('Token regenerated and copied to clipboard')) + return true + } - toast.error(response.message || i18next.t('Failed to generate token')) - return false - } catch (error) { - // eslint-disable-next-line no-console - console.error('Failed to generate token:', error) - toast.error(i18next.t('Failed to generate token')) - return false - } finally { - setGenerating(false) - } - }, [copyToClipboard]) + toast.error(response.message || i18next.t('Failed to generate token')) + return false + } catch (error) { + if ( + options.rethrowVerificationRequired && + isVerificationRequiredError(error) + ) { + throw error + } + + // eslint-disable-next-line no-console + console.error('Failed to generate token:', error) + toast.error(i18next.t('Failed to generate token')) + return false + } finally { + setGenerating(false) + } + }, + [copyToClipboard] + ) return { token, From f058a64ee1b5e9ed34da8c0cc69f56ae9a520aac Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 13 Jul 2026 02:31:42 +0800 Subject: [PATCH 7/9] v1.0.4 --- controller/passkey.go | 8 +- controller/secure_verification.go | 134 +++++++++-- controller/secure_verification_test.go | 215 ++++++++++++++++++ controller/user.go | 5 + docs/openapi/api.json | 31 ++- middleware/secure_verification.go | 64 +++++- middleware/secure_verification_test.go | 93 ++++++++ router/api-router.go | 3 +- router/api_router_test.go | 33 +++ .../features/auth/secure-verification/api.ts | 95 +++++--- .../components/secure-verification-dialog.tsx | 37 ++- .../hooks/use-secure-verification.ts | 54 +++-- .../secure-verification/method-selection.ts | 20 ++ .../auth/secure-verification/types.ts | 6 +- .../dialogs/access-token-dialog.tsx | 9 +- web/default/src/i18n/locales/en.json | 2 + web/default/src/i18n/locales/fr.json | 2 + web/default/src/i18n/locales/ja.json | 2 + web/default/src/i18n/locales/ru.json | 2 + web/default/src/i18n/locales/vi.json | 2 + web/default/src/i18n/locales/zh.json | 2 + ...cure-verification-method-selection.test.ts | 32 +++ 22 files changed, 760 insertions(+), 91 deletions(-) create mode 100644 controller/secure_verification_test.go create mode 100644 middleware/secure_verification_test.go create mode 100644 web/default/src/features/auth/secure-verification/method-selection.ts create mode 100644 web/default/tests/secure-verification-method-selection.test.ts diff --git a/controller/passkey.go b/controller/passkey.go index 9ad2faf..d142ab8 100644 --- a/controller/passkey.go +++ b/controller/passkey.go @@ -498,6 +498,8 @@ func PasskeyVerifyFinish(c *gin.Context) { session.Set(PasskeyReadySessionKey, time.Now().Unix()) session.Delete(SecureVerificationSessionKey) session.Delete(secureVerificationMethodSessionKey) + session.Delete(secureVerificationUserSessionKey) + session.Delete(secureVerificationScopeSessionKey) if err := session.Save(); err != nil { common.ApiError(c, fmt.Errorf("保存验证状态失败: %v", err)) return @@ -570,9 +572,13 @@ func requirePasskeyDeleteVerification(c *gin.Context, userID int) bool { func requireSecureVerificationMethod(c *gin.Context, method string) bool { session := sessions.Default(c) verifiedAt, ok := session.Get(SecureVerificationSessionKey).(int64) - if !ok || time.Now().Unix()-verifiedAt >= SecureVerificationTimeout { + verifiedUserID, userOK := session.Get(secureVerificationUserSessionKey).(int) + elapsed := time.Now().Unix() - verifiedAt + if !ok || !userOK || verifiedUserID != c.GetInt("id") || elapsed < 0 || elapsed >= SecureVerificationTimeout { session.Delete(SecureVerificationSessionKey) session.Delete(secureVerificationMethodSessionKey) + session.Delete(secureVerificationUserSessionKey) + session.Delete(secureVerificationScopeSessionKey) _ = session.Save() common.ApiErrorMsg(c, "请先完成安全验证") return false diff --git a/controller/secure_verification.go b/controller/secure_verification.go index 700b608..0d5727e 100644 --- a/controller/secure_verification.go +++ b/controller/secure_verification.go @@ -1,6 +1,7 @@ package controller import ( + "errors" "fmt" "net/http" "time" @@ -15,8 +16,12 @@ const ( // SecureVerificationSessionKey means the user has fully passed secure verification. SecureVerificationSessionKey = "secure_verified_at" secureVerificationMethodSessionKey = "secure_verified_method" + secureVerificationUserSessionKey = "secure_verified_user_id" + secureVerificationScopeSessionKey = "secure_verified_scope" secureVerificationMethod2FA = "2fa" secureVerificationMethodPasskey = "passkey" + secureVerificationMethodPassword = "password" + secureVerificationScopeAccessToken = "access_token" // PasskeyReadySessionKey means WebAuthn finished and /api/verify can finalize step-up verification. PasskeyReadySessionKey = "secure_passkey_ready_at" // SecureVerificationTimeout 验证有效期(秒) @@ -26,13 +31,81 @@ const ( ) type UniversalVerifyRequest struct { - Method string `json:"method"` // "2fa" 或 "passkey" - Code string `json:"code,omitempty"` + Method string `json:"method"` // "2fa"、"passkey" 或 "password" + Code string `json:"code,omitempty"` + Password string `json:"password,omitempty"` + Scope string `json:"scope,omitempty"` } -type VerificationStatusResponse struct { - Verified bool `json:"verified"` - ExpiresAt int64 `json:"expires_at,omitempty"` +type VerificationMethodsResponse struct { + Has2FA bool `json:"has_2fa"` + HasPasskey bool `json:"has_passkey"` + HasPassword bool `json:"has_password"` +} + +type secureVerificationMethods struct { + twoFA *model.TwoFA + has2FA bool + hasPasskey bool + hasPassword bool +} + +func loadSecureVerificationMethods(user *model.User, allowPassword bool) (secureVerificationMethods, error) { + methods := secureVerificationMethods{ + hasPassword: user != nil && user.Password != "", + } + if user == nil || user.Id == 0 { + return methods, errors.New("用户不存在") + } + + twoFA, err := model.GetTwoFAByUserId(user.Id) + if err != nil { + return methods, err + } + methods.twoFA = twoFA + methods.has2FA = twoFA != nil && twoFA.IsEnabled + + _, err = model.GetPasskeyByUserID(user.Id) + switch { + case err == nil: + methods.hasPasskey = true + case errors.Is(err, model.ErrPasskeyNotFound): + methods.hasPasskey = false + default: + return methods, err + } + + // Password is a fallback only when no stronger verification method is enrolled. + methods.hasPassword = allowPassword && methods.hasPassword && !methods.has2FA && !methods.hasPasskey + return methods, nil +} + +func GetVerificationMethods(c *gin.Context) { + userId := c.GetInt("id") + user, err := model.GetUserById(userId, true) + if err != nil { + common.ApiError(c, err) + return + } + if user.Status != common.UserStatusEnabled { + common.ApiError(c, fmt.Errorf("该用户已被禁用")) + return + } + + methods, err := loadSecureVerificationMethods(user, c.Query("scope") == secureVerificationScopeAccessToken) + if err != nil { + common.ApiError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": VerificationMethodsResponse{ + Has2FA: methods.has2FA, + HasPasskey: methods.hasPasskey, + HasPassword: methods.hasPassword, + }, + }) } // UniversalVerify 通用验证接口 @@ -54,8 +127,8 @@ func UniversalVerify(c *gin.Context) { } // 获取用户信息 - user := &model.User{Id: userId} - if err := user.FillUserById(); err != nil { + user, err := model.GetUserById(userId, true) + if err != nil { common.ApiError(c, fmt.Errorf("获取用户信息失败: %v", err)) return } @@ -65,26 +138,19 @@ func UniversalVerify(c *gin.Context) { return } - // 检查用户的验证方式 - twoFA, _ := model.GetTwoFAByUserId(userId) - has2FA := twoFA != nil && twoFA.IsEnabled - - passkey, passkeyErr := model.GetPasskeyByUserID(userId) - hasPasskey := passkeyErr == nil && passkey != nil - - if !has2FA && !hasPasskey { - common.ApiError(c, fmt.Errorf("用户未启用2FA或Passkey")) + methods, err := loadSecureVerificationMethods(user, req.Scope == secureVerificationScopeAccessToken) + if err != nil { + common.ApiError(c, err) return } // 根据验证方式进行验证 var verified bool var verifyMethod string - var err error switch req.Method { case "2fa": - if !has2FA { + if !methods.has2FA { common.ApiError(c, fmt.Errorf("用户未启用2FA")) return } @@ -92,11 +158,11 @@ func UniversalVerify(c *gin.Context) { common.ApiError(c, fmt.Errorf("验证码不能为空")) return } - verified = validateTwoFactorAuth(twoFA, req.Code) + verified = validateTwoFactorAuth(methods.twoFA, req.Code) verifyMethod = "2FA" case "passkey": - if !hasPasskey { + if !methods.hasPasskey { common.ApiError(c, fmt.Errorf("用户未启用Passkey")) return } @@ -112,18 +178,34 @@ func UniversalVerify(c *gin.Context) { } verifyMethod = "Passkey" + case "password": + if !methods.hasPassword { + common.ApiError(c, fmt.Errorf("密码验证不可用,请使用2FA或Passkey")) + return + } + if req.Password == "" { + common.ApiError(c, fmt.Errorf("密码不能为空")) + return + } + verified = common.ValidatePasswordAndHash(req.Password, user.Password) + verifyMethod = "Password" + default: common.ApiError(c, fmt.Errorf("不支持的验证方式: %s", req.Method)) return } if !verified { - common.ApiError(c, fmt.Errorf("验证失败,请检查验证码")) + if req.Method == secureVerificationMethodPassword { + common.ApiError(c, fmt.Errorf("密码错误")) + } else { + common.ApiError(c, fmt.Errorf("验证失败,请检查验证码")) + } return } // 验证成功,在 session 中记录时间戳 - now, err := setSecureVerificationSession(c, req.Method) + now, err := setSecureVerificationSession(c, userId, req.Method, req.Scope) if err != nil { common.ApiError(c, fmt.Errorf("保存验证状态失败: %v", err)) return @@ -142,12 +224,18 @@ func UniversalVerify(c *gin.Context) { }) } -func setSecureVerificationSession(c *gin.Context, method string) (int64, error) { +func setSecureVerificationSession(c *gin.Context, userId int, method string, scope string) (int64, error) { session := sessions.Default(c) session.Delete(PasskeyReadySessionKey) now := time.Now().Unix() session.Set(SecureVerificationSessionKey, now) session.Set(secureVerificationMethodSessionKey, method) + session.Set(secureVerificationUserSessionKey, userId) + if scope == "" { + session.Delete(secureVerificationScopeSessionKey) + } else { + session.Set(secureVerificationScopeSessionKey, scope) + } if err := session.Save(); err != nil { return 0, err } diff --git a/controller/secure_verification_test.go b/controller/secure_verification_test.go new file mode 100644 index 0000000..1ad98d3 --- /dev/null +++ b/controller/secure_verification_test.go @@ -0,0 +1,215 @@ +package controller + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/gin-contrib/sessions" + "github.com/gin-contrib/sessions/cookie" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestUniversalVerifyPasswordCreatesUserBoundSession(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate( + &model.TwoFA{}, + &model.TwoFABackupCode{}, + &model.PasskeyCredential{}, + &model.Log{}, + )) + + const password = "Password123" + passwordHash, err := common.Password2Hash(password) + require.NoError(t, err) + user := model.User{ + Id: 1001, + Username: "password-verification-user", + Password: passwordHash, + DisplayName: "Password Verification User", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("password-verification-test")))) + router.POST("/verify", func(c *gin.Context) { + c.Set("id", user.Id) + UniversalVerify(c) + }) + router.GET("/verification-session", func(c *gin.Context) { + session := sessions.Default(c) + c.JSON(http.StatusOK, gin.H{ + "verified_at": session.Get(SecureVerificationSessionKey), + "verified_method": session.Get(secureVerificationMethodSessionKey), + "verified_user_id": session.Get(secureVerificationUserSessionKey), + "verified_scope": session.Get(secureVerificationScopeSessionKey), + }) + }) + + payload := []byte(`{"method":"password","password":"Password123","scope":"access_token"}`) + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/verify", bytes.NewReader(payload)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":true`) + + inspectRecorder := httptest.NewRecorder() + inspectRequest := httptest.NewRequest(http.MethodGet, "/verification-session", nil) + for _, sessionCookie := range recorder.Result().Cookies() { + inspectRequest.AddCookie(sessionCookie) + } + router.ServeHTTP(inspectRecorder, inspectRequest) + + require.Equal(t, http.StatusOK, inspectRecorder.Code) + require.Contains(t, inspectRecorder.Body.String(), `"verified_method":"password"`) + require.Contains(t, inspectRecorder.Body.String(), `"verified_user_id":1001`) + require.Contains(t, inspectRecorder.Body.String(), `"verified_scope":"access_token"`) +} + +func TestUniversalVerifyPasswordRequiresAccessTokenScope(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate( + &model.TwoFA{}, + &model.TwoFABackupCode{}, + &model.PasskeyCredential{}, + )) + + passwordHash, err := common.Password2Hash("Password123") + require.NoError(t, err) + user := model.User{ + Id: 1004, + Username: "unscoped-password-user", + Password: passwordHash, + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Set("id", user.Id) + ctx.Request = httptest.NewRequest( + http.MethodPost, + "/api/verify", + bytes.NewReader([]byte(`{"method":"password","password":"Password123"}`)), + ) + ctx.Request.Header.Set("Content-Type", "application/json") + + UniversalVerify(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":false`) + require.Contains(t, recorder.Body.String(), "密码验证不可用") +} + +func TestVerificationMethodsUsePasswordOnlyAsFallback(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate( + &model.TwoFA{}, + &model.TwoFABackupCode{}, + &model.PasskeyCredential{}, + )) + + passwordHash, err := common.Password2Hash("Password123") + require.NoError(t, err) + user := model.User{ + Id: 1002, + Username: "password-fallback-user", + Password: passwordHash, + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + methods, err := loadSecureVerificationMethods(&user, false) + require.NoError(t, err) + require.False(t, methods.hasPassword) + + methods, err = loadSecureVerificationMethods(&user, true) + require.NoError(t, err) + require.True(t, methods.hasPassword) + require.False(t, methods.has2FA) + require.False(t, methods.hasPasskey) + + require.NoError(t, db.Create(&model.TwoFA{ + UserId: user.Id, + Secret: "test-secret", + IsEnabled: true, + }).Error) + + methods, err = loadSecureVerificationMethods(&user, true) + require.NoError(t, err) + require.True(t, methods.has2FA) + require.False(t, methods.hasPassword) +} + +func TestSetupLoginClearsPreviousSecureVerification(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.Log{})) + + passwordHash, err := common.Password2Hash("Password123") + require.NoError(t, err) + user := model.User{ + Id: 1003, + Username: "new-session-user", + Password: passwordHash, + DisplayName: "New Session User", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("login-clears-verification-test")))) + router.GET("/login", func(c *gin.Context) { + session := sessions.Default(c) + session.Set(SecureVerificationSessionKey, int64(123)) + session.Set(secureVerificationMethodSessionKey, secureVerificationMethod2FA) + session.Set(secureVerificationUserSessionKey, 999) + session.Set(secureVerificationScopeSessionKey, secureVerificationScopeAccessToken) + session.Set(PasskeyReadySessionKey, int64(123)) + setupLogin(&user, c) + }) + router.GET("/verification-session", func(c *gin.Context) { + session := sessions.Default(c) + c.JSON(http.StatusOK, gin.H{ + "verified_at": session.Get(SecureVerificationSessionKey), + "verified_method": session.Get(secureVerificationMethodSessionKey), + "verified_user_id": session.Get(secureVerificationUserSessionKey), + "verified_scope": session.Get(secureVerificationScopeSessionKey), + "passkey_ready_at": session.Get(PasskeyReadySessionKey), + }) + }) + + loginRecorder := httptest.NewRecorder() + router.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/login", nil)) + require.Equal(t, http.StatusOK, loginRecorder.Code) + + inspectRecorder := httptest.NewRecorder() + inspectRequest := httptest.NewRequest(http.MethodGet, "/verification-session", nil) + for _, sessionCookie := range loginRecorder.Result().Cookies() { + inspectRequest.AddCookie(sessionCookie) + } + router.ServeHTTP(inspectRecorder, inspectRequest) + + require.Equal(t, http.StatusOK, inspectRecorder.Code) + require.JSONEq(t, `{ + "verified_at": null, + "verified_method": null, + "verified_user_id": null, + "verified_scope": null, + "passkey_ready_at": null + }`, inspectRecorder.Body.String()) +} diff --git a/controller/user.go b/controller/user.go index 29e2c93..131d84a 100644 --- a/controller/user.go +++ b/controller/user.go @@ -138,6 +138,11 @@ func recordLoginAudit(user *model.User, c *gin.Context) { func setupLogin(user *model.User, c *gin.Context) { model.UpdateUserLastLoginAt(user.Id) session := sessions.Default(c) + session.Delete(SecureVerificationSessionKey) + session.Delete(secureVerificationMethodSessionKey) + session.Delete(secureVerificationUserSessionKey) + session.Delete(secureVerificationScopeSessionKey) + session.Delete(PasskeyReadySessionKey) session.Set("id", user.Id) session.Set("username", user.Username) session.Set("role", user.Role) diff --git a/docs/openapi/api.json b/docs/openapi/api.json index 266bd04..e277abf 100644 --- a/docs/openapi/api.json +++ b/docs/openapi/api.json @@ -1180,10 +1180,10 @@ } }, "/api/user/token": { - "get": { + "post": { "summary": "生成访问令牌", "deprecated": false, - "description": "🔐 需要登录(User权限)", + "description": "🔐 需要登录并完成安全验证(User权限)", "tags": [ "用户管理" ], @@ -2310,6 +2310,31 @@ ] } }, + "/api/verify/methods": { + "get": { + "summary": "获取可用安全验证方式", + "deprecated": false, + "description": "🔐 需要登录(User权限)", + "tags": [ + "安全验证" + ], + "parameters": [], + "responses": { + "200": { + "description": "成功", + "headers": {} + } + }, + "security": [ + { + "Combination343": [] + }, + { + "Combination1243": [] + } + ] + } + }, "/api/verify/status": { "get": { "summary": "获取验证状态", @@ -7815,4 +7840,4 @@ "Combination1243": [] } ] -} \ No newline at end of file +} diff --git a/middleware/secure_verification.go b/middleware/secure_verification.go index b218f6b..6b9181e 100644 --- a/middleware/secure_verification.go +++ b/middleware/secure_verification.go @@ -12,6 +12,9 @@ const ( // SecureVerificationSessionKey 安全验证的 session key(与 controller 保持一致) SecureVerificationSessionKey = "secure_verified_at" secureVerificationMethodSessionKey = "secure_verified_method" + secureVerificationUserSessionKey = "secure_verified_user_id" + secureVerificationScopeSessionKey = "secure_verified_scope" + secureVerificationMethodPassword = "password" // SecureVerificationTimeout 验证有效期(秒) SecureVerificationTimeout = 300 // 5分钟 ) @@ -19,7 +22,7 @@ const ( // SecureVerificationRequired 安全验证中间件 // 检查用户是否在有效时间内通过了安全验证 // 如果未验证或验证已过期,返回 401 错误 -func SecureVerificationRequired() gin.HandlerFunc { +func SecureVerificationRequired(requiredScopes ...string) gin.HandlerFunc { return func(c *gin.Context) { // 检查用户是否已登录 userId := c.GetInt("id") @@ -58,10 +61,21 @@ func SecureVerificationRequired() gin.HandlerFunc { c.Abort() return } + verifiedUserID, ok := session.Get(secureVerificationUserSessionKey).(int) + if !ok || verifiedUserID != userId { + clearSecureVerificationSession(session) + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": "验证状态与当前用户不匹配,请重新验证", + "code": "VERIFICATION_INVALID", + }) + c.Abort() + return + } // 检查验证是否过期 elapsed := time.Now().Unix() - verifiedAt - if elapsed >= SecureVerificationTimeout { + if elapsed < 0 || elapsed >= SecureVerificationTimeout { // 验证已过期,清除 session clearSecureVerificationSession(session) c.JSON(http.StatusForbidden, gin.H{ @@ -73,6 +87,34 @@ func SecureVerificationRequired() gin.HandlerFunc { return } + verifiedMethod, ok := session.Get(secureVerificationMethodSessionKey).(string) + if !ok || verifiedMethod == "" { + clearSecureVerificationSession(session) + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": "验证状态异常,请重新验证", + "code": "VERIFICATION_INVALID", + }) + c.Abort() + return + } + if verifiedMethod == secureVerificationMethodPassword { + requiredScope := "" + if len(requiredScopes) > 0 { + requiredScope = requiredScopes[0] + } + verifiedScope, _ := session.Get(secureVerificationScopeSessionKey).(string) + if requiredScope == "" || verifiedScope != requiredScope { + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": "需要对应操作的安全验证", + "code": "VERIFICATION_REQUIRED", + }) + c.Abort() + return + } + } + c.Next() } } @@ -80,6 +122,8 @@ func SecureVerificationRequired() gin.HandlerFunc { func clearSecureVerificationSession(session sessions.Session) { session.Delete(SecureVerificationSessionKey) session.Delete(secureVerificationMethodSessionKey) + session.Delete(secureVerificationUserSessionKey) + session.Delete(secureVerificationScopeSessionKey) _ = session.Save() } @@ -106,18 +150,32 @@ func OptionalSecureVerification() gin.HandlerFunc { verifiedAt, ok := verifiedAtRaw.(int64) if !ok { + clearSecureVerificationSession(session) + c.Set("secure_verified", false) + c.Next() + return + } + verifiedUserID, ok := session.Get(secureVerificationUserSessionKey).(int) + if !ok || verifiedUserID != userId { + clearSecureVerificationSession(session) c.Set("secure_verified", false) c.Next() return } elapsed := time.Now().Unix() - verifiedAt - if elapsed >= SecureVerificationTimeout { + if elapsed < 0 || elapsed >= SecureVerificationTimeout { clearSecureVerificationSession(session) c.Set("secure_verified", false) c.Next() return } + verifiedMethod, ok := session.Get(secureVerificationMethodSessionKey).(string) + if !ok || verifiedMethod == "" || verifiedMethod == secureVerificationMethodPassword { + c.Set("secure_verified", false) + c.Next() + return + } c.Set("secure_verified", true) c.Set("secure_verified_at", verifiedAt) diff --git a/middleware/secure_verification_test.go b/middleware/secure_verification_test.go new file mode 100644 index 0000000..96c5daf --- /dev/null +++ b/middleware/secure_verification_test.go @@ -0,0 +1,93 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-contrib/sessions" + "github.com/gin-contrib/sessions/cookie" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestSecureVerificationRejectsMarkerFromAnotherUser(t *testing.T) { + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("secure-verification-user-binding-test")))) + router.GET("/seed", func(c *gin.Context) { + session := sessions.Default(c) + session.Set(SecureVerificationSessionKey, time.Now().Unix()) + session.Set("secure_verified_user_id", 1001) + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + router.GET( + "/sensitive", + func(c *gin.Context) { + c.Set("id", 2002) + c.Next() + }, + SecureVerificationRequired(), + func(c *gin.Context) { + c.Status(http.StatusNoContent) + }, + ) + + seedRecorder := httptest.NewRecorder() + router.ServeHTTP(seedRecorder, httptest.NewRequest(http.MethodGet, "/seed", nil)) + require.Equal(t, http.StatusNoContent, seedRecorder.Code) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/sensitive", nil) + for _, sessionCookie := range seedRecorder.Result().Cookies() { + request.AddCookie(sessionCookie) + } + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusForbidden, recorder.Code) + require.Contains(t, recorder.Body.String(), `"code":"VERIFICATION_INVALID"`) +} + +func TestPasswordVerificationIsRestrictedToMatchingScope(t *testing.T) { + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("secure-verification-scope-test")))) + router.GET("/seed", func(c *gin.Context) { + session := sessions.Default(c) + session.Set(SecureVerificationSessionKey, time.Now().Unix()) + session.Set(secureVerificationUserSessionKey, 1001) + session.Set(secureVerificationMethodSessionKey, secureVerificationMethodPassword) + session.Set(secureVerificationScopeSessionKey, "access_token") + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + setUser := func(c *gin.Context) { + c.Set("id", 1001) + c.Next() + } + router.GET("/token", setUser, SecureVerificationRequired("access_token"), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + router.GET("/other-sensitive", setUser, SecureVerificationRequired(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + seedRecorder := httptest.NewRecorder() + router.ServeHTTP(seedRecorder, httptest.NewRequest(http.MethodGet, "/seed", nil)) + require.Equal(t, http.StatusNoContent, seedRecorder.Code) + + perform := func(path string) *httptest.ResponseRecorder { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, path, nil) + for _, sessionCookie := range seedRecorder.Result().Cookies() { + request.AddCookie(sessionCookie) + } + router.ServeHTTP(recorder, request) + return recorder + } + + require.Equal(t, http.StatusNoContent, perform("/token").Code) + otherRecorder := perform("/other-sensitive") + require.Equal(t, http.StatusForbidden, otherRecorder.Code) + require.Contains(t, otherRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`) +} diff --git a/router/api-router.go b/router/api-router.go index 1a16aab..e960195 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -62,6 +62,7 @@ func SetApiRouter(router *gin.Engine) { apiRouter.POST("/waffo-pancake/webhook/:env", anonymousRequestBodyLimit, controller.WaffoPancakeWebhook) // Universal secure verification routes + apiRouter.GET("/verify/methods", middleware.UserAuth(), middleware.DisableCache(), controller.GetVerificationMethods) apiRouter.POST("/verify", middleware.UserAuth(), middleware.CriticalRateLimit(), controller.UniversalVerify) userRoute := apiRouter.Group("/user") @@ -85,7 +86,7 @@ func SetApiRouter(router *gin.Engine) { selfRoute.GET("/models", controller.GetUserModels) selfRoute.PUT("/self", middleware.CriticalRateLimit(), controller.UpdateSelf) selfRoute.DELETE("/self", controller.DeleteSelf) - selfRoute.POST("/token", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(), controller.GenerateAccessToken) + selfRoute.POST("/token", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("access_token"), controller.GenerateAccessToken) selfRoute.GET("/passkey", controller.PasskeyStatus) selfRoute.POST("/passkey/register/begin", controller.PasskeyRegisterBegin) selfRoute.POST("/passkey/register/finish", controller.PasskeyRegisterFinish) diff --git a/router/api_router_test.go b/router/api_router_test.go index 332b25d..26601d6 100644 --- a/router/api_router_test.go +++ b/router/api_router_test.go @@ -1,9 +1,13 @@ package router import ( + "net/http" + "os" "testing" + "github.com/MAX-API-Next/MAX-API/common" "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" ) func TestUserTokenRouteUsesPostOnly(t *testing.T) { @@ -33,3 +37,32 @@ func TestUserTokenRouteUsesPostOnly(t *testing.T) { t.Fatal("did not expect GET /api/user/token route to be registered") } } + +func TestUserTokenOpenAPIUsesPostOnly(t *testing.T) { + documentBytes, err := os.ReadFile("../docs/openapi/api.json") + require.NoError(t, err) + + var document struct { + Paths map[string]map[string]any `json:"paths"` + } + require.NoError(t, common.Unmarshal(documentBytes, &document)) + + tokenPath, ok := document.Paths["/api/user/token"] + require.True(t, ok, "expected /api/user/token in OpenAPI document") + require.Contains(t, tokenPath, "post") + require.NotContains(t, tokenPath, "get") +} + +func TestVerificationMethodsRouteIsRegistered(t *testing.T) { + gin.SetMode(gin.TestMode) + + engine := gin.New() + SetApiRouter(engine) + + for _, route := range engine.Routes() { + if route.Method == http.MethodGet && route.Path == "/api/verify/methods" { + return + } + } + t.Fatal("expected GET /api/verify/methods route to be registered") +} diff --git a/web/default/src/features/auth/secure-verification/api.ts b/web/default/src/features/auth/secure-verification/api.ts index 37c661a..4d5415e 100644 --- a/web/default/src/features/auth/secure-verification/api.ts +++ b/web/default/src/features/auth/secure-verification/api.ts @@ -16,40 +16,53 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { api, get2FAStatus } from '@/lib/api' +import i18next from 'i18next' +import { api } from '@/lib/api' import { buildAssertionResult, prepareCredentialRequestOptions, isPasskeySupported as detectPasskeySupport, } from '@/lib/passkey' -import { - beginPasskeyVerification, - finishPasskeyVerification, - getPasskeyStatus, -} from '../passkey' -import type { VerificationMethod, VerificationMethods } from './types' +import { beginPasskeyVerification, finishPasskeyVerification } from '../passkey' +import type { + VerificationMethod, + VerificationMethods, + VerificationScope, +} from './types' + +interface VerificationMethodsApiResponse { + success: boolean + message?: string + data?: { + has_2fa: boolean + has_passkey: boolean + has_password: boolean + } +} /** * Fetch available verification methods for the current user. */ -export async function checkVerificationMethods(): Promise { +export async function checkVerificationMethods( + scope?: VerificationScope +): Promise { try { - const [twoFAResponse, passkeyResponse, passkeySupported] = - await Promise.all([ - get2FAStatus(), - getPasskeyStatus(), - detectPasskeySupport(), - ]) - - const has2FA = - Boolean(twoFAResponse?.success) && Boolean(twoFAResponse?.data?.enabled) - const hasPasskey = - Boolean(passkeyResponse?.success) && - Boolean(passkeyResponse?.data?.enabled) + const [methodsResponse, passkeySupported] = await Promise.all([ + api.get('/api/verify/methods', { + params: scope ? { scope } : undefined, + }), + detectPasskeySupport(), + ]) + const methods = methodsResponse.data?.data return { - has2FA, - hasPasskey, + has2FA: + Boolean(methodsResponse.data?.success) && Boolean(methods?.has_2fa), + hasPasskey: + Boolean(methodsResponse.data?.success) && Boolean(methods?.has_passkey), + hasPassword: + Boolean(methodsResponse.data?.success) && + Boolean(methods?.has_password), passkeySupported, } } catch (error) { @@ -58,6 +71,7 @@ export async function checkVerificationMethods(): Promise { return { has2FA: false, hasPasskey: false, + hasPassword: false, passkeySupported: false, } } @@ -68,22 +82,47 @@ export async function checkVerificationMethods(): Promise { */ export async function verify( method: VerificationMethod, - code?: string + code?: string, + scope?: VerificationScope ): Promise { switch (method) { case '2fa': - return verifyTwoFA(code) + return verifyTwoFA(code, scope) case 'passkey': - return verifyPasskey() + return verifyPasskey(scope) + case 'password': + return verifyPassword(code, scope) default: throw new Error(`Unsupported verification method: ${method}`) } } +async function verifyPassword( + password?: string | null, + scope?: VerificationScope +): Promise { + if (!password) { + throw new Error(i18next.t('Please enter your password')) + } + + const res = await api.post('/api/verify', { + method: 'password', + password, + scope, + }) + + if (!res.data?.success) { + throw new Error(res.data?.message || 'Verification failed') + } +} + /** * Perform 2FA verification flow. */ -async function verifyTwoFA(code?: string | null): Promise { +async function verifyTwoFA( + code?: string | null, + scope?: VerificationScope +): Promise { const trimmed = code?.trim() if (!trimmed) { throw new Error('Please enter the verification code or backup code') @@ -92,6 +131,7 @@ async function verifyTwoFA(code?: string | null): Promise { const res = await api.post('/api/verify', { method: '2fa', code: trimmed, + scope, }) if (!res.data?.success) { @@ -102,7 +142,7 @@ async function verifyTwoFA(code?: string | null): Promise { /** * Perform Passkey verification flow. */ -async function verifyPasskey(): Promise { +async function verifyPasskey(scope?: VerificationScope): Promise { if (typeof navigator === 'undefined' || !navigator.credentials) { throw new Error('Passkey verification is not supported in this environment') } @@ -137,6 +177,7 @@ async function verifyPasskey(): Promise { const verifyResponse = await api.post('/api/verify', { method: 'passkey', + scope, }) if (!verifyResponse.data?.success) { diff --git a/web/default/src/features/auth/secure-verification/components/secure-verification-dialog.tsx b/web/default/src/features/auth/secure-verification/components/secure-verification-dialog.tsx index 9551bc9..034167f 100644 --- a/web/default/src/features/auth/secure-verification/components/secure-verification-dialog.tsx +++ b/web/default/src/features/auth/secure-verification/components/secure-verification-dialog.tsx @@ -17,7 +17,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ import { useMemo } from 'react' -import { ShieldCheck, KeyRound, Loader2 } from 'lucide-react' +import { ShieldCheck, KeyRound, Loader2, LockKeyhole } from 'lucide-react' import { useTranslation } from 'react-i18next' import { Button } from '@/components/ui/button' import { @@ -62,6 +62,7 @@ export function SecureVerificationDialog({ const tabs: VerificationMethod[] = [] if (methods.has2FA) tabs.push('2fa') if (methods.hasPasskey && methods.passkeySupported) tabs.push('passkey') + if (methods.hasPassword) tabs.push('password') return tabs }, [methods]) @@ -82,13 +83,14 @@ export function SecureVerificationDialog({ const handleVerify = () => { if (!activeMethod) return - const payload = activeMethod === '2fa' ? state.code : undefined + const payload = activeMethod === 'passkey' ? undefined : state.code onVerify(activeMethod, payload) } const verifyDisabled = state.loading || - (activeMethod === '2fa' && (!state.code.trim() || state.code.length < 6)) + (activeMethod === '2fa' && (!state.code.trim() || state.code.length < 6)) || + (activeMethod === 'password' && !state.code) return ( @@ -136,6 +138,9 @@ export function SecureVerificationDialog({ {methods.hasPasskey && methods.passkeySupported && ( {t('Passkey')} )} + {methods.hasPassword && ( + {t('Password')} + )} @@ -183,6 +188,32 @@ export function SecureVerificationDialog({

)}
+ + +
+ +

+ {t( + 'Enter your account password to confirm this sensitive action.' + )} +

+
+ onCodeChange(event.target.value)} + placeholder={t('Password')} + disabled={state.loading} + autoFocus={activeMethod === 'password'} + onKeyDown={(event) => { + if (event.key === 'Enter' && !verifyDisabled) { + event.preventDefault() + handleVerify() + } + }} + /> +
)} diff --git a/web/default/src/features/auth/secure-verification/hooks/use-secure-verification.ts b/web/default/src/features/auth/secure-verification/hooks/use-secure-verification.ts index abc69fd..eaf4b74 100644 --- a/web/default/src/features/auth/secure-verification/hooks/use-secure-verification.ts +++ b/web/default/src/features/auth/secure-verification/hooks/use-secure-verification.ts @@ -16,7 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useCallback, useEffect, useMemo, useState } from 'react' +import { useCallback, useMemo, useState } from 'react' import i18next from 'i18next' import { toast } from 'sonner' import { @@ -24,12 +24,14 @@ import { isVerificationRequiredError, } from '@/lib/secure-verification' import { checkVerificationMethods, verify } from '../api' +import { selectVerificationMethod } from '../method-selection' import type { SecureVerificationState, StartVerificationOptions, UseSecureVerificationOptions, VerificationMethod, VerificationMethods, + VerificationScope, } from '../types' type ApiCall = (() => Promise) | null @@ -41,6 +43,7 @@ interface InternalState extends SecureVerificationState { const defaultMethods: VerificationMethods = { has2FA: false, hasPasskey: false, + hasPassword: false, passkeySupported: false, } @@ -62,15 +65,14 @@ export function useSecureVerification( const [state, setState] = useState(initialState) const [open, setOpen] = useState(false) - const fetchVerificationMethods = useCallback(async () => { - const result = await checkVerificationMethods() - setMethods(result) - return result - }, []) - - useEffect(() => { - fetchVerificationMethods() - }, [fetchVerificationMethods]) + const fetchVerificationMethods = useCallback( + async (scope?: VerificationScope) => { + const result = await checkVerificationMethods(scope) + setMethods(result) + return result + }, + [] + ) const reset = useCallback(() => { setState(initialState) @@ -82,10 +84,14 @@ export function useSecureVerification( apiCall: () => Promise, config: StartVerificationOptions = {} ) => { - const { preferredMethod, title, description } = config - const availableMethods = await fetchVerificationMethods() - - if (!availableMethods.has2FA && !availableMethods.hasPasskey) { + const { preferredMethod, title, description, scope } = config + const availableMethods = await fetchVerificationMethods(scope) + + if ( + !availableMethods.has2FA && + !availableMethods.hasPasskey && + !availableMethods.hasPassword + ) { toast.error( i18next.t( 'Please enable Two-factor Authentication or Passkey before proceeding' @@ -99,14 +105,10 @@ export function useSecureVerification( return false } - let defaultMethod: VerificationMethod | null = preferredMethod ?? null - if (!defaultMethod) { - if (availableMethods.hasPasskey && availableMethods.passkeySupported) { - defaultMethod = 'passkey' - } else if (availableMethods.has2FA) { - defaultMethod = '2fa' - } - } + const defaultMethod = selectVerificationMethod( + availableMethods, + preferredMethod + ) setState((prev) => ({ ...prev, @@ -114,6 +116,7 @@ export function useSecureVerification( method: defaultMethod, title, description, + scope, })) setOpen(true) return true @@ -137,7 +140,7 @@ export function useSecureVerification( setState((prev) => ({ ...prev, loading: true })) try { - await verify(actualMethod, code ?? state.code) + await verify(actualMethod, code ?? state.code, state.scope) const result = await state.apiCall() if (successMessage) { @@ -204,7 +207,7 @@ export function useSecureVerification( if (method === 'passkey') { return methods.hasPasskey && methods.passkeySupported } - return false + return methods.hasPassword }, [methods] ) @@ -212,6 +215,7 @@ export function useSecureVerification( const recommendedMethod = useMemo(() => { if (methods.hasPasskey && methods.passkeySupported) return 'passkey' if (methods.has2FA) return '2fa' + if (methods.hasPassword) return 'password' return null }, [methods]) @@ -230,7 +234,7 @@ export function useSecureVerification( fetchVerificationMethods, canUseMethod, recommendedMethod, - hasAnyMethod: methods.has2FA || methods.hasPasskey, + hasAnyMethod: methods.has2FA || methods.hasPasskey || methods.hasPassword, isLoading: state.loading, currentMethod: state.method, code: state.code, diff --git a/web/default/src/features/auth/secure-verification/method-selection.ts b/web/default/src/features/auth/secure-verification/method-selection.ts new file mode 100644 index 0000000..45bfe0e --- /dev/null +++ b/web/default/src/features/auth/secure-verification/method-selection.ts @@ -0,0 +1,20 @@ +import type { VerificationMethod, VerificationMethods } from './types' + +export function selectVerificationMethod( + methods: VerificationMethods, + preferredMethod?: VerificationMethod +): VerificationMethod | null { + const isAvailable = (method: VerificationMethod) => { + if (method === '2fa') return methods.has2FA + if (method === 'passkey') { + return methods.hasPasskey && methods.passkeySupported + } + return methods.hasPassword + } + + if (preferredMethod && isAvailable(preferredMethod)) return preferredMethod + if (methods.hasPasskey && methods.passkeySupported) return 'passkey' + if (methods.has2FA) return '2fa' + if (methods.hasPassword) return 'password' + return null +} diff --git a/web/default/src/features/auth/secure-verification/types.ts b/web/default/src/features/auth/secure-verification/types.ts index 2b9b6e1..60c3cae 100644 --- a/web/default/src/features/auth/secure-verification/types.ts +++ b/web/default/src/features/auth/secure-verification/types.ts @@ -16,11 +16,13 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -export type VerificationMethod = '2fa' | 'passkey' +export type VerificationMethod = '2fa' | 'passkey' | 'password' +export type VerificationScope = 'access_token' export interface VerificationMethods { has2FA: boolean hasPasskey: boolean + hasPassword: boolean passkeySupported: boolean } @@ -30,6 +32,7 @@ export interface SecureVerificationState { code: string title?: string description?: string + scope?: VerificationScope } export interface UseSecureVerificationOptions { @@ -43,4 +46,5 @@ export interface StartVerificationOptions { preferredMethod?: VerificationMethod title?: string description?: string + scope?: VerificationScope } diff --git a/web/default/src/features/profile/components/dialogs/access-token-dialog.tsx b/web/default/src/features/profile/components/dialogs/access-token-dialog.tsx index c1c2f81..b285725 100644 --- a/web/default/src/features/profile/components/dialogs/access-token-dialog.tsx +++ b/web/default/src/features/profile/components/dialogs/access-token-dialog.tsx @@ -20,10 +20,6 @@ import { useCallback, useEffect } from 'react' import { RefreshCw, Loader2 } from 'lucide-react' import { useTranslation } from 'react-i18next' import { Button } from '@/components/ui/button' -import { - SecureVerificationDialog, - useSecureVerification, -} from '@/features/auth/secure-verification' import { Dialog, DialogContent, @@ -35,6 +31,10 @@ import { import { Input } from '@/components/ui/input' import { Label } from '@/components/ui/label' import { CopyButton } from '@/components/copy-button' +import { + SecureVerificationDialog, + useSecureVerification, +} from '@/features/auth/secure-verification' import { useAccessToken } from '../../hooks' // ============================================================================ @@ -68,6 +68,7 @@ export function AccessTokenDialog({ () => generate({ rethrowVerificationRequired: true }), { preferredMethod: 'passkey', + scope: 'access_token', title: t('Security verification'), description: t( "Your system access token for API authentication. Keep it secure and don't share it with others." diff --git a/web/default/src/i18n/locales/en.json b/web/default/src/i18n/locales/en.json index 786fb69..85f3d3e 100644 --- a/web/default/src/i18n/locales/en.json +++ b/web/default/src/i18n/locales/en.json @@ -1694,6 +1694,7 @@ "Enter tag name (optional)": "Enter tag name (optional)", "Enter the 6-digit code from your authenticator app": "Enter the 6-digit code from your authenticator app", "Enter the 6-digit Time-based One-Time Password or 8-character backup code from your authenticator app.": "Enter the 6-digit Time-based One-Time Password or 8-character backup code from your authenticator app.", + "Enter your account password to confirm this sensitive action.": "Enter your account password to confirm this sensitive action.", "Enter the complete URL, supports": "Enter the complete URL, supports", "Enter the Coze agent ID": "Enter the Coze agent ID", "Enter the full URL of your Gotify server": "Enter the full URL of your Gotify server", @@ -3375,6 +3376,7 @@ "Please confirm that you understand the consequences": "Please confirm that you understand the consequences", "Please enable io.net model deployment service and configure an API key in System Settings.": "Please enable io.net model deployment service and configure an API key in System Settings.", "Please enable Two-factor Authentication or Passkey before proceeding": "Please enable Two-factor Authentication or Passkey before proceeding", + "Please enter your password": "Please enter your password", "Please enter a name": "Please enter a name", "Please enter a new password": "Please enter a new password", "Please enter a redemption code": "Please enter a redemption code", diff --git a/web/default/src/i18n/locales/fr.json b/web/default/src/i18n/locales/fr.json index 8de60bc..2c385fa 100644 --- a/web/default/src/i18n/locales/fr.json +++ b/web/default/src/i18n/locales/fr.json @@ -1694,6 +1694,7 @@ "Enter tag name (optional)": "Saisir le nom du tag (facultatif)", "Enter the 6-digit code from your authenticator app": "Saisir le code à 6 chiffres de votre application d'authentification", "Enter the 6-digit Time-based One-Time Password or 8-character backup code from your authenticator app.": "Saisir le mot de passe à usage unique basé sur le temps à 6 chiffres ou le code de secours à 8 caractères de votre application d'authentification.", + "Enter your account password to confirm this sensitive action.": "Saisissez le mot de passe de votre compte pour confirmer cette action sensible.", "Enter the complete URL, supports": "Saisir l'URL complète, prend en charge", "Enter the Coze agent ID": "Saisir l'ID de l'agent Coze", "Enter the full URL of your Gotify server": "Saisir l'URL complète de votre serveur Gotify", @@ -3375,6 +3376,7 @@ "Please confirm that you understand the consequences": "Veuillez confirmer que vous comprenez les conséquences", "Please enable io.net model deployment service and configure an API key in System Settings.": "Veuillez activer le service de déploiement de modèles io.net et configurer une clé API dans les Paramètres système.", "Please enable Two-factor Authentication or Passkey before proceeding": "Veuillez activer l'authentification à deux facteurs ou une Passkey avant de continuer", + "Please enter your password": "Veuillez saisir votre mot de passe", "Please enter a name": "Veuillez saisir un nom", "Please enter a new password": "Veuillez saisir un nouveau mot de passe", "Please enter a redemption code": "Veuillez saisir un code de rédemption", diff --git a/web/default/src/i18n/locales/ja.json b/web/default/src/i18n/locales/ja.json index 33870bd..869c3f5 100644 --- a/web/default/src/i18n/locales/ja.json +++ b/web/default/src/i18n/locales/ja.json @@ -1694,6 +1694,7 @@ "Enter tag name (optional)": "タグ名を入力 (オプション)", "Enter the 6-digit code from your authenticator app": "認証アプリからの6桁のコードを入力", "Enter the 6-digit Time-based One-Time Password or 8-character backup code from your authenticator app.": "認証アプリからの6桁のワンタイムパスワードまたは8文字のバックアップコードを入力してください。", + "Enter your account password to confirm this sensitive action.": "この機密操作を確認するため、アカウントのパスワードを入力してください。", "Enter the complete URL, supports": "完全なURLを入力、サポート", "Enter the Coze agent ID": "CozeエージェントIDを入力", "Enter the full URL of your Gotify server": "Gotifyサーバーの完全なURLを入力", @@ -3375,6 +3376,7 @@ "Please confirm that you understand the consequences": "結果を理解したことを確認してください", "Please enable io.net model deployment service and configure an API key in System Settings.": "システム設定で io.net モデルデプロイサービスを有効にし、API キーを設定してください。", "Please enable Two-factor Authentication or Passkey before proceeding": "続行する前に、二要素認証またはパスキーを有効にしてください", + "Please enter your password": "パスワードを入力してください", "Please enter a name": "名前を入力してください", "Please enter a new password": "新しいパスワードを入力してください", "Please enter a redemption code": "引き換えコードを入力してください", diff --git a/web/default/src/i18n/locales/ru.json b/web/default/src/i18n/locales/ru.json index f14eda6..ca8c9d7 100644 --- a/web/default/src/i18n/locales/ru.json +++ b/web/default/src/i18n/locales/ru.json @@ -1694,6 +1694,7 @@ "Enter tag name (optional)": "Введите имя тега (необязательно)", "Enter the 6-digit code from your authenticator app": "Введите 6-значный код из вашего приложения-аутентификатора", "Enter the 6-digit Time-based One-Time Password or 8-character backup code from your authenticator app.": "Введите 6-значный одноразовый пароль на основе времени или 8-значный резервный код из вашего приложения-аутентификатора.", + "Enter your account password to confirm this sensitive action.": "Введите пароль учетной записи, чтобы подтвердить это конфиденциальное действие.", "Enter the complete URL, supports": "Введите полный URL, поддерживает", "Enter the Coze agent ID": "Введите ID агента Coze", "Enter the full URL of your Gotify server": "Введите полный URL вашего сервера Gotify", @@ -3375,6 +3376,7 @@ "Please confirm that you understand the consequences": "Пожалуйста, подтвердите, что понимаете последствия", "Please enable io.net model deployment service and configure an API key in System Settings.": "Пожалуйста, включите сервис развертывания моделей io.net и настройте API-ключ в системных настройках.", "Please enable Two-factor Authentication or Passkey before proceeding": "Пожалуйста, включите двухфакторную аутентификацию или Passkey перед продолжением", + "Please enter your password": "Введите пароль", "Please enter a name": "Пожалуйста, введите имя", "Please enter a new password": "Пожалуйста, введите новый пароль", "Please enter a redemption code": "Пожалуйста, введите код активации", diff --git a/web/default/src/i18n/locales/vi.json b/web/default/src/i18n/locales/vi.json index 39af4f6..9eedd0b 100644 --- a/web/default/src/i18n/locales/vi.json +++ b/web/default/src/i18n/locales/vi.json @@ -1694,6 +1694,7 @@ "Enter tag name (optional)": "Nhập tên thẻ (tùy chọn)", "Enter the 6-digit code from your authenticator app": "Nhập mã 6 chữ số từ ứng dụng xác thực của bạn", "Enter the 6-digit Time-based One-Time Password or 8-character backup code from your authenticator app.": "Nhập Mật khẩu dùng một lần dựa trên thời gian gồm 6 chữ số hoặc mã dự phòng gồm 8 ký tự từ ứng dụng xác thực của bạn.", + "Enter your account password to confirm this sensitive action.": "Nhập mật khẩu tài khoản để xác nhận thao tác nhạy cảm này.", "Enter the complete URL, supports": "Nhập địa chỉ URL hoàn chỉnh, hỗ trợ", "Enter the Coze agent ID": "Nhập ID tác nhân Coze", "Enter the full URL of your Gotify server": "Nhập URL đầy đủ của máy chủ Gotify của bạn", @@ -3375,6 +3376,7 @@ "Please confirm that you understand the consequences": "Vui lòng xác nhận rằng bạn hiểu hậu quả", "Please enable io.net model deployment service and configure an API key in System Settings.": "Vui lòng bật dịch vụ triển khai mô hình io.net và cấu hình khóa API trong Cài đặt hệ thống.", "Please enable Two-factor Authentication or Passkey before proceeding": "Vui lòng bật Xác thực hai yếu tố hoặc Passkey trước khi tiếp tục", + "Please enter your password": "Vui lòng nhập mật khẩu", "Please enter a name": "Vui lòng nhập tên", "Please enter a new password": "Vui lòng nhập mật khẩu mới", "Please enter a redemption code": "Vui lòng nhập mã đổi thưởng", diff --git a/web/default/src/i18n/locales/zh.json b/web/default/src/i18n/locales/zh.json index 2e0e5ad..0ae9c1c 100644 --- a/web/default/src/i18n/locales/zh.json +++ b/web/default/src/i18n/locales/zh.json @@ -1694,6 +1694,7 @@ "Enter tag name (optional)": "输入标签名称(可选)", "Enter the 6-digit code from your authenticator app": "输入来自身份验证器应用的 6 位代码", "Enter the 6-digit Time-based One-Time Password or 8-character backup code from your authenticator app.": "输入来自身份验证器应用的 6 位基于时间的单次密码或 8 位备份代码。", + "Enter your account password to confirm this sensitive action.": "输入账户密码以确认此敏感操作。", "Enter the complete URL, supports": "输入完整的 URL,支持", "Enter the Coze agent ID": "输入 Coze 代理 ID", "Enter the full URL of your Gotify server": "输入您的 Gotify 服务器的完整 URL", @@ -3375,6 +3376,7 @@ "Please confirm that you understand the consequences": "请确认您了解后果", "Please enable io.net model deployment service and configure an API key in System Settings.": "请先在系统设置中启用 io.net 模型部署服务,并配置 API Key。", "Please enable Two-factor Authentication or Passkey before proceeding": "请先启用双因素认证或通行密钥", + "Please enter your password": "请输入密码", "Please enter a name": "请输入名称", "Please enter a new password": "请输入新密码", "Please enter a redemption code": "请输入兑换码", diff --git a/web/default/tests/secure-verification-method-selection.test.ts b/web/default/tests/secure-verification-method-selection.test.ts new file mode 100644 index 0000000..f70a3a1 --- /dev/null +++ b/web/default/tests/secure-verification-method-selection.test.ts @@ -0,0 +1,32 @@ +import assert from 'node:assert/strict' +import { test } from 'bun:test' + +import { selectVerificationMethod } from '../src/features/auth/secure-verification/method-selection' + +test('falls back to 2FA when preferred Passkey is unavailable', () => { + const method = selectVerificationMethod( + { + has2FA: true, + hasPasskey: false, + hasPassword: false, + passkeySupported: false, + }, + 'passkey' + ) + + assert.equal(method, '2fa') +}) + +test('uses password only when it is the available fallback', () => { + const method = selectVerificationMethod( + { + has2FA: false, + hasPasskey: false, + hasPassword: true, + passkeySupported: false, + }, + 'passkey' + ) + + assert.equal(method, 'password') +}) From 1528dc989d72d701b45c0040cd21f7f62a1ed757 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 13 Jul 2026 02:48:19 +0800 Subject: [PATCH 8/9] v1.0.4 --- controller/secure_verification.go | 2 + docs/openapi/api.json | 55 +++++++++++- middleware/rate-limit.go | 9 ++ router/api-router.go | 2 +- router/api_router_test.go | 135 ++++++++++++++++++++++++++++++ 5 files changed, 201 insertions(+), 2 deletions(-) diff --git a/controller/secure_verification.go b/controller/secure_verification.go index 0d5727e..64e1e61 100644 --- a/controller/secure_verification.go +++ b/controller/secure_verification.go @@ -51,6 +51,8 @@ type secureVerificationMethods struct { } func loadSecureVerificationMethods(user *model.User, allowPassword bool) (secureVerificationMethods, error) { + // PasswordLoginEnabled gates new password logins; it does not invalidate an + // existing credential used to re-verify an already authenticated session. methods := secureVerificationMethods{ hasPassword: user != nil && user.Password != "", } diff --git a/docs/openapi/api.json b/docs/openapi/api.json index e277abf..01f39ce 100644 --- a/docs/openapi/api.json +++ b/docs/openapi/api.json @@ -2294,6 +2294,46 @@ "安全验证" ], "parameters": [], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "type": "object", + "required": [ + "method" + ], + "properties": { + "method": { + "type": "string", + "description": "安全验证方式", + "enum": [ + "2fa", + "passkey", + "password" + ] + }, + "code": { + "type": "string", + "description": "2FA验证码或备用码,仅method=2fa时使用" + }, + "password": { + "type": "string", + "format": "password", + "description": "当前账户密码,仅method=password且scope=access_token时使用" + }, + "scope": { + "type": "string", + "description": "验证作用域;密码验证仅支持access_token", + "enum": [ + "access_token" + ] + } + } + } + } + } + }, "responses": { "200": { "description": "成功", @@ -2318,7 +2358,20 @@ "tags": [ "安全验证" ], - "parameters": [], + "parameters": [ + { + "name": "scope", + "in": "query", + "required": false, + "description": "验证作用域;传access_token时,无2FA或Passkey的用户可使用密码验证", + "schema": { + "type": "string", + "enum": [ + "access_token" + ] + } + } + ], "responses": { "200": { "description": "成功", diff --git a/middleware/rate-limit.go b/middleware/rate-limit.go index 8b9dd71..c356caf 100644 --- a/middleware/rate-limit.go +++ b/middleware/rate-limit.go @@ -108,6 +108,15 @@ func CriticalRateLimit() func(c *gin.Context) { return defNext } +// UserCriticalRateLimit applies the critical-operation limit per authenticated +// user. Use it together with CriticalRateLimit so both account and IP bursts are bounded. +func UserCriticalRateLimit() func(c *gin.Context) { + if !common.CriticalRateLimitEnable { + return defNext + } + return userRateLimitFactory(common.CriticalRateLimitNum, common.CriticalRateLimitDuration, "CT") +} + func DownloadRateLimit() func(c *gin.Context) { return rateLimitFactory(common.DownloadRateLimitNum, common.DownloadRateLimitDuration, "DW") } diff --git a/router/api-router.go b/router/api-router.go index e960195..809c2ce 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -63,7 +63,7 @@ func SetApiRouter(router *gin.Engine) { // Universal secure verification routes apiRouter.GET("/verify/methods", middleware.UserAuth(), middleware.DisableCache(), controller.GetVerificationMethods) - apiRouter.POST("/verify", middleware.UserAuth(), middleware.CriticalRateLimit(), controller.UniversalVerify) + apiRouter.POST("/verify", middleware.UserAuth(), middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), controller.UniversalVerify) userRoute := apiRouter.Group("/user") { diff --git a/router/api_router_test.go b/router/api_router_test.go index 26601d6..5d67018 100644 --- a/router/api_router_test.go +++ b/router/api_router_test.go @@ -2,12 +2,19 @@ package router import ( "net/http" + "net/http/httptest" "os" + "strings" "testing" "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/gin-contrib/sessions" + "github.com/gin-contrib/sessions/cookie" "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" "github.com/stretchr/testify/require" + "gorm.io/gorm" ) func TestUserTokenRouteUsesPostOnly(t *testing.T) { @@ -66,3 +73,131 @@ func TestVerificationMethodsRouteIsRegistered(t *testing.T) { } t.Fatal("expected GET /api/verify/methods route to be registered") } + +func TestUniversalVerifyRateLimitFollowsUserAcrossIPs(t *testing.T) { + gin.SetMode(gin.TestMode) + + oldDB := model.DB + oldLogDB := model.LOG_DB + oldRedisEnabled := common.RedisEnabled + oldMemoryCacheEnabled := common.MemoryCacheEnabled + oldGlobalAPIRateLimitEnable := common.GlobalApiRateLimitEnable + oldCriticalRateLimitEnable := common.CriticalRateLimitEnable + oldCriticalRateLimitNum := common.CriticalRateLimitNum + oldCriticalRateLimitDuration := common.CriticalRateLimitDuration + + common.RedisEnabled = false + common.MemoryCacheEnabled = false + common.GlobalApiRateLimitEnable = false + common.CriticalRateLimitEnable = true + common.CriticalRateLimitNum = 1 + common.CriticalRateLimitDuration = 20 * 60 + + db, err := gorm.Open(sqlite.Open("file:router_verify_rate_limit?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + model.LOG_DB = db + require.NoError(t, db.AutoMigrate( + &model.User{}, + &model.TwoFA{}, + &model.TwoFABackupCode{}, + &model.PasskeyCredential{}, + )) + user := model.User{ + Id: 987654, + Username: "verify-rate-limit-user", + Password: "hashed-password", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + t.Cleanup(func() { + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + model.DB = oldDB + model.LOG_DB = oldLogDB + common.RedisEnabled = oldRedisEnabled + common.MemoryCacheEnabled = oldMemoryCacheEnabled + common.GlobalApiRateLimitEnable = oldGlobalAPIRateLimitEnable + common.CriticalRateLimitEnable = oldCriticalRateLimitEnable + common.CriticalRateLimitNum = oldCriticalRateLimitNum + common.CriticalRateLimitDuration = oldCriticalRateLimitDuration + }) + + engine := gin.New() + engine.Use(sessions.Sessions("session", cookie.NewStore([]byte("verify-rate-limit-session")))) + engine.GET("/test/login", func(c *gin.Context) { + session := sessions.Default(c) + session.Set("id", user.Id) + session.Set("username", user.Username) + session.Set("role", user.Role) + session.Set("status", user.Status) + session.Set("group", user.Group) + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + SetApiRouter(engine) + + loginRecorder := httptest.NewRecorder() + engine.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/test/login", nil)) + require.Equal(t, http.StatusNoContent, loginRecorder.Code) + + performVerify := func(remoteAddr string) *httptest.ResponseRecorder { + recorder := httptest.NewRecorder() + request := httptest.NewRequest( + http.MethodPost, + "/api/verify", + strings.NewReader(`{"method":"invalid"}`), + ) + request.RemoteAddr = remoteAddr + request.Header.Set("Content-Type", "application/json") + for _, sessionCookie := range loginRecorder.Result().Cookies() { + request.AddCookie(sessionCookie) + } + engine.ServeHTTP(recorder, request) + return recorder + } + + require.Equal(t, http.StatusOK, performVerify("192.0.2.10:1234").Code) + require.Equal(t, http.StatusTooManyRequests, performVerify("192.0.2.11:1234").Code) +} + +func TestSecureVerificationOpenAPIIncludesScopeAndRequestBody(t *testing.T) { + documentBytes, err := os.ReadFile("../docs/openapi/api.json") + require.NoError(t, err) + + var document struct { + Paths map[string]map[string]map[string]any `json:"paths"` + } + require.NoError(t, common.Unmarshal(documentBytes, &document)) + + methodsOperation := document.Paths["/api/verify/methods"]["get"] + parameters, ok := methodsOperation["parameters"].([]any) + require.True(t, ok) + require.Condition(t, func() bool { + for _, parameter := range parameters { + value, isMap := parameter.(map[string]any) + if !isMap || value["name"] != "scope" || value["in"] != "query" { + continue + } + schema, schemaOK := value["schema"].(map[string]any) + values, valuesOK := schema["enum"].([]any) + return schemaOK && valuesOK && len(values) == 1 && values[0] == "access_token" + } + return false + }, "expected scope query parameter for GET /api/verify/methods") + + verifyOperation := document.Paths["/api/verify"]["post"] + requestBody, ok := verifyOperation["requestBody"].(map[string]any) + require.True(t, ok, "expected requestBody for POST /api/verify") + content := requestBody["content"].(map[string]any) + mediaType := content["application/json"].(map[string]any) + schema := mediaType["schema"].(map[string]any) + properties := schema["properties"].(map[string]any) + for _, property := range []string{"method", "code", "password", "scope"} { + require.Contains(t, properties, property) + } +} From b688fb59c7a7f93cbff84581c16a1934d4933d63 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 13 Jul 2026 02:54:44 +0800 Subject: [PATCH 9/9] v1.0.4-1 --- router/api_router_test.go | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/router/api_router_test.go b/router/api_router_test.go index 5d67018..10ac911 100644 --- a/router/api_router_test.go +++ b/router/api_router_test.go @@ -1,10 +1,12 @@ package router import ( + "fmt" "net/http" "net/http/httptest" "os" "strings" + "sync/atomic" "testing" "github.com/MAX-API-Next/MAX-API/common" @@ -17,6 +19,8 @@ import ( "gorm.io/gorm" ) +var universalVerifyRateLimitTestRun uint64 + func TestUserTokenRouteUsesPostOnly(t *testing.T) { gin.SetMode(gin.TestMode) @@ -76,6 +80,7 @@ func TestVerificationMethodsRouteIsRegistered(t *testing.T) { func TestUniversalVerifyRateLimitFollowsUserAcrossIPs(t *testing.T) { gin.SetMode(gin.TestMode) + testRun := atomic.AddUint64(&universalVerifyRateLimitTestRun, 1) oldDB := model.DB oldLogDB := model.LOG_DB @@ -104,8 +109,8 @@ func TestUniversalVerifyRateLimitFollowsUserAcrossIPs(t *testing.T) { &model.PasskeyCredential{}, )) user := model.User{ - Id: 987654, - Username: "verify-rate-limit-user", + Id: 987654 + int(testRun), + Username: fmt.Sprintf("verify-rate-limit-user-%d", testRun), Password: "hashed-password", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, @@ -161,8 +166,8 @@ func TestUniversalVerifyRateLimitFollowsUserAcrossIPs(t *testing.T) { return recorder } - require.Equal(t, http.StatusOK, performVerify("192.0.2.10:1234").Code) - require.Equal(t, http.StatusTooManyRequests, performVerify("192.0.2.11:1234").Code) + require.Equal(t, http.StatusOK, performVerify(fmt.Sprintf("[2001:db8:%x::1]:1234", testRun)).Code) + require.Equal(t, http.StatusTooManyRequests, performVerify(fmt.Sprintf("[2001:db8:%x::2]:1234", testRun)).Code) } func TestSecureVerificationOpenAPIIncludesScopeAndRequestBody(t *testing.T) {