diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 86721fbac0..3700e2257e 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -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) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 1c51713656..7a23984d90 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -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, diff --git a/backend/internal/handler/gateway_key_billing.go b/backend/internal/handler/gateway_key_billing.go new file mode 100644 index 0000000000..2c075d779d --- /dev/null +++ b/backend/internal/handler/gateway_key_billing.go @@ -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 +} diff --git a/backend/internal/handler/gateway_key_billing_test.go b/backend/internal/handler/gateway_key_billing_test.go new file mode 100644 index 0000000000..ecd85bd56f --- /dev/null +++ b/backend/internal/handler/gateway_key_billing_test.go @@ -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) + }) + } +} diff --git a/backend/internal/server/middleware/api_key_auth.go b/backend/internal/server/middleware/api_key_auth.go index 082766bfd5..33bf65ea46 100644 --- a/backend/internal/server/middleware/api_key_auth.go +++ b/backend/internal/server/middleware/api_key_auth.go @@ -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() } diff --git a/backend/internal/server/middleware/api_key_auth_test.go b/backend/internal/server/middleware/api_key_auth_test.go index d9b5c7efc1..2a2e2f5588 100644 --- a/backend/internal/server/middleware/api_key_auth_test.go +++ b/backend/internal/server/middleware/api_key_auth_test.go @@ -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 } diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index f73379981e..61b4c54ced 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -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 diff --git a/backend/internal/server/routes/gateway_key_billing_test.go b/backend/internal/server/routes/gateway_key_billing_test.go new file mode 100644 index 0000000000..ef69a5eb44 --- /dev/null +++ b/backend/internal/server/routes/gateway_key_billing_test.go @@ -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()), "") + 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()), "") + 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()), "") + 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) + }) +} diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 61ab3abd2f..71236b69bd 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -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 倍率缓存。 diff --git a/backend/internal/service/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index 410b4d944d..677239c3c6 100644 --- a/backend/internal/service/openai_gateway_usage.go +++ b/backend/internal/service/openai_gateway_usage.go @@ -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 倍率缓存。