mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-08-31 01:13:06 +08:00
feat: 增加 API Key 计费倍率自省接口
This commit is contained in:
@@ -254,7 +254,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
|
||||
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
|
||||
userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig)
|
||||
gatewayHandler := handler.NewGatewayHandler(gatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService)
|
||||
gatewayHandler := handler.NewGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService)
|
||||
openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, configConfig)
|
||||
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
|
||||
totpHandler := handler.NewTotpHandler(totpService)
|
||||
|
||||
@@ -39,6 +39,7 @@ var gatewayCompatibilityMetricsLogCounter atomic.Uint64
|
||||
// GatewayHandler handles API gateway requests
|
||||
type GatewayHandler struct {
|
||||
gatewayService *service.GatewayService
|
||||
openAIGatewayService *service.OpenAIGatewayService
|
||||
geminiCompatService *service.GeminiMessagesCompatService
|
||||
antigravityGatewayService *service.AntigravityGatewayService
|
||||
userService *service.UserService
|
||||
@@ -59,6 +60,7 @@ type GatewayHandler struct {
|
||||
// NewGatewayHandler creates a new GatewayHandler
|
||||
func NewGatewayHandler(
|
||||
gatewayService *service.GatewayService,
|
||||
openAIGatewayService *service.OpenAIGatewayService,
|
||||
geminiCompatService *service.GeminiMessagesCompatService,
|
||||
antigravityGatewayService *service.AntigravityGatewayService,
|
||||
userService *service.UserService,
|
||||
@@ -94,6 +96,7 @@ func NewGatewayHandler(
|
||||
|
||||
return &GatewayHandler{
|
||||
gatewayService: gatewayService,
|
||||
openAIGatewayService: openAIGatewayService,
|
||||
geminiCompatService: geminiCompatService,
|
||||
antigravityGatewayService: antigravityGatewayService,
|
||||
userService: userService,
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const keyBillingInfoSchemaVersion = 1
|
||||
|
||||
type keyBillingInfoResponse struct {
|
||||
Object string `json:"object"`
|
||||
SchemaVersion int `json:"schema_version"`
|
||||
BillingScope string `json:"billing_scope"`
|
||||
GroupRateMultiplier float64 `json:"group_rate_multiplier"`
|
||||
UserRateMultiplier *float64 `json:"user_rate_multiplier,omitempty"`
|
||||
ResolvedRateMultiplier float64 `json:"resolved_rate_multiplier"`
|
||||
PeakRateEnabled bool `json:"peak_rate_enabled"`
|
||||
PeakStart *string `json:"peak_start,omitempty"`
|
||||
PeakEnd *string `json:"peak_end,omitempty"`
|
||||
PeakRateMultiplier *float64 `json:"peak_rate_multiplier,omitempty"`
|
||||
AppliedPeakMultiplier *float64 `json:"applied_peak_multiplier,omitempty"`
|
||||
EffectiveRateMultiplier float64 `json:"effective_rate_multiplier"`
|
||||
Timezone *string `json:"timezone,omitempty"`
|
||||
ObservedAt time.Time `json:"observed_at"`
|
||||
}
|
||||
|
||||
// KeyBillingInfo returns the token billing multiplier effective for the authenticated API key.
|
||||
// GET /v1/sub2api/billing
|
||||
func (h *GatewayHandler) KeyBillingInfo(c *gin.Context) {
|
||||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||||
if !ok {
|
||||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||||
return
|
||||
}
|
||||
if h.cfg != nil && h.cfg.RunMode == config.RunModeSimple {
|
||||
h.errorResponse(c, http.StatusNotFound, "not_found_error", "Billing information is not supported in simple mode")
|
||||
return
|
||||
}
|
||||
if apiKey.GroupID == nil {
|
||||
h.errorResponse(c, http.StatusForbidden, "permission_error", "API key is not assigned to a group")
|
||||
return
|
||||
}
|
||||
if apiKey.Group == nil {
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Billing information is unavailable")
|
||||
return
|
||||
}
|
||||
|
||||
resolvedRate, ok := h.resolveKeyBillingRate(c, apiKey)
|
||||
if !ok {
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Billing information is unavailable")
|
||||
return
|
||||
}
|
||||
|
||||
c.Header("Cache-Control", "no-store")
|
||||
c.JSON(http.StatusOK, buildKeyBillingInfo(apiKey, resolvedRate, timezone.Now()))
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) resolveKeyBillingRate(c *gin.Context, apiKey *service.APIKey) (float64, bool) {
|
||||
groupRate := apiKey.Group.RateMultiplier
|
||||
switch apiKey.Group.Platform {
|
||||
case service.PlatformOpenAI, service.PlatformGrok:
|
||||
if h.openAIGatewayService == nil {
|
||||
return 0, false
|
||||
}
|
||||
return h.openAIGatewayService.ResolveUserGroupRateMultiplier(c.Request.Context(), apiKey.UserID, *apiKey.GroupID, groupRate), true
|
||||
default:
|
||||
if h.gatewayService == nil {
|
||||
return 0, false
|
||||
}
|
||||
return h.gatewayService.ResolveUserGroupRateMultiplier(c.Request.Context(), apiKey.UserID, *apiKey.GroupID, groupRate), true
|
||||
}
|
||||
}
|
||||
|
||||
func buildKeyBillingInfo(apiKey *service.APIKey, resolvedRate float64, now time.Time) keyBillingInfoResponse {
|
||||
groupRate := apiKey.Group.RateMultiplier
|
||||
var userRate *float64
|
||||
if resolvedRate != groupRate {
|
||||
userRate = &resolvedRate
|
||||
}
|
||||
appliedPeak := apiKey.Group.PeakMultiplierAt(now)
|
||||
|
||||
response := keyBillingInfoResponse{
|
||||
Object: "sub2api.key_billing",
|
||||
SchemaVersion: keyBillingInfoSchemaVersion,
|
||||
BillingScope: "token",
|
||||
GroupRateMultiplier: groupRate,
|
||||
UserRateMultiplier: userRate,
|
||||
ResolvedRateMultiplier: resolvedRate,
|
||||
PeakRateEnabled: apiKey.Group.PeakRateEnabled,
|
||||
EffectiveRateMultiplier: resolvedRate * appliedPeak,
|
||||
ObservedAt: now.UTC(),
|
||||
}
|
||||
if apiKey.Group.PeakRateEnabled {
|
||||
response.PeakStart = &apiKey.Group.PeakStart
|
||||
response.PeakEnd = &apiKey.Group.PeakEnd
|
||||
response.PeakRateMultiplier = &apiKey.Group.PeakRateMultiplier
|
||||
response.AppliedPeakMultiplier = &appliedPeak
|
||||
tz := timezone.Location().String()
|
||||
response.Timezone = &tz
|
||||
}
|
||||
return response
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type keyBillingUserGroupRateRepo struct {
|
||||
service.UserGroupRateRepository
|
||||
rate *float64
|
||||
err error
|
||||
gotUserID int64
|
||||
gotGroupID int64
|
||||
lookupCalls int
|
||||
}
|
||||
|
||||
func (r *keyBillingUserGroupRateRepo) GetByUserAndGroup(_ context.Context, userID, groupID int64) (*float64, error) {
|
||||
r.gotUserID = userID
|
||||
r.gotGroupID = groupID
|
||||
r.lookupCalls++
|
||||
return r.rate, r.err
|
||||
}
|
||||
|
||||
func newKeyBillingHandler(repo service.UserGroupRateRepository) *GatewayHandler {
|
||||
return &GatewayHandler{
|
||||
gatewayService: newKeyBillingGatewayService(repo),
|
||||
openAIGatewayService: newKeyBillingOpenAIGatewayService(repo),
|
||||
}
|
||||
}
|
||||
|
||||
func newKeyBillingGatewayService(repo service.UserGroupRateRepository) *service.GatewayService {
|
||||
return service.NewGatewayService(
|
||||
nil, nil, nil, nil, nil, nil, repo, nil, nil, nil, nil, nil, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
}
|
||||
|
||||
func newKeyBillingOpenAIGatewayService(repo service.UserGroupRateRepository) *service.OpenAIGatewayService {
|
||||
return service.NewOpenAIGatewayService(
|
||||
nil, nil, nil, nil, nil, repo, nil, nil, nil, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
}
|
||||
|
||||
func newKeyBillingContext(apiKey *service.APIKey) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
|
||||
if apiKey != nil {
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), apiKey)
|
||||
}
|
||||
return c, w
|
||||
}
|
||||
|
||||
func TestGatewayHandlerKeyBillingInfoUsesGroupRate(t *testing.T) {
|
||||
groupID := int64(7)
|
||||
apiKey := &service.APIKey{
|
||||
UserID: 11,
|
||||
GroupID: &groupID,
|
||||
Key: "sk-sensitive-value",
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Name: "private-group-name",
|
||||
RateMultiplier: 0.75,
|
||||
},
|
||||
}
|
||||
c, w := newKeyBillingContext(apiKey)
|
||||
|
||||
newKeyBillingHandler(nil).KeyBillingInfo(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, "no-store", w.Header().Get("Cache-Control"))
|
||||
var got keyBillingInfoResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
|
||||
require.Equal(t, "sub2api.key_billing", got.Object)
|
||||
require.Equal(t, 1, got.SchemaVersion)
|
||||
require.Equal(t, "token", got.BillingScope)
|
||||
require.Equal(t, 0.75, got.GroupRateMultiplier)
|
||||
require.Nil(t, got.UserRateMultiplier)
|
||||
require.Equal(t, 0.75, got.ResolvedRateMultiplier)
|
||||
require.False(t, got.PeakRateEnabled)
|
||||
require.Nil(t, got.PeakStart)
|
||||
require.Nil(t, got.PeakEnd)
|
||||
require.Nil(t, got.PeakRateMultiplier)
|
||||
require.Nil(t, got.AppliedPeakMultiplier)
|
||||
require.Equal(t, 0.75, got.EffectiveRateMultiplier)
|
||||
require.Nil(t, got.Timezone)
|
||||
require.False(t, got.ObservedAt.IsZero())
|
||||
var fields map[string]json.RawMessage
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &fields))
|
||||
require.NotContains(t, fields, "user_rate_multiplier")
|
||||
require.NotContains(t, fields, "peak_start")
|
||||
require.NotContains(t, fields, "peak_end")
|
||||
require.NotContains(t, fields, "peak_rate_multiplier")
|
||||
require.NotContains(t, fields, "applied_peak_multiplier")
|
||||
require.NotContains(t, fields, "timezone")
|
||||
require.NotContains(t, w.Body.String(), apiKey.Key)
|
||||
require.NotContains(t, w.Body.String(), apiKey.Group.Name)
|
||||
}
|
||||
|
||||
func TestGatewayHandlerKeyBillingInfoUsesUserOverride(t *testing.T) {
|
||||
groupID := int64(7)
|
||||
userRate := 0.5
|
||||
apiKey := &service.APIKey{
|
||||
UserID: 11,
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, RateMultiplier: 0.75},
|
||||
}
|
||||
c, w := newKeyBillingContext(apiKey)
|
||||
repo := &keyBillingUserGroupRateRepo{rate: &userRate}
|
||||
|
||||
newKeyBillingHandler(repo).KeyBillingInfo(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, 1, repo.lookupCalls)
|
||||
require.Equal(t, apiKey.UserID, repo.gotUserID)
|
||||
require.Equal(t, groupID, repo.gotGroupID)
|
||||
var got keyBillingInfoResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
|
||||
require.NotNil(t, got.UserRateMultiplier)
|
||||
require.Equal(t, 0.5, *got.UserRateMultiplier)
|
||||
require.Equal(t, 0.5, got.ResolvedRateMultiplier)
|
||||
require.Equal(t, 0.5, got.EffectiveRateMultiplier)
|
||||
}
|
||||
|
||||
func TestBuildKeyBillingInfoAppliesPeakMultiplier(t *testing.T) {
|
||||
groupID := int64(7)
|
||||
apiKey := &service.APIKey{
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
RateMultiplier: 1.2,
|
||||
SubscriptionType: service.SubscriptionTypeSubscription,
|
||||
PeakRateEnabled: true,
|
||||
PeakStart: "09:00",
|
||||
PeakEnd: "18:00",
|
||||
PeakRateMultiplier: 1.5,
|
||||
},
|
||||
}
|
||||
now := time.Date(2026, time.July, 12, 10, 0, 0, 0, timezone.Location())
|
||||
userRate := 0.8
|
||||
|
||||
got := buildKeyBillingInfo(apiKey, userRate, now)
|
||||
|
||||
require.Equal(t, 1.2, got.GroupRateMultiplier)
|
||||
require.NotNil(t, got.UserRateMultiplier)
|
||||
require.Equal(t, 0.8, *got.UserRateMultiplier)
|
||||
require.Equal(t, 0.8, got.ResolvedRateMultiplier)
|
||||
require.True(t, got.PeakRateEnabled)
|
||||
require.NotNil(t, got.PeakStart)
|
||||
require.Equal(t, "09:00", *got.PeakStart)
|
||||
require.NotNil(t, got.PeakEnd)
|
||||
require.Equal(t, "18:00", *got.PeakEnd)
|
||||
require.NotNil(t, got.PeakRateMultiplier)
|
||||
require.Equal(t, 1.5, *got.PeakRateMultiplier)
|
||||
require.NotNil(t, got.AppliedPeakMultiplier)
|
||||
require.Equal(t, 1.5, *got.AppliedPeakMultiplier)
|
||||
require.InDelta(t, 1.2, got.EffectiveRateMultiplier, 1e-12)
|
||||
require.NotNil(t, got.Timezone)
|
||||
require.Equal(t, timezone.Location().String(), *got.Timezone)
|
||||
require.Equal(t, now.UTC(), got.ObservedAt)
|
||||
|
||||
encoded, err := json.Marshal(got)
|
||||
require.NoError(t, err)
|
||||
var fields map[string]json.RawMessage
|
||||
require.NoError(t, json.Unmarshal(encoded, &fields))
|
||||
for _, field := range []string{
|
||||
"user_rate_multiplier",
|
||||
"peak_start",
|
||||
"peak_end",
|
||||
"peak_rate_multiplier",
|
||||
"applied_peak_multiplier",
|
||||
"timezone",
|
||||
} {
|
||||
require.Contains(t, fields, field)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyBillingInfoJSONKeepsZeroPeakMultiplierWhenEnabled(t *testing.T) {
|
||||
groupID := int64(7)
|
||||
apiKey := &service.APIKey{
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
SubscriptionType: service.SubscriptionTypeSubscription,
|
||||
PeakRateEnabled: true,
|
||||
PeakStart: "00:00",
|
||||
PeakEnd: "23:59",
|
||||
PeakRateMultiplier: 0,
|
||||
},
|
||||
}
|
||||
now := time.Date(2026, time.July, 12, 12, 0, 0, 0, timezone.Location())
|
||||
encoded, err := json.Marshal(buildKeyBillingInfo(apiKey, apiKey.Group.RateMultiplier, now))
|
||||
require.NoError(t, err)
|
||||
|
||||
var fields map[string]json.RawMessage
|
||||
require.NoError(t, json.Unmarshal(encoded, &fields))
|
||||
require.JSONEq(t, "0", string(fields["peak_rate_multiplier"]))
|
||||
require.JSONEq(t, "0", string(fields["applied_peak_multiplier"]))
|
||||
}
|
||||
|
||||
func TestGatewayHandlerKeyBillingInfoErrorsAreSafe(t *testing.T) {
|
||||
t.Run("missing API key", func(t *testing.T) {
|
||||
c, w := newKeyBillingContext(nil)
|
||||
newKeyBillingHandler(nil).KeyBillingInfo(c)
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
})
|
||||
|
||||
t.Run("ungrouped API key", func(t *testing.T) {
|
||||
c, w := newKeyBillingContext(&service.APIKey{})
|
||||
newKeyBillingHandler(nil).KeyBillingInfo(c)
|
||||
require.Equal(t, http.StatusForbidden, w.Code)
|
||||
})
|
||||
|
||||
t.Run("missing billing service", func(t *testing.T) {
|
||||
groupID := int64(7)
|
||||
c, w := newKeyBillingContext(&service.APIKey{
|
||||
UserID: 11,
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, RateMultiplier: 1},
|
||||
})
|
||||
(&GatewayHandler{}).KeyBillingInfo(c)
|
||||
require.Equal(t, http.StatusInternalServerError, w.Code)
|
||||
})
|
||||
|
||||
t.Run("rate lookup failure matches billing fallback", func(t *testing.T) {
|
||||
groupID := int64(7)
|
||||
c, w := newKeyBillingContext(&service.APIKey{
|
||||
UserID: 11,
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, RateMultiplier: 1},
|
||||
})
|
||||
newKeyBillingHandler(&keyBillingUserGroupRateRepo{err: errors.New("database password leaked")}).KeyBillingInfo(c)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
var got keyBillingInfoResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
|
||||
require.Equal(t, 1.0, got.ResolvedRateMultiplier)
|
||||
require.NotContains(t, w.Body.String(), "database password leaked")
|
||||
})
|
||||
}
|
||||
|
||||
func TestGatewayHandlerKeyBillingInfoSharesBillingResolverCacheByPlatform(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
platform string
|
||||
openAI bool
|
||||
}{
|
||||
{name: "anthropic", platform: service.PlatformAnthropic},
|
||||
{name: "openai", platform: service.PlatformOpenAI, openAI: true},
|
||||
{name: "grok", platform: service.PlatformGrok, openAI: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
groupID := int64(7)
|
||||
oldRate, newRate := 0.5, 1.8
|
||||
repo := &keyBillingUserGroupRateRepo{rate: &oldRate}
|
||||
gatewayService := newKeyBillingGatewayService(repo)
|
||||
openAIGatewayService := newKeyBillingOpenAIGatewayService(repo)
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gatewayService,
|
||||
openAIGatewayService: openAIGatewayService,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
UserID: 11,
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: tc.platform,
|
||||
RateMultiplier: 0.75,
|
||||
},
|
||||
}
|
||||
|
||||
if tc.openAI {
|
||||
require.Equal(t, oldRate, openAIGatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier))
|
||||
} else {
|
||||
require.Equal(t, oldRate, gatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier))
|
||||
}
|
||||
repo.rate = &newRate
|
||||
|
||||
for range 2 {
|
||||
c, w := newKeyBillingContext(apiKey)
|
||||
h.KeyBillingInfo(c)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
var got keyBillingInfoResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
|
||||
require.Equal(t, oldRate, got.ResolvedRateMultiplier)
|
||||
require.Equal(t, oldRate, got.EffectiveRateMultiplier)
|
||||
}
|
||||
|
||||
var billedRate float64
|
||||
if tc.openAI {
|
||||
billedRate = openAIGatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier)
|
||||
} else {
|
||||
billedRate = gatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier)
|
||||
}
|
||||
require.Equal(t, oldRate, billedRate)
|
||||
require.Equal(t, 1, repo.lookupCalls)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -25,7 +25,8 @@ func NewAPIKeyAuthMiddleware(apiKeyService *service.APIKeyService, subscriptionS
|
||||
// - 鉴权(Authentication):验证 Key 有效性、用户状态、IP 限制 —— 始终执行
|
||||
// - 计费执行(Billing Enforcement):过期/配额/订阅/余额检查 —— skipBilling 时整块跳过
|
||||
//
|
||||
// /v1/usage 端点只需鉴权,不需要计费执行(允许过期/配额耗尽的 Key 查询自身用量)。
|
||||
// /v1/usage 和 /v1/sub2api/billing 端点只需鉴权,不需要计费执行。
|
||||
// 前者允许过期/配额耗尽的 Key 查询自身用量,后者用于读取当前 Key 的倍率配置。
|
||||
func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// ── 1. 提取 API Key ──────────────────────────────────────────
|
||||
@@ -128,6 +129,8 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
|
||||
}
|
||||
ctx := context.WithValue(c.Request.Context(), ctxkey.UserID, apiKey.User.ID)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
billingInfoRequest := c.Request.URL.Path == "/v1/sub2api/billing"
|
||||
skipBilling := c.Request.URL.Path == "/v1/usage" || billingInfoRequest
|
||||
|
||||
// ── 4. SimpleMode → early return ─────────────────────────────
|
||||
|
||||
@@ -139,20 +142,20 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
if !billingInfoRequest {
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
}
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// ── 5. 加载订阅(订阅模式时始终加载) ───────────────────────
|
||||
|
||||
// skipBilling: /v1/usage 只需鉴权,跳过所有计费执行
|
||||
skipBilling := c.Request.URL.Path == "/v1/usage"
|
||||
// ── 5. 按端点需要加载订阅 ───────────────────────────────────
|
||||
|
||||
var subscription *service.UserSubscription
|
||||
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
|
||||
|
||||
if isSubscriptionType && subscriptionService != nil {
|
||||
// 倍率自省不需要订阅数据;/v1/usage 仍保留原有订阅读取行为。
|
||||
if isSubscriptionType && subscriptionService != nil && !billingInfoRequest {
|
||||
sub, subErr := subscriptionService.GetActiveSubscription(
|
||||
c.Request.Context(),
|
||||
apiKey.User.ID,
|
||||
@@ -237,7 +240,9 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
if !billingInfoRequest {
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
|
||||
@@ -1066,6 +1066,129 @@ func TestAPIKeyAuthTouchesLastUsedInStandardMode(t *testing.T) {
|
||||
require.Equal(t, 1, touchCalls)
|
||||
}
|
||||
|
||||
func TestAPIKeyAuthBillingInfoSkipsBillingAndSideEffects(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
group := &service.Group{
|
||||
ID: 42,
|
||||
Name: "subscription",
|
||||
Status: service.StatusActive,
|
||||
Hydrated: true,
|
||||
SubscriptionType: service.SubscriptionTypeSubscription,
|
||||
}
|
||||
user := &service.User{
|
||||
ID: 7,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 0,
|
||||
Concurrency: 3,
|
||||
}
|
||||
expiredAt := time.Now().Add(-time.Hour)
|
||||
apiKey := &service.APIKey{
|
||||
ID: 100,
|
||||
UserID: user.ID,
|
||||
Key: "billing-info-auth-only",
|
||||
Status: service.StatusAPIKeyQuotaExhausted,
|
||||
User: user,
|
||||
GroupID: &group.ID,
|
||||
Group: group,
|
||||
Quota: 1,
|
||||
QuotaUsed: 1,
|
||||
ExpiresAt: &expiredAt,
|
||||
}
|
||||
|
||||
touchCalls := 0
|
||||
subscriptionCalls := 0
|
||||
apiKeyRepo := &stubApiKeyRepo{
|
||||
getByKey: func(context.Context, string) (*service.APIKey, error) {
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
updateLastUsed: func(context.Context, int64, time.Time) error {
|
||||
touchCalls++
|
||||
return nil
|
||||
},
|
||||
}
|
||||
subscriptionRepo := &stubUserSubscriptionRepo{
|
||||
getActive: func(context.Context, int64, int64) (*service.UserSubscription, error) {
|
||||
subscriptionCalls++
|
||||
return nil, service.ErrSubscriptionNotFound
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{RunMode: config.RunModeStandard}
|
||||
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
|
||||
subscriptionService := service.NewSubscriptionService(nil, subscriptionRepo, nil, nil, cfg)
|
||||
t.Cleanup(subscriptionService.Stop)
|
||||
router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
|
||||
req.Header.Set("x-api-key", apiKey.Key)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Zero(t, subscriptionCalls)
|
||||
require.Zero(t, touchCalls)
|
||||
}
|
||||
|
||||
func TestAPIKeyAuthBillingInfoSkipsLastUsedInSimpleMode(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive}
|
||||
apiKey := &service.APIKey{ID: 100, UserID: user.ID, Key: "billing-info-simple", Status: service.StatusActive, User: user}
|
||||
touchCalls := 0
|
||||
apiKeyRepo := &stubApiKeyRepo{
|
||||
getByKey: func(context.Context, string) (*service.APIKey, error) {
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
updateLastUsed: func(context.Context, int64, time.Time) error {
|
||||
touchCalls++
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
|
||||
router := newAuthTestRouter(apiKeyService, nil, cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
|
||||
req.Header.Set("x-api-key", apiKey.Key)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Zero(t, touchCalls)
|
||||
}
|
||||
|
||||
func TestAPIKeyAuthUsageStillTouchesLastUsed(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive, Balance: 10}
|
||||
apiKey := &service.APIKey{ID: 100, UserID: user.ID, Key: "usage-touch", Status: service.StatusActive, User: user}
|
||||
touchCalls := 0
|
||||
apiKeyRepo := &stubApiKeyRepo{
|
||||
getByKey: func(context.Context, string) (*service.APIKey, error) {
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
updateLastUsed: func(context.Context, int64, time.Time) error {
|
||||
touchCalls++
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{RunMode: config.RunModeStandard}
|
||||
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
|
||||
router := newAuthTestRouter(apiKeyService, nil, cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/usage", nil)
|
||||
req.Header.Set("x-api-key", apiKey.Key)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, 1, touchCalls)
|
||||
}
|
||||
|
||||
func TestAPIKeyAuthAllowsBalanceBelowMinimumReserve(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -1155,9 +1278,12 @@ func TestAPIKeyAuthRejectsExhaustedBalance(t *testing.T) {
|
||||
func newAuthTestRouter(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) *gin.Engine {
|
||||
router := gin.New()
|
||||
router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, cfg)))
|
||||
router.GET("/t", func(c *gin.Context) {
|
||||
ok := func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
}
|
||||
router.GET("/t", ok)
|
||||
router.GET("/v1/usage", ok)
|
||||
router.GET("/v1/sub2api/billing", ok)
|
||||
return router
|
||||
}
|
||||
|
||||
|
||||
@@ -114,6 +114,7 @@ func RegisterGatewayRoutes(
|
||||
gateway.Use(opsErrorLogger)
|
||||
gateway.Use(endpointNorm)
|
||||
gateway.Use(gin.HandlerFunc(apiKeyAuth))
|
||||
gateway.GET("/sub2api/billing", h.Gateway.KeyBillingInfo)
|
||||
gateway.Use(requireGroupAnthropic)
|
||||
{
|
||||
// /v1/messages: auto-route based on group platform
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/Wei-Shaw/sub2api/internal/web"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type keyBillingRouteAPIKeyRepo struct {
|
||||
service.APIKeyRepository
|
||||
apiKey *service.APIKey
|
||||
}
|
||||
|
||||
func (r *keyBillingRouteAPIKeyRepo) GetByKeyForAuth(_ context.Context, key string) (*service.APIKey, error) {
|
||||
if r.apiKey == nil || key != r.apiKey.Key {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
clone := *r.apiKey
|
||||
return &clone, nil
|
||||
}
|
||||
|
||||
type keyBillingRouteRateRepo struct {
|
||||
service.UserGroupRateRepository
|
||||
lookupCalls int
|
||||
}
|
||||
|
||||
func (r *keyBillingRouteRateRepo) GetByUserAndGroup(context.Context, int64, int64) (*float64, error) {
|
||||
r.lookupCalls++
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *keyBillingRouteRateRepo) GetRPMOverrideByUserAndGroup(context.Context, int64, int64) (*int, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func newKeyBillingRouteTestRouter(runMode string) (*gin.Engine, *keyBillingRouteRateRepo, string) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
group := &service.Group{
|
||||
ID: 42,
|
||||
Status: service.StatusActive,
|
||||
Hydrated: true,
|
||||
Platform: service.PlatformOpenAI,
|
||||
SubscriptionType: service.SubscriptionTypeStandard,
|
||||
RateMultiplier: 0.75,
|
||||
}
|
||||
user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive, Balance: 10}
|
||||
var groupID *int64
|
||||
var apiKeyGroup *service.Group
|
||||
if runMode != config.RunModeSimple {
|
||||
groupID = &group.ID
|
||||
apiKeyGroup = group
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 100,
|
||||
UserID: user.ID,
|
||||
Key: "billing-route-test-key",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
GroupID: groupID,
|
||||
Group: apiKeyGroup,
|
||||
}
|
||||
cfg := &config.Config{RunMode: runMode}
|
||||
rateRepo := &keyBillingRouteRateRepo{}
|
||||
apiKeyService := service.NewAPIKeyService(
|
||||
&keyBillingRouteAPIKeyRepo{apiKey: apiKey}, nil, nil, nil, rateRepo, nil, cfg,
|
||||
)
|
||||
gatewayService := service.NewGatewayService(
|
||||
nil, nil, nil, nil, nil, nil, rateRepo, nil, cfg, nil, nil, nil, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
openAIGatewayService := service.NewOpenAIGatewayService(
|
||||
nil, nil, nil, nil, nil, rateRepo, nil, cfg, nil, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
gatewayHandler := handler.NewGatewayHandler(
|
||||
gatewayService, openAIGatewayService, nil, nil, nil, nil, nil, nil,
|
||||
apiKeyService, nil, nil, nil, nil, cfg, nil,
|
||||
)
|
||||
|
||||
router := gin.New()
|
||||
if web.HasEmbeddedFrontend() {
|
||||
router.Use(web.ServeEmbeddedFrontend())
|
||||
}
|
||||
RegisterGatewayRoutes(
|
||||
router,
|
||||
&handler.Handlers{Gateway: gatewayHandler, OpenAIGateway: &handler.OpenAIGatewayHandler{}},
|
||||
servermiddleware.NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg),
|
||||
apiKeyService,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
cfg,
|
||||
)
|
||||
return router, rateRepo, apiKey.Key
|
||||
}
|
||||
|
||||
func TestGatewayRoutesKeyBillingInfoPathIsRegistered(t *testing.T) {
|
||||
router := newGatewayRoutesTestRouter()
|
||||
|
||||
for _, route := range router.Routes() {
|
||||
if route.Method == http.MethodGet && route.Path == "/v1/sub2api/billing" {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
t.Fatal("GET /v1/sub2api/billing should be registered")
|
||||
}
|
||||
|
||||
func TestGatewayRoutesKeyBillingInfoEndToEnd(t *testing.T) {
|
||||
t.Run("missing credentials", func(t *testing.T) {
|
||||
router, rateRepo, _ := newKeyBillingRouteTestRouter(config.RunModeStandard)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil))
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
require.Contains(t, w.Header().Get("Content-Type"), "application/json")
|
||||
require.NotContains(t, strings.ToLower(w.Body.String()), "<!doctype html>")
|
||||
require.Zero(t, rateRepo.lookupCalls)
|
||||
})
|
||||
|
||||
t.Run("standard mode", func(t *testing.T) {
|
||||
router, rateRepo, key := newKeyBillingRouteTestRouter(config.RunModeStandard)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Contains(t, w.Header().Get("Content-Type"), "application/json")
|
||||
require.Equal(t, "no-store", w.Header().Get("Cache-Control"))
|
||||
require.NotContains(t, strings.ToLower(w.Body.String()), "<!doctype html>")
|
||||
var body map[string]any
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
require.Equal(t, "sub2api.key_billing", body["object"])
|
||||
require.Equal(t, 0.75, body["effective_rate_multiplier"])
|
||||
require.Equal(t, 1, rateRepo.lookupCalls)
|
||||
})
|
||||
|
||||
t.Run("simple mode", func(t *testing.T) {
|
||||
router, rateRepo, key := newKeyBillingRouteTestRouter(config.RunModeSimple)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
|
||||
req.Header.Set("x-api-key", key)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusNotFound, w.Code)
|
||||
require.Contains(t, w.Header().Get("Content-Type"), "application/json")
|
||||
require.NotContains(t, strings.ToLower(w.Body.String()), "<!doctype html>")
|
||||
require.JSONEq(t, `{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "not_found_error",
|
||||
"message": "Billing information is not supported in simple mode"
|
||||
}
|
||||
}`, w.Body.String())
|
||||
require.Zero(t, rateRepo.lookupCalls)
|
||||
})
|
||||
}
|
||||
@@ -29,6 +29,11 @@ func (s *GatewayService) getUserGroupRateMultiplier(ctx context.Context, userID,
|
||||
return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier)
|
||||
}
|
||||
|
||||
// ResolveUserGroupRateMultiplier resolves the same cached multiplier used by usage billing.
|
||||
func (s *GatewayService) ResolveUserGroupRateMultiplier(ctx context.Context, userID, groupID int64, groupDefaultMultiplier float64) float64 {
|
||||
return s.getUserGroupRateMultiplier(ctx, userID, groupID, groupDefaultMultiplier)
|
||||
}
|
||||
|
||||
// RecordUsageInput 记录使用量的输入参数。
|
||||
// 异步 worker 只接收计费所需快照,不能持有 ParsedRequest/RequestBodyRef 这类大请求体引用。
|
||||
type RecordUsageInput struct {
|
||||
@@ -662,7 +667,7 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
}
|
||||
if apiKey.GroupID != nil && apiKey.Group != nil {
|
||||
groupDefault := apiKey.Group.RateMultiplier
|
||||
multiplier = s.getUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault)
|
||||
multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault)
|
||||
}
|
||||
// token 倍率叠加高峰因子(token 计费含图片 token,图片按次倍率不受影响)。高峰因子按请求时刻现算,
|
||||
// 不并入上面的 getUserGroupRateMultiplier,以免污染 user:group 倍率缓存。
|
||||
|
||||
@@ -98,6 +98,18 @@ func (s *OpenAIGatewayService) RecordCyberPolicyUsageLog(ctx context.Context, in
|
||||
}
|
||||
}
|
||||
|
||||
// ResolveUserGroupRateMultiplier resolves the same cached multiplier used by OpenAI usage billing.
|
||||
func (s *OpenAIGatewayService) ResolveUserGroupRateMultiplier(ctx context.Context, userID, groupID int64, groupDefaultMultiplier float64) float64 {
|
||||
if s == nil {
|
||||
return groupDefaultMultiplier
|
||||
}
|
||||
resolver := s.userGroupRateResolver
|
||||
if resolver == nil {
|
||||
resolver = newUserGroupRateResolver(nil, nil, resolveUserGroupRateCacheTTL(s.cfg), nil, "service.openai_gateway")
|
||||
}
|
||||
return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier)
|
||||
}
|
||||
|
||||
// RecordUsage records usage and deducts balance
|
||||
func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRecordUsageInput) error {
|
||||
if input == nil {
|
||||
@@ -142,11 +154,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
|
||||
multiplier = s.cfg.Default.RateMultiplier
|
||||
}
|
||||
if apiKey.GroupID != nil && apiKey.Group != nil {
|
||||
resolver := s.userGroupRateResolver
|
||||
if resolver == nil {
|
||||
resolver = newUserGroupRateResolver(nil, nil, resolveUserGroupRateCacheTTL(s.cfg), nil, "service.openai_gateway")
|
||||
}
|
||||
multiplier = resolver.Resolve(ctx, user.ID, *apiKey.GroupID, apiKey.Group.RateMultiplier)
|
||||
multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, apiKey.Group.RateMultiplier)
|
||||
}
|
||||
// token 倍率叠加高峰因子(token 计费含图片 token,图片按次倍率不受影响)。高峰因子按请求时刻现算,
|
||||
// 不并入上面的 Resolve,以免污染 user:group 倍率缓存。
|
||||
|
||||
Reference in New Issue
Block a user