diff --git a/controller/pricing.go b/controller/pricing.go index 8252327244c4..177ddca4416e 100644 --- a/controller/pricing.go +++ b/controller/pricing.go @@ -23,12 +23,19 @@ func filterPricingByUsableGroups(pricing []model.Pricing, usableGroup map[string filtered = append(filtered, item) continue } + usableEnableGroups := make([]string, 0, len(item.EnableGroup)) for _, group := range item.EnableGroup { if _, ok := usableGroup[group]; ok { - filtered = append(filtered, item) - break + usableEnableGroups = append(usableEnableGroups, group) } } + if len(usableEnableGroups) == 0 { + continue + } + // item is a copy of the shared pricing cache entry; assign a fresh + // slice so the cached EnableGroup is never mutated across requests. + item.EnableGroup = usableEnableGroups + filtered = append(filtered, item) } return filtered } diff --git a/controller/pricing_test.go b/controller/pricing_test.go new file mode 100644 index 000000000000..a97d90689cd9 --- /dev/null +++ b/controller/pricing_test.go @@ -0,0 +1,42 @@ +package controller + +import ( + "testing" + + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestFilterPricingByUsableGroupsTrimsEnableGroups(t *testing.T) { + pricing := []model.Pricing{ + {ModelName: "model-mixed", EnableGroup: []string{"default", "internal"}}, + {ModelName: "model-internal-only", EnableGroup: []string{"internal"}}, + {ModelName: "model-all", EnableGroup: []string{"all", "internal"}}, + } + usableGroup := map[string]string{"default": "default group"} + + filtered := filterPricingByUsableGroups(pricing, usableGroup) + + require.Len(t, filtered, 2) + + assert.Equal(t, "model-mixed", filtered[0].ModelName) + assert.Equal(t, []string{"default"}, filtered[0].EnableGroup, + "groups outside the user's usable groups must not be exposed") + + assert.Equal(t, "model-all", filtered[1].ModelName) + assert.Equal(t, []string{"all", "internal"}, filtered[1].EnableGroup, + "entries enabled for all groups keep their original group list") + + assert.Equal(t, []string{"default", "internal"}, pricing[0].EnableGroup, + "the shared pricing cache must stay untouched") +} + +func TestFilterPricingByUsableGroupsEmptyInputs(t *testing.T) { + pricing := []model.Pricing{ + {ModelName: "model-a", EnableGroup: []string{"default"}}, + } + + assert.Empty(t, filterPricingByUsableGroups(pricing, map[string]string{})) + assert.Empty(t, filterPricingByUsableGroups(nil, map[string]string{"default": ""})) +}