mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
01842576b3 | ||
|
|
49b2122dc6 | ||
|
|
8d7d3185e9 | ||
|
|
bba85923b9 | ||
|
|
49dc6e2363 | ||
|
|
9d5e63a598 | ||
|
|
53c21fed25 | ||
|
|
7c5afa4c36 | ||
|
|
b669dc6daf | ||
|
|
c4bdc48475 | ||
|
|
26f70c6a87 | ||
|
|
cfbbfd6b3b | ||
|
|
43a3c197cb | ||
|
|
2023c290e1 | ||
|
|
ab80abc6d1 | ||
|
|
492c630599 | ||
|
|
ba2a717332 | ||
|
|
8b74741efe | ||
|
|
fc82e83102 | ||
|
|
9591136d45 | ||
|
|
26b31a83e4 | ||
|
|
7b03b1f75a | ||
|
|
43768680d3 | ||
|
|
0718581265 | ||
|
|
e468262294 | ||
|
|
5dbb8f7085 | ||
|
|
6d509c632d | ||
|
|
6a6d4728a0 | ||
|
|
832677a5e6 | ||
|
|
c2885c430f | ||
|
|
59fb22791e | ||
|
|
20b0be5b05 | ||
|
|
d70265306b | ||
|
|
f2b1556bd4 | ||
|
|
0ae1b371fa | ||
|
|
da9799542d | ||
|
|
6977bb63bb | ||
|
|
b63dafa426 | ||
|
|
daebb7a80c | ||
|
|
4426593e9d | ||
|
|
f0ee427f0a | ||
|
|
05d7946af8 | ||
|
|
c9e4b267b6 | ||
|
|
6f9b02f12c | ||
|
|
1c68f9996d | ||
|
|
2588675ce7 | ||
|
|
81e6ec5346 | ||
|
|
64331a9792 | ||
|
|
3abd646582 | ||
|
|
d5a802f84a | ||
|
|
3bfd40f5af | ||
|
|
0f189af8de | ||
|
|
ce6242fca1 | ||
|
|
66b7ce08e7 | ||
|
|
29b2d721f4 | ||
|
|
5753552b95 | ||
|
|
8e3d75b57b | ||
|
|
bd2f53b855 | ||
|
|
d01628975d | ||
|
|
2bfc49477c | ||
|
|
b686842716 | ||
|
|
9a51333101 | ||
|
|
3a6fd1c9db | ||
|
|
d56f5e7879 | ||
|
|
ba280b3d5e | ||
|
|
df3c38e10f | ||
|
|
07d51a6aef | ||
|
|
daa7915749 | ||
|
|
63800c50d4 | ||
|
|
f7f1112cda | ||
|
|
c15d858eca | ||
|
|
a46249aac5 | ||
|
|
5a25507dcb | ||
|
|
983cf297da | ||
|
|
69d6d106fd | ||
|
|
c4224045dc | ||
|
|
27de27ac64 | ||
|
|
a56720855c | ||
|
|
f48866958b |
@@ -5,6 +5,7 @@ FRONTEND_CRITICAL_VITEST := \
|
||||
src/api/__tests__/client.spec.ts \
|
||||
src/api/__tests__/tokenRefresh.spec.ts \
|
||||
src/api/__tests__/keys.bulkUpdate.spec.ts \
|
||||
src/composables/__tests__/useModelWhitelist.spec.ts \
|
||||
src/components/account/__tests__/OpenAIReferralCell.spec.ts \
|
||||
src/components/account/__tests__/OpenAIReferralCell.transport.spec.ts \
|
||||
src/components/account/__tests__/OpenAIQuotaResetCell.spark_shadow.spec.ts \
|
||||
|
||||
@@ -1 +1 @@
|
||||
0.2.14
|
||||
0.2.15
|
||||
|
||||
@@ -104,6 +104,7 @@ func provideCleanup(
|
||||
subscriptionExpiry *service.SubscriptionExpiryService,
|
||||
usageCleanup *service.UsageCleanupService,
|
||||
idempotencyCleanup *service.IdempotencyCleanupService,
|
||||
geminiCachedContent *service.GeminiCachedContentService,
|
||||
batchImageCleanup *service.BatchImageCleanupService,
|
||||
batchImageWorker *service.BatchImageWorkerRuntime,
|
||||
pricing *service.PricingService,
|
||||
@@ -244,6 +245,12 @@ func provideCleanup(
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"GeminiCachedContentService", func() error {
|
||||
if geminiCachedContent != nil {
|
||||
geminiCachedContent.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"BatchImageCleanupService", func() error {
|
||||
if batchImageCleanup != nil {
|
||||
batchImageCleanup.Stop()
|
||||
|
||||
@@ -296,7 +296,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig)
|
||||
legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService)
|
||||
coordinator := securityaudit.NewCoordinator(legacyEngine, promptService)
|
||||
gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator)
|
||||
geminiCachedContentRepository := repository.NewGeminiCachedContentRepository(db)
|
||||
geminiCachedContentService := service.ProvideGeminiCachedContentService(geminiCachedContentRepository, geminiMessagesCompatService, configConfig)
|
||||
gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator, geminiCachedContentService)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator, compositeRouteResolver)
|
||||
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
|
||||
totpHandler := handler.NewTotpHandler(totpService)
|
||||
@@ -353,7 +355,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService, channelMonitorQuotaFetcher)
|
||||
channelMonitorV2Aggregator := service.ProvideChannelMonitorV2Aggregator(channelMonitorV2Repository, db, settingService)
|
||||
userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, claudeCodeVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, claudeCodeVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, geminiCachedContentService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager)
|
||||
application := &Application{
|
||||
Server: httpServer,
|
||||
PromptAudit: promptService,
|
||||
@@ -413,6 +415,7 @@ func provideCleanup(
|
||||
subscriptionExpiry *service.SubscriptionExpiryService,
|
||||
usageCleanup *service.UsageCleanupService,
|
||||
idempotencyCleanup *service.IdempotencyCleanupService,
|
||||
geminiCachedContent *service.GeminiCachedContentService,
|
||||
batchImageCleanup *service.BatchImageCleanupService,
|
||||
batchImageWorker *service.BatchImageWorkerRuntime,
|
||||
pricing *service.PricingService,
|
||||
@@ -552,6 +555,12 @@ func provideCleanup(
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"GeminiCachedContentService", func() error {
|
||||
if geminiCachedContent != nil {
|
||||
geminiCachedContent.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"BatchImageCleanupService", func() error {
|
||||
if batchImageCleanup != nil {
|
||||
batchImageCleanup.Stop()
|
||||
|
||||
@@ -74,6 +74,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
|
||||
subscriptionExpirySvc,
|
||||
&service.UsageCleanupService{},
|
||||
idempotencyCleanupSvc,
|
||||
service.NewGeminiCachedContentService(nil, nil, cfg),
|
||||
&service.BatchImageCleanupService{},
|
||||
nil, // batchImageWorker
|
||||
pricingSvc,
|
||||
|
||||
+14
-1
@@ -30,6 +30,8 @@ type Group struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
// Description holds the value of the "description" field.
|
||||
Description *string `json:"description,omitempty"`
|
||||
// 名称与描述的界面语言译文:语言代码 -> {name, description};缺失的语言沿用 name/description
|
||||
I18n domain.GroupI18n `json:"i18n,omitempty"`
|
||||
// RateMultiplier holds the value of the "rate_multiplier" field.
|
||||
RateMultiplier float64 `json:"rate_multiplier,omitempty"`
|
||||
// 是否启用高峰时段倍率
|
||||
@@ -258,7 +260,7 @@ func (*Group) scanValues(columns []string) ([]any, error) {
|
||||
values := make([]any, len(columns))
|
||||
for i := range columns {
|
||||
switch columns[i] {
|
||||
case group.FieldVideoModelPrices, group.FieldModelPricing, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelAllowlist, group.FieldCodexModelsManifestConfig, group.FieldReasoningEffortMappings:
|
||||
case group.FieldI18n, group.FieldVideoModelPrices, group.FieldModelPricing, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelAllowlist, group.FieldCodexModelsManifestConfig, group.FieldReasoningEffortMappings:
|
||||
values[i] = new([]byte)
|
||||
case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldLongContextPricingEnabled, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldAllowLive, group.FieldForceOpenaiFast, group.FieldFreeOpenaiFast, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet, group.FieldProfitControlEnabled:
|
||||
values[i] = new(sql.NullBool)
|
||||
@@ -323,6 +325,14 @@ func (_m *Group) assignValues(columns []string, values []any) error {
|
||||
_m.Description = new(string)
|
||||
*_m.Description = value.String
|
||||
}
|
||||
case group.FieldI18n:
|
||||
if value, ok := values[i].(*[]byte); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field i18n", values[i])
|
||||
} else if value != nil && len(*value) > 0 {
|
||||
if err := json.Unmarshal(*value, &_m.I18n); err != nil {
|
||||
return fmt.Errorf("unmarshal field i18n: %w", err)
|
||||
}
|
||||
}
|
||||
case group.FieldRateMultiplier:
|
||||
if value, ok := values[i].(*sql.NullFloat64); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field rate_multiplier", values[i])
|
||||
@@ -817,6 +827,9 @@ func (_m *Group) String() string {
|
||||
builder.WriteString(*v)
|
||||
}
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("i18n=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.I18n))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("rate_multiplier=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.RateMultiplier))
|
||||
builder.WriteString(", ")
|
||||
|
||||
@@ -26,6 +26,8 @@ const (
|
||||
FieldName = "name"
|
||||
// FieldDescription holds the string denoting the description field in the database.
|
||||
FieldDescription = "description"
|
||||
// FieldI18n holds the string denoting the i18n field in the database.
|
||||
FieldI18n = "i18n"
|
||||
// FieldRateMultiplier holds the string denoting the rate_multiplier field in the database.
|
||||
FieldRateMultiplier = "rate_multiplier"
|
||||
// FieldPeakRateEnabled holds the string denoting the peak_rate_enabled field in the database.
|
||||
@@ -228,6 +230,7 @@ var Columns = []string{
|
||||
FieldDeletedAt,
|
||||
FieldName,
|
||||
FieldDescription,
|
||||
FieldI18n,
|
||||
FieldRateMultiplier,
|
||||
FieldPeakRateEnabled,
|
||||
FieldPeakStart,
|
||||
@@ -326,6 +329,8 @@ var (
|
||||
UpdateDefaultUpdatedAt func() time.Time
|
||||
// NameValidator is a validator for the "name" field. It is called by the builders before save.
|
||||
NameValidator func(string) error
|
||||
// DefaultI18n holds the default value on creation for the "i18n" field.
|
||||
DefaultI18n domain.GroupI18n
|
||||
// DefaultRateMultiplier holds the default value on creation for the "rate_multiplier" field.
|
||||
DefaultRateMultiplier float64
|
||||
// DefaultPeakRateEnabled holds the default value on creation for the "peak_rate_enabled" field.
|
||||
|
||||
@@ -92,6 +92,12 @@ func (_c *GroupCreate) SetNillableDescription(v *string) *GroupCreate {
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetI18n sets the "i18n" field.
|
||||
func (_c *GroupCreate) SetI18n(v domain.GroupI18n) *GroupCreate {
|
||||
_c.mutation.SetI18n(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetRateMultiplier sets the "rate_multiplier" field.
|
||||
func (_c *GroupCreate) SetRateMultiplier(v float64) *GroupCreate {
|
||||
_c.mutation.SetRateMultiplier(v)
|
||||
@@ -1047,6 +1053,10 @@ func (_c *GroupCreate) defaults() error {
|
||||
v := group.DefaultUpdatedAt()
|
||||
_c.mutation.SetUpdatedAt(v)
|
||||
}
|
||||
if _, ok := _c.mutation.I18n(); !ok {
|
||||
v := group.DefaultI18n
|
||||
_c.mutation.SetI18n(v)
|
||||
}
|
||||
if _, ok := _c.mutation.RateMultiplier(); !ok {
|
||||
v := group.DefaultRateMultiplier
|
||||
_c.mutation.SetRateMultiplier(v)
|
||||
@@ -1230,6 +1240,9 @@ func (_c *GroupCreate) check() error {
|
||||
return &ValidationError{Name: "name", err: fmt.Errorf(`ent: validator failed for field "Group.name": %w`, err)}
|
||||
}
|
||||
}
|
||||
if _, ok := _c.mutation.I18n(); !ok {
|
||||
return &ValidationError{Name: "i18n", err: errors.New(`ent: missing required field "Group.i18n"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.RateMultiplier(); !ok {
|
||||
return &ValidationError{Name: "rate_multiplier", err: errors.New(`ent: missing required field "Group.rate_multiplier"`)}
|
||||
}
|
||||
@@ -1465,6 +1478,10 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) {
|
||||
_spec.SetField(group.FieldDescription, field.TypeString, value)
|
||||
_node.Description = &value
|
||||
}
|
||||
if value, ok := _c.mutation.I18n(); ok {
|
||||
_spec.SetField(group.FieldI18n, field.TypeJSON, value)
|
||||
_node.I18n = value
|
||||
}
|
||||
if value, ok := _c.mutation.RateMultiplier(); ok {
|
||||
_spec.SetField(group.FieldRateMultiplier, field.TypeFloat64, value)
|
||||
_node.RateMultiplier = value
|
||||
@@ -1925,6 +1942,18 @@ func (u *GroupUpsert) ClearDescription() *GroupUpsert {
|
||||
return u
|
||||
}
|
||||
|
||||
// SetI18n sets the "i18n" field.
|
||||
func (u *GroupUpsert) SetI18n(v domain.GroupI18n) *GroupUpsert {
|
||||
u.Set(group.FieldI18n, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateI18n sets the "i18n" field to the value that was provided on create.
|
||||
func (u *GroupUpsert) UpdateI18n() *GroupUpsert {
|
||||
u.SetExcluded(group.FieldI18n)
|
||||
return u
|
||||
}
|
||||
|
||||
// SetRateMultiplier sets the "rate_multiplier" field.
|
||||
func (u *GroupUpsert) SetRateMultiplier(v float64) *GroupUpsert {
|
||||
u.Set(group.FieldRateMultiplier, v)
|
||||
@@ -3039,6 +3068,20 @@ func (u *GroupUpsertOne) ClearDescription() *GroupUpsertOne {
|
||||
})
|
||||
}
|
||||
|
||||
// SetI18n sets the "i18n" field.
|
||||
func (u *GroupUpsertOne) SetI18n(v domain.GroupI18n) *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetI18n(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateI18n sets the "i18n" field to the value that was provided on create.
|
||||
func (u *GroupUpsertOne) UpdateI18n() *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateI18n()
|
||||
})
|
||||
}
|
||||
|
||||
// SetRateMultiplier sets the "rate_multiplier" field.
|
||||
func (u *GroupUpsertOne) SetRateMultiplier(v float64) *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
@@ -4485,6 +4528,20 @@ func (u *GroupUpsertBulk) ClearDescription() *GroupUpsertBulk {
|
||||
})
|
||||
}
|
||||
|
||||
// SetI18n sets the "i18n" field.
|
||||
func (u *GroupUpsertBulk) SetI18n(v domain.GroupI18n) *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetI18n(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateI18n sets the "i18n" field to the value that was provided on create.
|
||||
func (u *GroupUpsertBulk) UpdateI18n() *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateI18n()
|
||||
})
|
||||
}
|
||||
|
||||
// SetRateMultiplier sets the "rate_multiplier" field.
|
||||
func (u *GroupUpsertBulk) SetRateMultiplier(v float64) *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
|
||||
@@ -97,6 +97,12 @@ func (_u *GroupUpdate) ClearDescription() *GroupUpdate {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetI18n sets the "i18n" field.
|
||||
func (_u *GroupUpdate) SetI18n(v domain.GroupI18n) *GroupUpdate {
|
||||
_u.mutation.SetI18n(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetRateMultiplier sets the "rate_multiplier" field.
|
||||
func (_u *GroupUpdate) SetRateMultiplier(v float64) *GroupUpdate {
|
||||
_u.mutation.ResetRateMultiplier()
|
||||
@@ -1581,6 +1587,9 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) {
|
||||
if _u.mutation.DescriptionCleared() {
|
||||
_spec.ClearField(group.FieldDescription, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.I18n(); ok {
|
||||
_spec.SetField(group.FieldI18n, field.TypeJSON, value)
|
||||
}
|
||||
if value, ok := _u.mutation.RateMultiplier(); ok {
|
||||
_spec.SetField(group.FieldRateMultiplier, field.TypeFloat64, value)
|
||||
}
|
||||
@@ -2291,6 +2300,12 @@ func (_u *GroupUpdateOne) ClearDescription() *GroupUpdateOne {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetI18n sets the "i18n" field.
|
||||
func (_u *GroupUpdateOne) SetI18n(v domain.GroupI18n) *GroupUpdateOne {
|
||||
_u.mutation.SetI18n(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetRateMultiplier sets the "rate_multiplier" field.
|
||||
func (_u *GroupUpdateOne) SetRateMultiplier(v float64) *GroupUpdateOne {
|
||||
_u.mutation.ResetRateMultiplier()
|
||||
@@ -3805,6 +3820,9 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error)
|
||||
if _u.mutation.DescriptionCleared() {
|
||||
_spec.ClearField(group.FieldDescription, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.I18n(); ok {
|
||||
_spec.SetField(group.FieldI18n, field.TypeJSON, value)
|
||||
}
|
||||
if value, ok := _u.mutation.RateMultiplier(); ok {
|
||||
_spec.SetField(group.FieldRateMultiplier, field.TypeFloat64, value)
|
||||
}
|
||||
|
||||
@@ -908,6 +908,7 @@ var (
|
||||
{Name: "deleted_at", Type: field.TypeTime, Nullable: true, SchemaType: map[string]string{"postgres": "timestamptz"}},
|
||||
{Name: "name", Type: field.TypeString, Size: 100},
|
||||
{Name: "description", Type: field.TypeString, Nullable: true, SchemaType: map[string]string{"postgres": "text"}},
|
||||
{Name: "i18n", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
|
||||
{Name: "rate_multiplier", Type: field.TypeFloat64, Default: 1, SchemaType: map[string]string{"postgres": "decimal(10,4)"}},
|
||||
{Name: "peak_rate_enabled", Type: field.TypeBool, Default: false},
|
||||
{Name: "peak_start", Type: field.TypeString, Size: 5, Default: ""},
|
||||
@@ -979,22 +980,22 @@ var (
|
||||
{
|
||||
Name: "group_status",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{GroupsColumns[12]},
|
||||
Columns: []*schema.Column{GroupsColumns[13]},
|
||||
},
|
||||
{
|
||||
Name: "group_platform",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{GroupsColumns[14]},
|
||||
Columns: []*schema.Column{GroupsColumns[15]},
|
||||
},
|
||||
{
|
||||
Name: "group_subscription_type",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{GroupsColumns[15]},
|
||||
Columns: []*schema.Column{GroupsColumns[16]},
|
||||
},
|
||||
{
|
||||
Name: "group_is_exclusive",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{GroupsColumns[11]},
|
||||
Columns: []*schema.Column{GroupsColumns[12]},
|
||||
},
|
||||
{
|
||||
Name: "group_deleted_at",
|
||||
@@ -1004,12 +1005,12 @@ var (
|
||||
{
|
||||
Name: "group_sort_order",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{GroupsColumns[49]},
|
||||
Columns: []*schema.Column{GroupsColumns[50]},
|
||||
},
|
||||
{
|
||||
Name: "idx_groups_duplicate_operation_id_active",
|
||||
Unique: true,
|
||||
Columns: []*schema.Column{GroupsColumns[13]},
|
||||
Columns: []*schema.Column{GroupsColumns[14]},
|
||||
Annotation: &entsql.IndexAnnotation{
|
||||
Where: "duplicate_operation_id IS NOT NULL AND deleted_at IS NULL",
|
||||
},
|
||||
|
||||
+55
-1
@@ -22086,6 +22086,7 @@ type GroupMutation struct {
|
||||
deleted_at *time.Time
|
||||
name *string
|
||||
description *string
|
||||
i18n *domain.GroupI18n
|
||||
rate_multiplier *float64
|
||||
addrate_multiplier *float64
|
||||
peak_rate_enabled *bool
|
||||
@@ -22505,6 +22506,42 @@ func (m *GroupMutation) ResetDescription() {
|
||||
delete(m.clearedFields, group.FieldDescription)
|
||||
}
|
||||
|
||||
// SetI18n sets the "i18n" field.
|
||||
func (m *GroupMutation) SetI18n(di domain.GroupI18n) {
|
||||
m.i18n = &di
|
||||
}
|
||||
|
||||
// I18n returns the value of the "i18n" field in the mutation.
|
||||
func (m *GroupMutation) I18n() (r domain.GroupI18n, exists bool) {
|
||||
v := m.i18n
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldI18n returns the old "i18n" field's value of the Group entity.
|
||||
// If the Group object wasn't provided to the builder, the object is fetched from the database.
|
||||
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
|
||||
func (m *GroupMutation) OldI18n(ctx context.Context) (v domain.GroupI18n, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldI18n is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldI18n requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldI18n: %w", err)
|
||||
}
|
||||
return oldValue.I18n, nil
|
||||
}
|
||||
|
||||
// ResetI18n resets all changes to the "i18n" field.
|
||||
func (m *GroupMutation) ResetI18n() {
|
||||
m.i18n = nil
|
||||
}
|
||||
|
||||
// SetRateMultiplier sets the "rate_multiplier" field.
|
||||
func (m *GroupMutation) SetRateMultiplier(f float64) {
|
||||
m.rate_multiplier = &f
|
||||
@@ -25921,7 +25958,7 @@ func (m *GroupMutation) Type() string {
|
||||
// order to get all numeric fields that were incremented/decremented, call
|
||||
// AddedFields().
|
||||
func (m *GroupMutation) Fields() []string {
|
||||
fields := make([]string, 0, 66)
|
||||
fields := make([]string, 0, 67)
|
||||
if m.created_at != nil {
|
||||
fields = append(fields, group.FieldCreatedAt)
|
||||
}
|
||||
@@ -25937,6 +25974,9 @@ func (m *GroupMutation) Fields() []string {
|
||||
if m.description != nil {
|
||||
fields = append(fields, group.FieldDescription)
|
||||
}
|
||||
if m.i18n != nil {
|
||||
fields = append(fields, group.FieldI18n)
|
||||
}
|
||||
if m.rate_multiplier != nil {
|
||||
fields = append(fields, group.FieldRateMultiplier)
|
||||
}
|
||||
@@ -26138,6 +26178,8 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) {
|
||||
return m.Name()
|
||||
case group.FieldDescription:
|
||||
return m.Description()
|
||||
case group.FieldI18n:
|
||||
return m.I18n()
|
||||
case group.FieldRateMultiplier:
|
||||
return m.RateMultiplier()
|
||||
case group.FieldPeakRateEnabled:
|
||||
@@ -26279,6 +26321,8 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e
|
||||
return m.OldName(ctx)
|
||||
case group.FieldDescription:
|
||||
return m.OldDescription(ctx)
|
||||
case group.FieldI18n:
|
||||
return m.OldI18n(ctx)
|
||||
case group.FieldRateMultiplier:
|
||||
return m.OldRateMultiplier(ctx)
|
||||
case group.FieldPeakRateEnabled:
|
||||
@@ -26445,6 +26489,13 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error {
|
||||
}
|
||||
m.SetDescription(v)
|
||||
return nil
|
||||
case group.FieldI18n:
|
||||
v, ok := value.(domain.GroupI18n)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetI18n(v)
|
||||
return nil
|
||||
case group.FieldRateMultiplier:
|
||||
v, ok := value.(float64)
|
||||
if !ok {
|
||||
@@ -27398,6 +27449,9 @@ func (m *GroupMutation) ResetField(name string) error {
|
||||
case group.FieldDescription:
|
||||
m.ResetDescription()
|
||||
return nil
|
||||
case group.FieldI18n:
|
||||
m.ResetI18n()
|
||||
return nil
|
||||
case group.FieldRateMultiplier:
|
||||
m.ResetRateMultiplier()
|
||||
return nil
|
||||
|
||||
@@ -1037,204 +1037,208 @@ func init() {
|
||||
return nil
|
||||
}
|
||||
}()
|
||||
// groupDescI18n is the schema descriptor for i18n field.
|
||||
groupDescI18n := groupFields[2].Descriptor()
|
||||
// group.DefaultI18n holds the default value on creation for the i18n field.
|
||||
group.DefaultI18n = groupDescI18n.Default.(domain.GroupI18n)
|
||||
// groupDescRateMultiplier is the schema descriptor for rate_multiplier field.
|
||||
groupDescRateMultiplier := groupFields[2].Descriptor()
|
||||
groupDescRateMultiplier := groupFields[3].Descriptor()
|
||||
// group.DefaultRateMultiplier holds the default value on creation for the rate_multiplier field.
|
||||
group.DefaultRateMultiplier = groupDescRateMultiplier.Default.(float64)
|
||||
// groupDescPeakRateEnabled is the schema descriptor for peak_rate_enabled field.
|
||||
groupDescPeakRateEnabled := groupFields[3].Descriptor()
|
||||
groupDescPeakRateEnabled := groupFields[4].Descriptor()
|
||||
// group.DefaultPeakRateEnabled holds the default value on creation for the peak_rate_enabled field.
|
||||
group.DefaultPeakRateEnabled = groupDescPeakRateEnabled.Default.(bool)
|
||||
// groupDescPeakStart is the schema descriptor for peak_start field.
|
||||
groupDescPeakStart := groupFields[4].Descriptor()
|
||||
groupDescPeakStart := groupFields[5].Descriptor()
|
||||
// group.DefaultPeakStart holds the default value on creation for the peak_start field.
|
||||
group.DefaultPeakStart = groupDescPeakStart.Default.(string)
|
||||
// group.PeakStartValidator is a validator for the "peak_start" field. It is called by the builders before save.
|
||||
group.PeakStartValidator = groupDescPeakStart.Validators[0].(func(string) error)
|
||||
// groupDescPeakEnd is the schema descriptor for peak_end field.
|
||||
groupDescPeakEnd := groupFields[5].Descriptor()
|
||||
groupDescPeakEnd := groupFields[6].Descriptor()
|
||||
// group.DefaultPeakEnd holds the default value on creation for the peak_end field.
|
||||
group.DefaultPeakEnd = groupDescPeakEnd.Default.(string)
|
||||
// group.PeakEndValidator is a validator for the "peak_end" field. It is called by the builders before save.
|
||||
group.PeakEndValidator = groupDescPeakEnd.Validators[0].(func(string) error)
|
||||
// groupDescPeakRateMultiplier is the schema descriptor for peak_rate_multiplier field.
|
||||
groupDescPeakRateMultiplier := groupFields[6].Descriptor()
|
||||
groupDescPeakRateMultiplier := groupFields[7].Descriptor()
|
||||
// group.DefaultPeakRateMultiplier holds the default value on creation for the peak_rate_multiplier field.
|
||||
group.DefaultPeakRateMultiplier = groupDescPeakRateMultiplier.Default.(float64)
|
||||
// groupDescIsExclusive is the schema descriptor for is_exclusive field.
|
||||
groupDescIsExclusive := groupFields[7].Descriptor()
|
||||
groupDescIsExclusive := groupFields[8].Descriptor()
|
||||
// group.DefaultIsExclusive holds the default value on creation for the is_exclusive field.
|
||||
group.DefaultIsExclusive = groupDescIsExclusive.Default.(bool)
|
||||
// groupDescStatus is the schema descriptor for status field.
|
||||
groupDescStatus := groupFields[8].Descriptor()
|
||||
groupDescStatus := groupFields[9].Descriptor()
|
||||
// group.DefaultStatus holds the default value on creation for the status field.
|
||||
group.DefaultStatus = groupDescStatus.Default.(string)
|
||||
// group.StatusValidator is a validator for the "status" field. It is called by the builders before save.
|
||||
group.StatusValidator = groupDescStatus.Validators[0].(func(string) error)
|
||||
// groupDescDuplicateOperationID is the schema descriptor for duplicate_operation_id field.
|
||||
groupDescDuplicateOperationID := groupFields[9].Descriptor()
|
||||
groupDescDuplicateOperationID := groupFields[10].Descriptor()
|
||||
// group.DuplicateOperationIDValidator is a validator for the "duplicate_operation_id" field. It is called by the builders before save.
|
||||
group.DuplicateOperationIDValidator = groupDescDuplicateOperationID.Validators[0].(func(string) error)
|
||||
// groupDescPlatform is the schema descriptor for platform field.
|
||||
groupDescPlatform := groupFields[10].Descriptor()
|
||||
groupDescPlatform := groupFields[11].Descriptor()
|
||||
// group.DefaultPlatform holds the default value on creation for the platform field.
|
||||
group.DefaultPlatform = groupDescPlatform.Default.(string)
|
||||
// group.PlatformValidator is a validator for the "platform" field. It is called by the builders before save.
|
||||
group.PlatformValidator = groupDescPlatform.Validators[0].(func(string) error)
|
||||
// groupDescSubscriptionType is the schema descriptor for subscription_type field.
|
||||
groupDescSubscriptionType := groupFields[11].Descriptor()
|
||||
groupDescSubscriptionType := groupFields[12].Descriptor()
|
||||
// group.DefaultSubscriptionType holds the default value on creation for the subscription_type field.
|
||||
group.DefaultSubscriptionType = groupDescSubscriptionType.Default.(string)
|
||||
// group.SubscriptionTypeValidator is a validator for the "subscription_type" field. It is called by the builders before save.
|
||||
group.SubscriptionTypeValidator = groupDescSubscriptionType.Validators[0].(func(string) error)
|
||||
// groupDescDefaultValidityDays is the schema descriptor for default_validity_days field.
|
||||
groupDescDefaultValidityDays := groupFields[15].Descriptor()
|
||||
groupDescDefaultValidityDays := groupFields[16].Descriptor()
|
||||
// group.DefaultDefaultValidityDays holds the default value on creation for the default_validity_days field.
|
||||
group.DefaultDefaultValidityDays = groupDescDefaultValidityDays.Default.(int)
|
||||
// groupDescAllowImageGeneration is the schema descriptor for allow_image_generation field.
|
||||
groupDescAllowImageGeneration := groupFields[16].Descriptor()
|
||||
groupDescAllowImageGeneration := groupFields[17].Descriptor()
|
||||
// group.DefaultAllowImageGeneration holds the default value on creation for the allow_image_generation field.
|
||||
group.DefaultAllowImageGeneration = groupDescAllowImageGeneration.Default.(bool)
|
||||
// groupDescAllowBatchImageGeneration is the schema descriptor for allow_batch_image_generation field.
|
||||
groupDescAllowBatchImageGeneration := groupFields[17].Descriptor()
|
||||
groupDescAllowBatchImageGeneration := groupFields[18].Descriptor()
|
||||
// group.DefaultAllowBatchImageGeneration holds the default value on creation for the allow_batch_image_generation field.
|
||||
group.DefaultAllowBatchImageGeneration = groupDescAllowBatchImageGeneration.Default.(bool)
|
||||
// groupDescImageRateIndependent is the schema descriptor for image_rate_independent field.
|
||||
groupDescImageRateIndependent := groupFields[18].Descriptor()
|
||||
groupDescImageRateIndependent := groupFields[19].Descriptor()
|
||||
// group.DefaultImageRateIndependent holds the default value on creation for the image_rate_independent field.
|
||||
group.DefaultImageRateIndependent = groupDescImageRateIndependent.Default.(bool)
|
||||
// groupDescImageRateMultiplier is the schema descriptor for image_rate_multiplier field.
|
||||
groupDescImageRateMultiplier := groupFields[19].Descriptor()
|
||||
groupDescImageRateMultiplier := groupFields[20].Descriptor()
|
||||
// group.DefaultImageRateMultiplier holds the default value on creation for the image_rate_multiplier field.
|
||||
group.DefaultImageRateMultiplier = groupDescImageRateMultiplier.Default.(float64)
|
||||
// groupDescBatchImageDiscountMultiplier is the schema descriptor for batch_image_discount_multiplier field.
|
||||
groupDescBatchImageDiscountMultiplier := groupFields[23].Descriptor()
|
||||
groupDescBatchImageDiscountMultiplier := groupFields[24].Descriptor()
|
||||
// group.DefaultBatchImageDiscountMultiplier holds the default value on creation for the batch_image_discount_multiplier field.
|
||||
group.DefaultBatchImageDiscountMultiplier = groupDescBatchImageDiscountMultiplier.Default.(float64)
|
||||
// groupDescBatchImageHoldMultiplier is the schema descriptor for batch_image_hold_multiplier field.
|
||||
groupDescBatchImageHoldMultiplier := groupFields[24].Descriptor()
|
||||
groupDescBatchImageHoldMultiplier := groupFields[25].Descriptor()
|
||||
// group.DefaultBatchImageHoldMultiplier holds the default value on creation for the batch_image_hold_multiplier field.
|
||||
group.DefaultBatchImageHoldMultiplier = groupDescBatchImageHoldMultiplier.Default.(float64)
|
||||
// groupDescVideoRateIndependent is the schema descriptor for video_rate_independent field.
|
||||
groupDescVideoRateIndependent := groupFields[25].Descriptor()
|
||||
groupDescVideoRateIndependent := groupFields[26].Descriptor()
|
||||
// group.DefaultVideoRateIndependent holds the default value on creation for the video_rate_independent field.
|
||||
group.DefaultVideoRateIndependent = groupDescVideoRateIndependent.Default.(bool)
|
||||
// groupDescVideoRateMultiplier is the schema descriptor for video_rate_multiplier field.
|
||||
groupDescVideoRateMultiplier := groupFields[26].Descriptor()
|
||||
groupDescVideoRateMultiplier := groupFields[27].Descriptor()
|
||||
// group.DefaultVideoRateMultiplier holds the default value on creation for the video_rate_multiplier field.
|
||||
group.DefaultVideoRateMultiplier = groupDescVideoRateMultiplier.Default.(float64)
|
||||
// groupDescSearchPricePer1k is the schema descriptor for search_price_per_1k field.
|
||||
groupDescSearchPricePer1k := groupFields[32].Descriptor()
|
||||
groupDescSearchPricePer1k := groupFields[33].Descriptor()
|
||||
// group.SearchPricePer1kValidator is a validator for the "search_price_per_1k" field. It is called by the builders before save.
|
||||
group.SearchPricePer1kValidator = groupDescSearchPricePer1k.Validators[0].(func(float64) error)
|
||||
// groupDescAudioRealtimePricePerMin is the schema descriptor for audio_realtime_price_per_min field.
|
||||
groupDescAudioRealtimePricePerMin := groupFields[33].Descriptor()
|
||||
groupDescAudioRealtimePricePerMin := groupFields[34].Descriptor()
|
||||
// group.AudioRealtimePricePerMinValidator is a validator for the "audio_realtime_price_per_min" field. It is called by the builders before save.
|
||||
group.AudioRealtimePricePerMinValidator = groupDescAudioRealtimePricePerMin.Validators[0].(func(float64) error)
|
||||
// groupDescAudioTtsPricePerMillionChars is the schema descriptor for audio_tts_price_per_million_chars field.
|
||||
groupDescAudioTtsPricePerMillionChars := groupFields[34].Descriptor()
|
||||
groupDescAudioTtsPricePerMillionChars := groupFields[35].Descriptor()
|
||||
// group.AudioTtsPricePerMillionCharsValidator is a validator for the "audio_tts_price_per_million_chars" field. It is called by the builders before save.
|
||||
group.AudioTtsPricePerMillionCharsValidator = groupDescAudioTtsPricePerMillionChars.Validators[0].(func(float64) error)
|
||||
// groupDescAudioSttPricePerHour is the schema descriptor for audio_stt_price_per_hour field.
|
||||
groupDescAudioSttPricePerHour := groupFields[35].Descriptor()
|
||||
groupDescAudioSttPricePerHour := groupFields[36].Descriptor()
|
||||
// group.AudioSttPricePerHourValidator is a validator for the "audio_stt_price_per_hour" field. It is called by the builders before save.
|
||||
group.AudioSttPricePerHourValidator = groupDescAudioSttPricePerHour.Validators[0].(func(float64) error)
|
||||
// groupDescLongContextPricingEnabled is the schema descriptor for long_context_pricing_enabled field.
|
||||
groupDescLongContextPricingEnabled := groupFields[36].Descriptor()
|
||||
groupDescLongContextPricingEnabled := groupFields[37].Descriptor()
|
||||
// group.DefaultLongContextPricingEnabled holds the default value on creation for the long_context_pricing_enabled field.
|
||||
group.DefaultLongContextPricingEnabled = groupDescLongContextPricingEnabled.Default.(bool)
|
||||
// groupDescClaudeCodeOnly is the schema descriptor for claude_code_only field.
|
||||
groupDescClaudeCodeOnly := groupFields[38].Descriptor()
|
||||
groupDescClaudeCodeOnly := groupFields[39].Descriptor()
|
||||
// group.DefaultClaudeCodeOnly holds the default value on creation for the claude_code_only field.
|
||||
group.DefaultClaudeCodeOnly = groupDescClaudeCodeOnly.Default.(bool)
|
||||
// groupDescModelRoutingEnabled is the schema descriptor for model_routing_enabled field.
|
||||
groupDescModelRoutingEnabled := groupFields[42].Descriptor()
|
||||
groupDescModelRoutingEnabled := groupFields[43].Descriptor()
|
||||
// group.DefaultModelRoutingEnabled holds the default value on creation for the model_routing_enabled field.
|
||||
group.DefaultModelRoutingEnabled = groupDescModelRoutingEnabled.Default.(bool)
|
||||
// groupDescMcpXMLInject is the schema descriptor for mcp_xml_inject field.
|
||||
groupDescMcpXMLInject := groupFields[43].Descriptor()
|
||||
groupDescMcpXMLInject := groupFields[44].Descriptor()
|
||||
// group.DefaultMcpXMLInject holds the default value on creation for the mcp_xml_inject field.
|
||||
group.DefaultMcpXMLInject = groupDescMcpXMLInject.Default.(bool)
|
||||
// groupDescSupportedModelScopes is the schema descriptor for supported_model_scopes field.
|
||||
groupDescSupportedModelScopes := groupFields[44].Descriptor()
|
||||
groupDescSupportedModelScopes := groupFields[45].Descriptor()
|
||||
// group.DefaultSupportedModelScopes holds the default value on creation for the supported_model_scopes field.
|
||||
group.DefaultSupportedModelScopes = groupDescSupportedModelScopes.Default.([]string)
|
||||
// groupDescSortOrder is the schema descriptor for sort_order field.
|
||||
groupDescSortOrder := groupFields[45].Descriptor()
|
||||
groupDescSortOrder := groupFields[46].Descriptor()
|
||||
// group.DefaultSortOrder holds the default value on creation for the sort_order field.
|
||||
group.DefaultSortOrder = groupDescSortOrder.Default.(int)
|
||||
// groupDescAllowMessagesDispatch is the schema descriptor for allow_messages_dispatch field.
|
||||
groupDescAllowMessagesDispatch := groupFields[46].Descriptor()
|
||||
groupDescAllowMessagesDispatch := groupFields[47].Descriptor()
|
||||
// group.DefaultAllowMessagesDispatch holds the default value on creation for the allow_messages_dispatch field.
|
||||
group.DefaultAllowMessagesDispatch = groupDescAllowMessagesDispatch.Default.(bool)
|
||||
// groupDescAllowLive is the schema descriptor for allow_live field.
|
||||
groupDescAllowLive := groupFields[47].Descriptor()
|
||||
groupDescAllowLive := groupFields[48].Descriptor()
|
||||
// group.DefaultAllowLive holds the default value on creation for the allow_live field.
|
||||
group.DefaultAllowLive = groupDescAllowLive.Default.(bool)
|
||||
// groupDescForceOpenaiFast is the schema descriptor for force_openai_fast field.
|
||||
groupDescForceOpenaiFast := groupFields[48].Descriptor()
|
||||
groupDescForceOpenaiFast := groupFields[49].Descriptor()
|
||||
// group.DefaultForceOpenaiFast holds the default value on creation for the force_openai_fast field.
|
||||
group.DefaultForceOpenaiFast = groupDescForceOpenaiFast.Default.(bool)
|
||||
// groupDescFreeOpenaiFast is the schema descriptor for free_openai_fast field.
|
||||
groupDescFreeOpenaiFast := groupFields[49].Descriptor()
|
||||
groupDescFreeOpenaiFast := groupFields[50].Descriptor()
|
||||
// group.DefaultFreeOpenaiFast holds the default value on creation for the free_openai_fast field.
|
||||
group.DefaultFreeOpenaiFast = groupDescFreeOpenaiFast.Default.(bool)
|
||||
// groupDescRequireOauthOnly is the schema descriptor for require_oauth_only field.
|
||||
groupDescRequireOauthOnly := groupFields[50].Descriptor()
|
||||
groupDescRequireOauthOnly := groupFields[51].Descriptor()
|
||||
// group.DefaultRequireOauthOnly holds the default value on creation for the require_oauth_only field.
|
||||
group.DefaultRequireOauthOnly = groupDescRequireOauthOnly.Default.(bool)
|
||||
// groupDescRequirePrivacySet is the schema descriptor for require_privacy_set field.
|
||||
groupDescRequirePrivacySet := groupFields[51].Descriptor()
|
||||
groupDescRequirePrivacySet := groupFields[52].Descriptor()
|
||||
// group.DefaultRequirePrivacySet holds the default value on creation for the require_privacy_set field.
|
||||
group.DefaultRequirePrivacySet = groupDescRequirePrivacySet.Default.(bool)
|
||||
// groupDescDefaultMappedModel is the schema descriptor for default_mapped_model field.
|
||||
groupDescDefaultMappedModel := groupFields[52].Descriptor()
|
||||
groupDescDefaultMappedModel := groupFields[53].Descriptor()
|
||||
// group.DefaultDefaultMappedModel holds the default value on creation for the default_mapped_model field.
|
||||
group.DefaultDefaultMappedModel = groupDescDefaultMappedModel.Default.(string)
|
||||
// group.DefaultMappedModelValidator is a validator for the "default_mapped_model" field. It is called by the builders before save.
|
||||
group.DefaultMappedModelValidator = groupDescDefaultMappedModel.Validators[0].(func(string) error)
|
||||
// groupDescMessagesDispatchModelConfig is the schema descriptor for messages_dispatch_model_config field.
|
||||
groupDescMessagesDispatchModelConfig := groupFields[53].Descriptor()
|
||||
groupDescMessagesDispatchModelConfig := groupFields[54].Descriptor()
|
||||
// group.DefaultMessagesDispatchModelConfig holds the default value on creation for the messages_dispatch_model_config field.
|
||||
group.DefaultMessagesDispatchModelConfig = groupDescMessagesDispatchModelConfig.Default.(domain.OpenAIMessagesDispatchModelConfig)
|
||||
// groupDescModelAllowlist is the schema descriptor for model_allowlist field.
|
||||
groupDescModelAllowlist := groupFields[54].Descriptor()
|
||||
groupDescModelAllowlist := groupFields[55].Descriptor()
|
||||
// group.DefaultModelAllowlist holds the default value on creation for the model_allowlist field.
|
||||
group.DefaultModelAllowlist = groupDescModelAllowlist.Default.(domain.GroupModelAllowlist)
|
||||
// groupDescCodexModelsManifestConfig is the schema descriptor for codex_models_manifest_config field.
|
||||
groupDescCodexModelsManifestConfig := groupFields[55].Descriptor()
|
||||
groupDescCodexModelsManifestConfig := groupFields[56].Descriptor()
|
||||
// group.DefaultCodexModelsManifestConfig holds the default value on creation for the codex_models_manifest_config field.
|
||||
group.DefaultCodexModelsManifestConfig = groupDescCodexModelsManifestConfig.Default.(domain.GroupCodexModelsManifestConfig)
|
||||
// groupDescRpmLimit is the schema descriptor for rpm_limit field.
|
||||
groupDescRpmLimit := groupFields[56].Descriptor()
|
||||
groupDescRpmLimit := groupFields[57].Descriptor()
|
||||
// group.DefaultRpmLimit holds the default value on creation for the rpm_limit field.
|
||||
group.DefaultRpmLimit = groupDescRpmLimit.Default.(int)
|
||||
// groupDescMaxReasoningEffort is the schema descriptor for max_reasoning_effort field.
|
||||
groupDescMaxReasoningEffort := groupFields[57].Descriptor()
|
||||
groupDescMaxReasoningEffort := groupFields[58].Descriptor()
|
||||
// group.DefaultMaxReasoningEffort holds the default value on creation for the max_reasoning_effort field.
|
||||
group.DefaultMaxReasoningEffort = groupDescMaxReasoningEffort.Default.(string)
|
||||
// group.MaxReasoningEffortValidator is a validator for the "max_reasoning_effort" field. It is called by the builders before save.
|
||||
group.MaxReasoningEffortValidator = groupDescMaxReasoningEffort.Validators[0].(func(string) error)
|
||||
// groupDescMaxReasoningEffortOverLimit is the schema descriptor for max_reasoning_effort_over_limit field.
|
||||
groupDescMaxReasoningEffortOverLimit := groupFields[58].Descriptor()
|
||||
groupDescMaxReasoningEffortOverLimit := groupFields[59].Descriptor()
|
||||
// group.DefaultMaxReasoningEffortOverLimit holds the default value on creation for the max_reasoning_effort_over_limit field.
|
||||
group.DefaultMaxReasoningEffortOverLimit = groupDescMaxReasoningEffortOverLimit.Default.(string)
|
||||
// group.MaxReasoningEffortOverLimitValidator is a validator for the "max_reasoning_effort_over_limit" field. It is called by the builders before save.
|
||||
group.MaxReasoningEffortOverLimitValidator = groupDescMaxReasoningEffortOverLimit.Validators[0].(func(string) error)
|
||||
// groupDescReasoningEffortMappings is the schema descriptor for reasoning_effort_mappings field.
|
||||
groupDescReasoningEffortMappings := groupFields[59].Descriptor()
|
||||
groupDescReasoningEffortMappings := groupFields[60].Descriptor()
|
||||
// group.DefaultReasoningEffortMappings holds the default value on creation for the reasoning_effort_mappings field.
|
||||
group.DefaultReasoningEffortMappings = groupDescReasoningEffortMappings.Default.([]domain.ReasoningEffortMapping)
|
||||
// groupDescProfitControlEnabled is the schema descriptor for profit_control_enabled field.
|
||||
groupDescProfitControlEnabled := groupFields[60].Descriptor()
|
||||
groupDescProfitControlEnabled := groupFields[61].Descriptor()
|
||||
// group.DefaultProfitControlEnabled holds the default value on creation for the profit_control_enabled field.
|
||||
group.DefaultProfitControlEnabled = groupDescProfitControlEnabled.Default.(bool)
|
||||
// groupDescProfitMinMargin is the schema descriptor for profit_min_margin field.
|
||||
groupDescProfitMinMargin := groupFields[61].Descriptor()
|
||||
groupDescProfitMinMargin := groupFields[62].Descriptor()
|
||||
// group.DefaultProfitMinMargin holds the default value on creation for the profit_min_margin field.
|
||||
group.DefaultProfitMinMargin = groupDescProfitMinMargin.Default.(float64)
|
||||
// groupDescProfitSafetyBuffer is the schema descriptor for profit_safety_buffer field.
|
||||
groupDescProfitSafetyBuffer := groupFields[62].Descriptor()
|
||||
groupDescProfitSafetyBuffer := groupFields[63].Descriptor()
|
||||
// group.DefaultProfitSafetyBuffer holds the default value on creation for the profit_safety_buffer field.
|
||||
group.DefaultProfitSafetyBuffer = groupDescProfitSafetyBuffer.Default.(float64)
|
||||
idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin()
|
||||
|
||||
@@ -44,6 +44,11 @@ func (Group) Fields() []ent.Field {
|
||||
Optional().
|
||||
Nillable().
|
||||
SchemaType(map[string]string{dialect.Postgres: "text"}),
|
||||
// 名称与描述的界面语言译文 (added by migration 242)
|
||||
field.JSON("i18n", domain.GroupI18n{}).
|
||||
Default(domain.GroupI18n{}).
|
||||
SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
|
||||
Comment("名称与描述的界面语言译文:语言代码 -> {name, description};缺失的语言沿用 name/description"),
|
||||
field.Float("rate_multiplier").
|
||||
SchemaType(map[string]string{dialect.Postgres: "decimal(10,4)"}).
|
||||
Default(1.0),
|
||||
|
||||
@@ -1005,6 +1005,8 @@ type GatewayConfig struct {
|
||||
ProxyProbeResponseReadMaxBytes int64 `mapstructure:"proxy_probe_response_read_max_bytes"`
|
||||
// Gemini 上游响应头调试日志开关(默认关闭,避免高频日志开销)
|
||||
GeminiDebugResponseHeaders bool `mapstructure:"gemini_debug_response_headers"`
|
||||
// GeminiCachedContentMaxTTLSeconds: Gemini 显式缓存单次可设置的最长有效期(秒),<=0 时按 86400 处理
|
||||
GeminiCachedContentMaxTTLSeconds int `mapstructure:"gemini_cached_content_max_ttl_seconds"`
|
||||
// ConnectionPoolIsolation: 上游连接池隔离策略(proxy/account/account_proxy)
|
||||
ConnectionPoolIsolation string `mapstructure:"connection_pool_isolation"`
|
||||
// ForceCodexCLI: 强制将 OpenAI `/v1/responses` 请求按 Codex CLI 处理。
|
||||
@@ -2090,6 +2092,7 @@ func setDefaults() {
|
||||
"api.openai.com",
|
||||
"api.anthropic.com",
|
||||
"api.kimi.com",
|
||||
"api.kimi.ai",
|
||||
"api.moonshot.ai",
|
||||
"api.moonshot.cn",
|
||||
"open.bigmodel.cn",
|
||||
@@ -2527,6 +2530,7 @@ func setDefaults() {
|
||||
viper.SetDefault("gateway.models_list_read_max_bytes", DefaultModelsListReadMaxBytes)
|
||||
viper.SetDefault("gateway.proxy_probe_response_read_max_bytes", int64(1024*1024))
|
||||
viper.SetDefault("gateway.gemini_debug_response_headers", false)
|
||||
viper.SetDefault("gateway.gemini_cached_content_max_ttl_seconds", 86400)
|
||||
viper.SetDefault("gateway.connection_pool_isolation", ConnectionPoolIsolationAccountProxy)
|
||||
// HTTP 上游连接池配置(针对 5000+ 并发用户优化)
|
||||
viper.SetDefault("gateway.max_idle_conns", 2560) // 最大空闲连接总数(高并发场景可调大)
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
|
||||
"github.com/spf13/viper"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
@@ -831,6 +832,7 @@ func TestLoadDefaultSecurityToggles(t *testing.T) {
|
||||
|
||||
wantHosts := []string{
|
||||
"api.kimi.com",
|
||||
"api.kimi.ai",
|
||||
"api.moonshot.ai",
|
||||
"api.moonshot.cn",
|
||||
}
|
||||
@@ -843,6 +845,19 @@ func TestLoadDefaultSecurityToggles(t *testing.T) {
|
||||
t.Fatalf("URLAllowlist.UpstreamHosts missing %q; got %v", want, cfg.Security.URLAllowlist.UpstreamHosts)
|
||||
}
|
||||
}
|
||||
opts := urlvalidator.ValidationOptions{
|
||||
AllowedHosts: cfg.Security.URLAllowlist.UpstreamHosts,
|
||||
RequireAllowlist: true,
|
||||
AllowPrivate: cfg.Security.URLAllowlist.AllowPrivateHosts,
|
||||
}
|
||||
for _, host := range []string{"api.kimi.com", "api.kimi.ai"} {
|
||||
target := "https://" + host + "/coding/v1/usages"
|
||||
normalized, err := urlvalidator.ValidateHTTPSURL(target, opts)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, target, normalized)
|
||||
_, err = urlvalidator.ValidateHTTPSURL("https://"+host+".attacker.example/coding/v1/usages", opts)
|
||||
require.Error(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadDefaultServerMode(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
package domain
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// GroupLocaleText 是分组名称与描述在某个界面语言下的译文。
|
||||
// 字段为空表示该语言沿用分组自身的 name/description。
|
||||
type GroupLocaleText struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
}
|
||||
|
||||
// GroupI18n 以界面语言代码(如 "en"、"zh")为键保存分组展示文案的译文。
|
||||
// 仅用于页面展示;分组的唯一性、搜索与各类日志仍以 name 为准。
|
||||
type GroupI18n map[string]GroupLocaleText
|
||||
|
||||
// Localize 返回 name、description 在请求首选语言下的版本;该语言没有译文的字段保持原值。
|
||||
// acceptLanguage 是 Accept-Language 请求头。只采用其中优先级最高的语言:原字段的语言
|
||||
// 未知,退到次选语言的译文不一定比原字段更合适。
|
||||
func (m GroupI18n) Localize(acceptLanguage, name, description string) (string, string) {
|
||||
text, ok := m.lookup(preferredLanguage(acceptLanguage))
|
||||
if !ok {
|
||||
return name, description
|
||||
}
|
||||
if text.Name != "" {
|
||||
name = text.Name
|
||||
}
|
||||
if text.Description != "" {
|
||||
description = text.Description
|
||||
}
|
||||
return name, description
|
||||
}
|
||||
|
||||
// LocalizeName 是只需要名称时的 Localize。
|
||||
func (m GroupI18n) LocalizeName(acceptLanguage, name string) string {
|
||||
name, _ = m.Localize(acceptLanguage, name, "")
|
||||
return name
|
||||
}
|
||||
|
||||
// CanonicalLocale 按 BCP 47 的惯用大小写书写语言代码:主语言小写,4 位文字代码首字母大写,
|
||||
// 2 位地区代码大写,其余小写(zh-hant-tw -> zh-Hant-TW);单字母扩展标记之后的部分全部小写。
|
||||
// 只调整大小写,不校验代码是否合法。
|
||||
func CanonicalLocale(locale string) string {
|
||||
subtags := strings.Split(locale, "-")
|
||||
extension := false
|
||||
for i, subtag := range subtags {
|
||||
subtag = strings.ToLower(subtag)
|
||||
switch {
|
||||
case i == 0 || extension:
|
||||
case len(subtag) == 1:
|
||||
extension = true
|
||||
case len(subtag) == 2:
|
||||
subtag = strings.ToUpper(subtag)
|
||||
case len(subtag) == 4:
|
||||
subtag = strings.ToUpper(subtag[:1]) + subtag[1:]
|
||||
}
|
||||
subtags[i] = subtag
|
||||
}
|
||||
return strings.Join(subtags, "-")
|
||||
}
|
||||
|
||||
// lookup 先按完整语言代码匹配(zh-TW),再按主语言匹配(zh-CN -> zh),不区分大小写。
|
||||
// 写入时语言代码已统一为 CanonicalLocale,同一语言只有一个键;若存有仅大小写不同的多个键,
|
||||
// 规范写法的键优先,其余按字节序取最小者,保证同一请求语言的结果固定。
|
||||
func (m GroupI18n) lookup(tag string) (GroupLocaleText, bool) {
|
||||
if tag == "" || len(m) == 0 {
|
||||
return GroupLocaleText{}, false
|
||||
}
|
||||
primary, _, _ := strings.Cut(tag, "-")
|
||||
for _, candidate := range []string{tag, primary} {
|
||||
candidate = CanonicalLocale(candidate)
|
||||
if text, ok := m[candidate]; ok {
|
||||
return text, true
|
||||
}
|
||||
matched := ""
|
||||
for locale := range m {
|
||||
if strings.EqualFold(locale, candidate) && (matched == "" || locale < matched) {
|
||||
matched = locale
|
||||
}
|
||||
}
|
||||
if matched != "" {
|
||||
return m[matched], true
|
||||
}
|
||||
}
|
||||
return GroupLocaleText{}, false
|
||||
}
|
||||
|
||||
// preferredLanguage 取 Accept-Language 中权重最高的语言代码,权重相同取靠前者;
|
||||
// 通配符 * 与 q=0 的条目不参与。
|
||||
func preferredLanguage(acceptLanguage string) string {
|
||||
best, bestQ := "", 0.0
|
||||
for _, part := range strings.Split(acceptLanguage, ",") {
|
||||
tag, params, _ := strings.Cut(part, ";")
|
||||
tag = strings.ReplaceAll(strings.TrimSpace(tag), "_", "-")
|
||||
if tag == "" || tag == "*" {
|
||||
continue
|
||||
}
|
||||
q := 1.0
|
||||
if value, ok := strings.CutPrefix(strings.TrimSpace(params), "q="); ok {
|
||||
parsed, err := strconv.ParseFloat(strings.TrimSpace(value), 64)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
q = parsed
|
||||
}
|
||||
if q > bestQ {
|
||||
best, bestQ = tag, q
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package domain
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestGroupI18nLocalize(t *testing.T) {
|
||||
i18n := GroupI18n{
|
||||
"en": {Name: "Standard", Description: "Pro + Max pool"},
|
||||
"zh-TW": {Name: "標準"},
|
||||
"ja": {Description: "説明のみ"},
|
||||
}
|
||||
const name, description = "标准", "号池"
|
||||
|
||||
tests := []struct {
|
||||
title string
|
||||
acceptLanguage string
|
||||
wantName string
|
||||
wantDescription string
|
||||
}{
|
||||
{"exact locale", "en", "Standard", "Pro + Max pool"},
|
||||
{"region falls back to the primary language", "en-US", "Standard", "Pro + Max pool"},
|
||||
{"locale match ignores case and underscore", "ZH_tw", "標準", description},
|
||||
{"a missing field keeps the original text", "ja", name, "説明のみ"},
|
||||
{"no translation for the language keeps both", "zh", name, description},
|
||||
{"primary language does not match a regional translation", "zh-CN", name, description},
|
||||
{"empty header keeps both", "", name, description},
|
||||
{"wildcard only keeps both", "*", name, description},
|
||||
// 原字段的语言未知,不退到次选语言:浏览器的 zh 用户不应拿到英文译文。
|
||||
{"only the preferred language is considered", "zh-CN,zh;q=0.9,en;q=0.8", name, description},
|
||||
{"highest weight wins over position", "zh;q=0.5, en;q=0.9", "Standard", "Pro + Max pool"},
|
||||
{"equal weight keeps the first", "en, zh", "Standard", "Pro + Max pool"},
|
||||
{"entries with q=0 are ignored", "en;q=0, ja", name, "説明のみ"},
|
||||
{"malformed weight is skipped", "en;q=abc, ja", name, "説明のみ"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.title, func(t *testing.T) {
|
||||
gotName, gotDescription := i18n.Localize(tt.acceptLanguage, name, description)
|
||||
if gotName != tt.wantName || gotDescription != tt.wantDescription {
|
||||
t.Fatalf("Localize(%q) = (%q, %q), want (%q, %q)",
|
||||
tt.acceptLanguage, gotName, gotDescription, tt.wantName, tt.wantDescription)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupI18nLocalize_NilMap(t *testing.T) {
|
||||
var i18n GroupI18n
|
||||
|
||||
name, description := i18n.Localize("en", "标准", "号池")
|
||||
|
||||
if name != "标准" || description != "号池" {
|
||||
t.Fatalf("nil translations must keep the original text, got (%q, %q)", name, description)
|
||||
}
|
||||
if got := i18n.LocalizeName("en", "标准"); got != "标准" {
|
||||
t.Fatalf("LocalizeName on nil translations = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupI18nLocalizeName(t *testing.T) {
|
||||
i18n := GroupI18n{"en": {Name: "Standard"}}
|
||||
|
||||
if got := i18n.LocalizeName("en", "标准"); got != "Standard" {
|
||||
t.Fatalf("LocalizeName(en) = %q", got)
|
||||
}
|
||||
if got := i18n.LocalizeName("zh", "标准"); got != "标准" {
|
||||
t.Fatalf("LocalizeName(zh) = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCanonicalLocale(t *testing.T) {
|
||||
tests := map[string]string{
|
||||
"en": "en",
|
||||
"EN": "en",
|
||||
"zh-tw": "zh-TW",
|
||||
"ZH-Hant-tw": "zh-Hant-TW",
|
||||
"es-419": "es-419",
|
||||
"de-CH-1901": "de-CH-1901",
|
||||
"en-US-x-Twain": "en-US-x-twain",
|
||||
}
|
||||
for in, want := range tests {
|
||||
if got := CanonicalLocale(in); got != want {
|
||||
t.Errorf("CanonicalLocale(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 写入路径已拒绝仅大小写不同的语言代码;已存在的此类数据读取时也必须得到固定结果。
|
||||
func TestGroupI18nLocalize_CaseCollidingKeysAreDeterministic(t *testing.T) {
|
||||
tests := []struct {
|
||||
title string
|
||||
i18n GroupI18n
|
||||
acceptLanguage string
|
||||
want string
|
||||
}{
|
||||
{"canonical key wins", GroupI18n{"EN": {Name: "B"}, "en": {Name: "A"}, "En": {Name: "C"}}, "en", "A"},
|
||||
{"otherwise the smallest key wins", GroupI18n{"eN": {Name: "B"}, "EN": {Name: "A"}}, "en", "A"},
|
||||
{"primary language fallback is deterministic too", GroupI18n{"ZH": {Name: "B"}, "zh": {Name: "A"}}, "zh-CN", "A"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.title, func(t *testing.T) {
|
||||
for range 200 {
|
||||
if got := tt.i18n.LocalizeName(tt.acceptLanguage, "原名"); got != tt.want {
|
||||
t.Fatalf("LocalizeName(%q) = %q, want %q", tt.acceptLanguage, got, tt.want)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -142,6 +142,8 @@ type simpleModeGroupResponse struct {
|
||||
Platform string `json:"platform"`
|
||||
Status string `json:"status"`
|
||||
|
||||
I18n service.GroupI18n `json:"i18n,omitempty"`
|
||||
|
||||
AccountCount int64 `json:"account_count,omitempty"`
|
||||
ActiveAccountCount int64 `json:"active_account_count,omitempty"`
|
||||
RateLimitedAccountCount int64 `json:"rate_limited_account_count,omitempty"`
|
||||
@@ -157,6 +159,7 @@ func groupForSimpleMode(group *service.Group) *simpleModeGroupResponse {
|
||||
return &simpleModeGroupResponse{
|
||||
ID: group.ID, Name: group.Name, Description: group.Description, Platform: group.Platform,
|
||||
Status: group.Status,
|
||||
I18n: group.I18n,
|
||||
AccountCount: group.AccountCount,
|
||||
ActiveAccountCount: group.ActiveAccountCount, RateLimitedAccountCount: group.RateLimitedAccountCount,
|
||||
SortOrder: group.SortOrder, CreatedAt: group.CreatedAt, UpdatedAt: group.UpdatedAt,
|
||||
@@ -167,7 +170,7 @@ func sanitizeCreateGroupRequestForSimpleMode(req *CreateGroupRequest) {
|
||||
if req == nil {
|
||||
return
|
||||
}
|
||||
allowed := CreateGroupRequest{Name: req.Name, Description: req.Description, Platform: req.Platform}
|
||||
allowed := CreateGroupRequest{Name: req.Name, Description: req.Description, I18n: req.I18n, Platform: req.Platform}
|
||||
allowed.RateMultiplier = 1
|
||||
allowed.SubscriptionType = service.SubscriptionTypeStandard
|
||||
*req = allowed
|
||||
@@ -177,13 +180,14 @@ func sanitizeUpdateGroupRequestForSimpleMode(req *UpdateGroupRequest) {
|
||||
if req == nil {
|
||||
return
|
||||
}
|
||||
*req = UpdateGroupRequest{Name: req.Name, Description: req.Description}
|
||||
*req = UpdateGroupRequest{Name: req.Name, Description: req.Description, I18n: req.I18n}
|
||||
}
|
||||
|
||||
// CreateGroupRequest represents create group request
|
||||
type CreateGroupRequest struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
I18n service.GroupI18n `json:"i18n"` // 名称与描述的界面语言译文
|
||||
Platform string `json:"platform" binding:"omitempty,group_platform"`
|
||||
RateMultiplier float64 `json:"rate_multiplier"`
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
@@ -258,6 +262,7 @@ type CreateGroupRequest struct {
|
||||
type UpdateGroupRequest struct {
|
||||
Name string `json:"name"`
|
||||
Description *string `json:"description"`
|
||||
I18n *service.GroupI18n `json:"i18n"` // 缺省不修改,传入则整体替换
|
||||
Platform string `json:"platform" binding:"omitempty,group_platform"`
|
||||
RateMultiplier *float64 `json:"rate_multiplier"`
|
||||
IsExclusive *bool `json:"is_exclusive"`
|
||||
@@ -664,6 +669,7 @@ func (h *GroupHandler) Create(c *gin.Context) {
|
||||
group, err := h.adminService.CreateGroup(c.Request.Context(), &service.CreateGroupInput{
|
||||
Name: req.Name,
|
||||
Description: req.Description,
|
||||
I18n: req.I18n,
|
||||
Platform: req.Platform,
|
||||
RateMultiplier: req.RateMultiplier,
|
||||
IsExclusive: req.IsExclusive,
|
||||
@@ -809,6 +815,7 @@ func (h *GroupHandler) Update(c *gin.Context) {
|
||||
group, err := h.adminService.UpdateGroup(c.Request.Context(), groupID, &service.UpdateGroupInput{
|
||||
Name: req.Name,
|
||||
Description: req.Description,
|
||||
I18n: req.I18n,
|
||||
Platform: req.Platform,
|
||||
RateMultiplier: req.RateMultiplier,
|
||||
IsExclusive: req.IsExclusive,
|
||||
|
||||
@@ -142,7 +142,7 @@ func (h *APIKeyHandler) List(c *gin.Context) {
|
||||
|
||||
out := make([]dto.APIKey, 0, len(keys))
|
||||
for i := range keys {
|
||||
out = append(out, *dto.APIKeyFromService(&keys[i]))
|
||||
out = append(out, *dto.APIKeyFromService(&keys[i]).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
response.Paginated(c, out, result.Total, page, pageSize)
|
||||
}
|
||||
@@ -174,7 +174,7 @@ func (h *APIKeyHandler) GetByID(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, dto.APIKeyFromService(key))
|
||||
response.Success(c, dto.APIKeyFromService(key).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
|
||||
// Create handles creating a new API key
|
||||
@@ -222,7 +222,7 @@ func (h *APIKeyHandler) Create(c *gin.Context) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return dto.APIKeyFromService(key), nil
|
||||
return dto.APIKeyFromService(key).Localized(c.GetHeader("Accept-Language")), nil
|
||||
})
|
||||
}
|
||||
|
||||
@@ -290,7 +290,7 @@ func (h *APIKeyHandler) Update(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, dto.APIKeyFromService(key))
|
||||
response.Success(c, dto.APIKeyFromService(key).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
|
||||
// Delete handles deleting an API key
|
||||
@@ -334,7 +334,7 @@ func (h *APIKeyHandler) GetAvailableGroups(c *gin.Context) {
|
||||
|
||||
out := make([]dto.Group, 0, len(groups))
|
||||
for i := range groups {
|
||||
out = append(out, *dto.GroupFromService(&groups[i]))
|
||||
out = append(out, *dto.GroupFromService(&groups[i]).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
response.Success(c, out)
|
||||
}
|
||||
|
||||
@@ -62,6 +62,9 @@ type userAvailableGroup struct {
|
||||
PeakEnd string `json:"peak_end"`
|
||||
PeakRateMultiplier float64 `json:"peak_rate_multiplier"`
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
|
||||
// 不对外输出;List 按请求语言把 Name 换成译文后丢弃。
|
||||
i18n service.GroupI18n
|
||||
}
|
||||
|
||||
// userSupportedModelPricing 用户可见的定价字段白名单。
|
||||
@@ -154,6 +157,7 @@ func (h *AvailableChannelHandler) List(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
acceptLanguage := c.GetHeader("Accept-Language")
|
||||
out := make([]userAvailableChannel, 0, len(channels))
|
||||
for _, ch := range channels {
|
||||
if ch.Status != service.StatusActive {
|
||||
@@ -163,6 +167,7 @@ func (h *AvailableChannelHandler) List(c *gin.Context) {
|
||||
if len(visibleGroups) == 0 {
|
||||
continue
|
||||
}
|
||||
localizeUserAvailableGroups(visibleGroups, acceptLanguage)
|
||||
sections := buildPlatformSections(ch, visibleGroups)
|
||||
if len(sections) == 0 {
|
||||
continue
|
||||
@@ -177,6 +182,13 @@ func (h *AvailableChannelHandler) List(c *gin.Context) {
|
||||
response.Success(c, out)
|
||||
}
|
||||
|
||||
// localizeUserAvailableGroups 把分组名称就地换成请求首选语言的版本。
|
||||
func localizeUserAvailableGroups(groups []userAvailableGroup, acceptLanguage string) {
|
||||
for i := range groups {
|
||||
groups[i].Name = groups[i].i18n.LocalizeName(acceptLanguage, groups[i].Name)
|
||||
}
|
||||
}
|
||||
|
||||
// buildPlatformSections 把一个渠道按 visibleGroups 的平台集合拆成有序的 section 列表:
|
||||
// 每个 section 对应一个具体平台,只包含该平台的 groups 和 supported_models。
|
||||
//
|
||||
@@ -254,6 +266,7 @@ func filterUserVisibleGroups(
|
||||
visible = append(visible, userAvailableGroup{
|
||||
ID: g.ID,
|
||||
Name: g.Name,
|
||||
i18n: g.I18n,
|
||||
Platform: g.Platform,
|
||||
SubscriptionType: g.SubscriptionType,
|
||||
RateMultiplier: g.RateMultiplier,
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
package dto
|
||||
|
||||
// 用户侧接口在输出前调用 Localized:把分组的名称与描述换成请求首选语言
|
||||
// (Accept-Language)的版本,并去掉译文表。管理端接口不调用,保留原字段与译文表,
|
||||
// 分组编辑表单依赖原字段回写。
|
||||
|
||||
// Localized 就地本地化分组并返回自身。
|
||||
func (g *Group) Localized(acceptLanguage string) *Group {
|
||||
if g == nil {
|
||||
return nil
|
||||
}
|
||||
g.Name, g.Description = g.I18n.Localize(acceptLanguage, g.Name, g.Description)
|
||||
g.I18n = nil
|
||||
return g
|
||||
}
|
||||
|
||||
// Localized 本地化 API Key 所属分组并返回自身。
|
||||
func (k *APIKey) Localized(acceptLanguage string) *APIKey {
|
||||
if k != nil {
|
||||
k.Group.Localized(acceptLanguage)
|
||||
}
|
||||
return k
|
||||
}
|
||||
|
||||
// Localized 本地化兑换码关联分组并返回自身。
|
||||
func (r *RedeemCode) Localized(acceptLanguage string) *RedeemCode {
|
||||
if r != nil {
|
||||
r.Group.Localized(acceptLanguage)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// Localized 本地化使用记录关联的分组(含其 API Key 与订阅各自携带的分组)并返回自身。
|
||||
func (l *UsageLog) Localized(acceptLanguage string) *UsageLog {
|
||||
if l != nil {
|
||||
l.Group.Localized(acceptLanguage)
|
||||
l.APIKey.Localized(acceptLanguage)
|
||||
l.Subscription.Localized(acceptLanguage)
|
||||
}
|
||||
return l
|
||||
}
|
||||
|
||||
// Localized 本地化订阅关联分组并返回自身。
|
||||
func (s *UserSubscription) Localized(acceptLanguage string) *UserSubscription {
|
||||
if s != nil {
|
||||
s.Group.Localized(acceptLanguage)
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package dto
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func localizableGroup() *Group {
|
||||
return &Group{
|
||||
ID: 1, Name: "标准", Description: "号池",
|
||||
I18n: domain.GroupI18n{"en": {Name: "Standard", Description: "Pro + Max pool"}},
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupLocalized_ReplacesTextAndDropsTranslations(t *testing.T) {
|
||||
g := localizableGroup().Localized("en")
|
||||
|
||||
require.Equal(t, "Standard", g.Name)
|
||||
require.Equal(t, "Pro + Max pool", g.Description)
|
||||
|
||||
raw, err := json.Marshal(g)
|
||||
require.NoError(t, err)
|
||||
var decoded map[string]any
|
||||
require.NoError(t, json.Unmarshal(raw, &decoded))
|
||||
require.NotContains(t, decoded, "i18n", "user-facing output must not carry the translation table")
|
||||
}
|
||||
|
||||
func TestGroupLocalized_KeepsOriginalWhenLanguageHasNoTranslation(t *testing.T) {
|
||||
g := localizableGroup().Localized("zh")
|
||||
|
||||
require.Equal(t, "标准", g.Name)
|
||||
require.Equal(t, "号池", g.Description)
|
||||
require.Nil(t, g.I18n)
|
||||
}
|
||||
|
||||
func TestLocalized_NilSafe(t *testing.T) {
|
||||
require.Nil(t, (*Group)(nil).Localized("en"))
|
||||
require.Nil(t, (*APIKey)(nil).Localized("en"))
|
||||
require.Nil(t, (*RedeemCode)(nil).Localized("en"))
|
||||
require.Nil(t, (*UsageLog)(nil).Localized("en"))
|
||||
require.Nil(t, (*UserSubscription)(nil).Localized("en"))
|
||||
|
||||
// 未携带分组的记录原样返回
|
||||
require.Nil(t, (&APIKey{}).Localized("en").Group)
|
||||
}
|
||||
|
||||
func TestLocalized_NestedGroups(t *testing.T) {
|
||||
key := (&APIKey{Group: localizableGroup()}).Localized("en")
|
||||
require.Equal(t, "Standard", key.Group.Name)
|
||||
|
||||
code := (&RedeemCode{Group: localizableGroup()}).Localized("en")
|
||||
require.Equal(t, "Standard", code.Group.Name)
|
||||
|
||||
sub := (&UserSubscription{Group: localizableGroup()}).Localized("en")
|
||||
require.Equal(t, "Standard", sub.Group.Name)
|
||||
|
||||
// 使用记录除自身分组外,还带着 API Key 与订阅各自的分组
|
||||
log := (&UsageLog{
|
||||
Group: localizableGroup(),
|
||||
APIKey: &APIKey{Group: localizableGroup()},
|
||||
Subscription: &UserSubscription{Group: localizableGroup()},
|
||||
}).Localized("en")
|
||||
require.Equal(t, "Standard", log.Group.Name)
|
||||
require.Equal(t, "Standard", log.APIKey.Group.Name)
|
||||
require.Equal(t, "Standard", log.Subscription.Group.Name)
|
||||
require.Nil(t, log.APIKey.Group.I18n)
|
||||
require.Nil(t, log.Subscription.Group.I18n)
|
||||
}
|
||||
@@ -182,6 +182,7 @@ func groupFromServiceBase(g *service.Group) Group {
|
||||
ID: g.ID,
|
||||
Name: g.Name,
|
||||
Description: g.Description,
|
||||
I18n: g.I18n,
|
||||
Platform: g.Platform,
|
||||
RateMultiplier: g.RateMultiplier,
|
||||
IsExclusive: g.IsExclusive,
|
||||
|
||||
@@ -99,6 +99,10 @@ type Group struct {
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
Status string `json:"status"`
|
||||
|
||||
// I18n 名称与描述的各语言译文。只有管理端输出;用户侧接口经 Localized 把
|
||||
// name/description 换成请求语言的版本后清空本字段。
|
||||
I18n domain.GroupI18n `json:"i18n,omitempty"`
|
||||
|
||||
SubscriptionType string `json:"subscription_type"`
|
||||
DailyLimitUSD *float64 `json:"daily_limit_usd"`
|
||||
WeeklyLimitUSD *float64 `json:"weekly_limit_usd"`
|
||||
|
||||
@@ -32,6 +32,7 @@ const (
|
||||
EndpointVideos = "/v1/videos"
|
||||
EndpointSeedanceTasks = "/api/v3/contents/generations/tasks"
|
||||
EndpointGeminiModels = "/v1beta/models"
|
||||
EndpointGeminiCachedContents = "/v1beta/cachedContents"
|
||||
)
|
||||
|
||||
const EndpointAntigravityGenerateContent = "/v1internal:streamGenerateContent"
|
||||
@@ -53,6 +54,7 @@ const (
|
||||
// "/v1/chat/completions" → "/v1/chat/completions"
|
||||
// "/openai/v1/responses/foo" → "/v1/responses"
|
||||
// "/v1beta/models/gemini:gen" → "/v1beta/models"
|
||||
// "/v1beta/cachedContents/abc" → "/v1beta/cachedContents"
|
||||
//
|
||||
// The OpenAI Responses API is also exposed via a few bare/alias
|
||||
// routes that do not carry a "/v1/" prefix (top-level bare route and
|
||||
@@ -117,6 +119,8 @@ func NormalizeInboundEndpoint(path string) string {
|
||||
return EndpointResponses
|
||||
case strings.Contains(path, EndpointGeminiModels):
|
||||
return EndpointGeminiModels
|
||||
case strings.Contains(path, EndpointGeminiCachedContents):
|
||||
return EndpointGeminiCachedContents
|
||||
default:
|
||||
return path
|
||||
}
|
||||
@@ -194,7 +198,7 @@ func isBareOrSubpathOf(path, root string) bool {
|
||||
// retain their paths. Grok raw Chat requests override this through the
|
||||
// forwarding result consumed by resolveOpenAIUpstreamEndpoint.
|
||||
// - Anthropic → /v1/messages
|
||||
// - Gemini → /v1beta/models
|
||||
// - Gemini → /v1beta/models(显式缓存管理接口保留 /v1beta/cachedContents)
|
||||
// - Antigravity → /v1/messages (Claude) or gemini (Gemini)
|
||||
// - Antigravity routes may target either Claude or Gemini, so the
|
||||
// inbound endpoint is used to distinguish.
|
||||
@@ -230,6 +234,9 @@ func DeriveUpstreamEndpoint(inbound, rawRequestPath, platform string) string {
|
||||
return EndpointSystemOne
|
||||
|
||||
case service.PlatformGemini:
|
||||
if inbound == EndpointGeminiCachedContents {
|
||||
return EndpointGeminiCachedContents
|
||||
}
|
||||
return EndpointGeminiModels
|
||||
|
||||
case service.PlatformAntigravity:
|
||||
|
||||
@@ -37,6 +37,8 @@ func TestNormalizeInboundEndpoint(t *testing.T) {
|
||||
{"/v1/videos/generations", EndpointVideosGenerations},
|
||||
{"/v1/videos/req_123", EndpointVideos},
|
||||
{"/v1beta/models", EndpointGeminiModels},
|
||||
{"/v1beta/cachedContents", EndpointGeminiCachedContents},
|
||||
{"/v1beta/cachedContents/gfo2i2haszm3wbwiyivyok5cyddubfwsr77efg7h", EndpointGeminiCachedContents},
|
||||
|
||||
// Prefixed paths (antigravity, openai) — root Responses.
|
||||
{"/antigravity/v1/messages", EndpointMessages},
|
||||
@@ -102,6 +104,7 @@ func TestDeriveUpstreamEndpoint(t *testing.T) {
|
||||
|
||||
// Gemini.
|
||||
{"gemini models", EndpointGeminiModels, "/v1beta/models/gemini:gen", service.PlatformGemini, EndpointGeminiModels},
|
||||
{"gemini cached contents", EndpointGeminiCachedContents, "/v1beta/cachedContents/abc", service.PlatformGemini, EndpointGeminiCachedContents},
|
||||
|
||||
// OpenAI — root Responses.
|
||||
{"openai responses root", EndpointResponses, "/v1/responses", service.PlatformOpenAI, EndpointResponses},
|
||||
|
||||
@@ -52,12 +52,14 @@ type GatewayHandler struct {
|
||||
errorPassthroughService *service.ErrorPassthroughService
|
||||
contentModerationService *service.ContentModerationService
|
||||
securityAuditCoordinator *securityaudit.Coordinator
|
||||
concurrencyHelper *ConcurrencyHelper
|
||||
userMsgQueueHelper *UserMsgQueueHelper
|
||||
maxAccountSwitches int
|
||||
maxAccountSwitchesGemini int
|
||||
cfg *config.Config
|
||||
settingService *service.SettingService
|
||||
// geminiCachedContentService 管理 Gemini 显式缓存;为 nil 时相关接口不可用。
|
||||
geminiCachedContentService *service.GeminiCachedContentService
|
||||
concurrencyHelper *ConcurrencyHelper
|
||||
userMsgQueueHelper *UserMsgQueueHelper
|
||||
maxAccountSwitches int
|
||||
maxAccountSwitchesGemini int
|
||||
cfg *config.Config
|
||||
settingService *service.SettingService
|
||||
}
|
||||
|
||||
// NewGatewayHandler creates a new GatewayHandler
|
||||
|
||||
@@ -0,0 +1,955 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
middleware "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"
|
||||
)
|
||||
|
||||
const (
|
||||
cachedBillingInputPrice = 0.00000125
|
||||
cachedBillingStoragePrice = 0.0000045
|
||||
cachedBillingAliasStoragePrice = 0.000009
|
||||
cachedBillingTokens = 1_000_000
|
||||
)
|
||||
|
||||
type cachedBillingRepo struct {
|
||||
mu sync.Mutex
|
||||
records map[string]*service.GeminiCachedContent
|
||||
// mutations 按记录 id 保存修改意图。
|
||||
mutations map[int64]cachedBillingMutation
|
||||
failCommit bool
|
||||
nextID int64
|
||||
deleted []int64
|
||||
}
|
||||
|
||||
type cachedBillingMutation struct {
|
||||
token string
|
||||
leaseUntil time.Time
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) Create(_ context.Context, record *service.GeminiCachedContent) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.nextID++
|
||||
record.ID = r.nextID
|
||||
cp := *record
|
||||
r.records[record.PublicID] = &cp
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) GetForOwner(_ context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
record, ok := r.records[publicID]
|
||||
if !ok || record.APIKeyID != apiKeyID || r.liveLocked(record.ID) == nil {
|
||||
return nil, service.ErrGeminiCachedContentNotFound
|
||||
}
|
||||
cp := *record
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) ListForOwner(context.Context, int64, int64, time.Time, int64, int) ([]*service.GeminiCachedContent, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// liveLocked 返回未删除的记录;调用方须持有 mu。
|
||||
func (r *cachedBillingRepo) liveLocked(id int64) *service.GeminiCachedContent {
|
||||
for _, deleted := range r.deleted {
|
||||
if deleted == id {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
for _, record := range r.records {
|
||||
if record.ID == id {
|
||||
return record
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) BeginMutation(_ context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
record := r.liveLocked(id)
|
||||
if record == nil {
|
||||
return false, service.ErrGeminiCachedContentNotFound
|
||||
}
|
||||
if _, pending := r.mutations[id]; pending || !record.ExpireTime.Equal(baseline) {
|
||||
return false, nil
|
||||
}
|
||||
r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) TakeOverStaleMutation(_ context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
record := r.liveLocked(id)
|
||||
mutation, pending := r.mutations[id]
|
||||
if record == nil || !pending || mutation.leaseUntil.After(time.Now()) {
|
||||
return nil, nil
|
||||
}
|
||||
r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)}
|
||||
cp := *record
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) CommitMutation(_ context.Context, id int64, token string, expireTime time.Time) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.failCommit {
|
||||
return errors.New("db unavailable")
|
||||
}
|
||||
record := r.liveLocked(id)
|
||||
mutation, pending := r.mutations[id]
|
||||
if record == nil || !pending || mutation.token != token {
|
||||
return service.ErrGeminiCachedContentMutationConflict
|
||||
}
|
||||
record.ExpireTime = expireTime
|
||||
delete(r.mutations, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) ReleaseMutation(_ context.Context, id int64, token string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if mutation, pending := r.mutations[id]; pending && mutation.token == token {
|
||||
delete(r.mutations, id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) ListStaleMutations(_ context.Context, limit int) ([]int64, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var ids []int64
|
||||
for id, mutation := range r.mutations {
|
||||
if len(ids) < limit && mutation.leaseUntil.Before(time.Now()) && r.liveLocked(id) != nil {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingRepo) SoftDelete(_ context.Context, id int64) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.deleted = append(r.deleted, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
// pendingMutation 返回记录是否有未清除的修改意图。
|
||||
func (r *cachedBillingRepo) pendingMutation(id int64) bool {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
_, pending := r.mutations[id]
|
||||
return pending
|
||||
}
|
||||
|
||||
// expireLease 让记录上的修改意图租约立即过期。
|
||||
func (r *cachedBillingRepo) expireLease(id int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
mutation := r.mutations[id]
|
||||
mutation.leaseUntil = time.Now().Add(-time.Second)
|
||||
r.mutations[id] = mutation
|
||||
}
|
||||
|
||||
type cachedBillingAccountRepo struct {
|
||||
service.AccountRepository
|
||||
accounts map[int64]*service.Account
|
||||
}
|
||||
|
||||
func (r *cachedBillingAccountRepo) GetByID(_ context.Context, id int64) (*service.Account, error) {
|
||||
account, ok := r.accounts[id]
|
||||
if !ok {
|
||||
return nil, service.ErrAccountNotFound
|
||||
}
|
||||
cp := *account
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
type cachedBillingUpstream struct {
|
||||
service.HTTPUpstream
|
||||
mu sync.Mutex
|
||||
createBody string
|
||||
calls []string
|
||||
// patchTTLs 按顺序记录 PATCH 请求体中的 ttl。
|
||||
patchTTLs []string
|
||||
beforePost func(n int)
|
||||
// failures 按请求方法返回上游错误响应;transportErrors 按请求方法返回传输层错误(上游结果未知)。
|
||||
failures map[string]cachedBillingUpstreamFailure
|
||||
transportErrors map[string]error
|
||||
}
|
||||
|
||||
type cachedBillingUpstreamFailure struct {
|
||||
status int
|
||||
body string
|
||||
}
|
||||
|
||||
func (u *cachedBillingUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
var patch struct {
|
||||
TTL string `json:"ttl"`
|
||||
}
|
||||
if req.Method == http.MethodPatch {
|
||||
_ = json.NewDecoder(req.Body).Decode(&patch)
|
||||
}
|
||||
u.mu.Lock()
|
||||
u.calls = append(u.calls, req.Method+" "+req.URL.Path)
|
||||
n := len(u.calls)
|
||||
if req.Method == http.MethodPatch {
|
||||
u.patchTTLs = append(u.patchTTLs, patch.TTL)
|
||||
}
|
||||
failure, failed := u.failures[req.Method]
|
||||
transportErr := u.transportErrors[req.Method]
|
||||
u.mu.Unlock()
|
||||
if transportErr != nil {
|
||||
return nil, transportErr
|
||||
}
|
||||
if failed {
|
||||
return &http.Response{StatusCode: failure.status, Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(failure.body))}, nil
|
||||
}
|
||||
body := "{}"
|
||||
switch req.Method {
|
||||
case http.MethodPost:
|
||||
body = u.createBody
|
||||
case http.MethodPatch:
|
||||
if u.beforePost != nil {
|
||||
u.beforePost(n)
|
||||
}
|
||||
seconds, _ := strconv.ParseInt(strings.TrimSuffix(patch.TTL, "s"), 10, 64)
|
||||
body = fmt.Sprintf(`{"name":"cachedContents/up1","expireTime":%q}`, time.Now().Add(time.Duration(seconds)*time.Second).UTC().Format(time.RFC3339Nano))
|
||||
}
|
||||
return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(body))}, nil
|
||||
}
|
||||
|
||||
func (u *cachedBillingUpstream) methodCalls(method string) int {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
n := 0
|
||||
for _, call := range u.calls {
|
||||
if strings.HasPrefix(call, method+" ") {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
type cachedBillingUserRepo struct {
|
||||
service.UserRepository
|
||||
balance float64
|
||||
}
|
||||
|
||||
func (r *cachedBillingUserRepo) GetByID(_ context.Context, id int64) (*service.User, error) {
|
||||
return &service.User{ID: id, Balance: r.balance}, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingUserRepo) DeductBalance(context.Context, int64, float64) error { return nil }
|
||||
|
||||
type cachedBillingUsageRepo struct {
|
||||
service.UsageLogRepository
|
||||
mu sync.Mutex
|
||||
logs []*service.UsageLog
|
||||
}
|
||||
|
||||
func (r *cachedBillingUsageRepo) Create(_ context.Context, log *service.UsageLog) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.logs = append(r.logs, log)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *cachedBillingUsageRepo) first() *service.UsageLog {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.logs[0]
|
||||
}
|
||||
|
||||
func (r *cachedBillingUsageRepo) totals() []float64 {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make([]float64, 0, len(r.logs))
|
||||
for _, log := range r.logs {
|
||||
out = append(out, log.TotalCost)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
type cachedBillingEnv struct {
|
||||
h *GatewayHandler
|
||||
repo *cachedBillingRepo
|
||||
upstream *cachedBillingUpstream
|
||||
usage *cachedBillingUsageRepo
|
||||
apiKey *service.APIKey
|
||||
group *service.Group
|
||||
// accounts 是调度快照桩直接返回的账号对象,测试可改其凭据。
|
||||
accounts map[int64]*service.Account
|
||||
}
|
||||
|
||||
func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *cachedBillingEnv {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
require.NoError(t, os.WriteFile(filepath.Join(dir, "model_pricing.json"), []byte(fmt.Sprintf(`{"gemini-2.5-pro":{
|
||||
"litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001,
|
||||
"cache_storage_cost_per_token_per_hour":%g},"gemini-alias":{
|
||||
"litellm_provider":"gemini","mode":"chat","input_cost_per_token":%g,"output_cost_per_token":0.00001,
|
||||
"cache_storage_cost_per_token_per_hour":%g}}`,
|
||||
cachedBillingInputPrice, cachedBillingStoragePrice, cachedBillingInputPrice, cachedBillingAliasStoragePrice)), 0o644))
|
||||
cfg := &config.Config{}
|
||||
cfg.Pricing.DataDir = dir
|
||||
cfg.Pricing.UpdateIntervalHours = 100000
|
||||
cfg.Default.RateMultiplier = 1
|
||||
pricing := service.NewPricingService(cfg, nil)
|
||||
require.NoError(t, pricing.Initialize())
|
||||
billing := service.NewBillingService(cfg, pricing)
|
||||
|
||||
groupID := int64(42)
|
||||
group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformGemini, Status: service.StatusActive, RateMultiplier: 1}
|
||||
accounts := map[int64]*service.Account{}
|
||||
var scheduled []*service.Account
|
||||
for _, id := range accountIDs {
|
||||
account := &service.Account{ID: id, Platform: service.PlatformGemini, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "k"}, Status: service.StatusActive, Schedulable: true,
|
||||
Concurrency: 5, Priority: 1, AccountGroups: []service.AccountGroup{{AccountID: id, GroupID: groupID}}}
|
||||
accounts[id] = account
|
||||
scheduled = append(scheduled, account)
|
||||
}
|
||||
userRepo := &cachedBillingUserRepo{balance: balance}
|
||||
usage := &cachedBillingUsageRepo{}
|
||||
bcs := service.NewBillingCacheService(nil, userRepo, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(bcs.Stop)
|
||||
snapshot := service.NewSchedulerSnapshotService(&fakeSchedulerCache{accounts: scheduled}, nil, nil, nil, nil)
|
||||
gw := service.NewGatewayService(nil, &fakeGroupRepo{group: group}, usage, nil, userRepo, nil, nil, nil, cfg, snapshot, nil,
|
||||
billing, nil, &service.BillingCacheService{}, nil, nil, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
upstream := &cachedBillingUpstream{}
|
||||
compat := service.NewGeminiMessagesCompatService(&cachedBillingAccountRepo{accounts: accounts}, nil, nil, nil, nil, nil, upstream, nil, cfg)
|
||||
repo := &cachedBillingRepo{records: map[string]*service.GeminiCachedContent{}, mutations: map[int64]cachedBillingMutation{}}
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gw,
|
||||
billingCacheService: bcs,
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatNone, 0),
|
||||
geminiCachedContentService: service.NewGeminiCachedContentService(repo, compat, cfg),
|
||||
}
|
||||
apiKey := &service.APIKey{ID: 7, UserID: 3, GroupID: &groupID, Group: group, Status: service.StatusActive,
|
||||
User: &service.User{ID: 3, Concurrency: 10, Balance: balance}}
|
||||
return &cachedBillingEnv{h: h, repo: repo, upstream: upstream, usage: usage, apiKey: apiKey, group: group, accounts: accounts}
|
||||
}
|
||||
|
||||
// TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot 候选过滤直接按快照里的凭据判定:
|
||||
// 第三方地址且未声明上游显式缓存的 Key 账号在选号前就被排除,官方地址与声明了上游显式缓存的账号保留。
|
||||
func TestGeminiExplicitCacheExclusions_FiltersUnsupportedAccountsFromSnapshot(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 500, 900, 901, 902)
|
||||
env.accounts[901].Credentials["base_url"] = "https://third.example.com"
|
||||
env.accounts[902].Credentials["base_url"] = "https://gw.example.com"
|
||||
env.accounts[902].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true
|
||||
|
||||
groupID := env.group.ID
|
||||
excluded, err := env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 0)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, map[int64]struct{}{901: {}}, excluded)
|
||||
|
||||
excluded, err = env.h.gatewayService.GeminiExplicitCacheExclusions(context.Background(), &groupID, 902)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, map[int64]struct{}{900: {}, 901: {}}, excluded, "绑定缓存的请求只保留持有账号")
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) serve(method, target, body string, params gin.Params, call func(*gin.Context)) *httptest.ResponseRecorder {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(method, target, strings.NewReader(body))
|
||||
c.Request = req.WithContext(context.WithValue(req.Context(), ctxkey.Group, e.group))
|
||||
c.Params = params
|
||||
c.Set(string(middleware.ContextKeyAPIKey), e.apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: e.apiKey.UserID, Concurrency: 10})
|
||||
call(c)
|
||||
return rec
|
||||
}
|
||||
|
||||
// waitForUsage 等待异步计费任务落库(无 worker 池时同步执行,直接返回)。
|
||||
func (e *cachedBillingEnv) waitForUsage(t *testing.T, n int) []float64 {
|
||||
t.Helper()
|
||||
require.Eventually(t, func() bool { return len(e.usage.totals()) >= n }, time.Second, 10*time.Millisecond)
|
||||
return e.usage.totals()
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) create() *httptest.ResponseRecorder {
|
||||
return e.serve(http.MethodPost, "/v1beta/cachedContents",
|
||||
`{"model":"models/gemini-2.5-pro","ttl":"86400s","contents":[{"role":"user","parts":[{"text":"doc"}]}]}`,
|
||||
nil, e.h.GeminiCachedContentsCreate)
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) seed(expire time.Time) {
|
||||
e.repo.records["abc"] = &service.GeminiCachedContent{
|
||||
ID: 100, PublicID: "abc", UserID: 3, APIKeyID: 7, GroupID: 42, AccountID: 900,
|
||||
UpstreamName: "cachedContents/up1", Model: "gemini-2.5-pro", UpstreamModel: "models/gemini-2.5-pro",
|
||||
TotalTokenCount: cachedBillingTokens, ExpireTime: expire,
|
||||
}
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) patch() *httptest.ResponseRecorder {
|
||||
return e.serve(http.MethodPatch, "/v1beta/cachedContents/abc?updateMask=ttl", `{"ttl":"86400s"}`,
|
||||
gin.Params{{Key: "cacheID", Value: "abc"}}, e.h.GeminiCachedContentsPatch)
|
||||
}
|
||||
|
||||
// TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate 延长有效期沿用创建时的渠道计费口径:
|
||||
// 计费模型来源为请求模型时按别名定价,用量记录带渠道 ID。
|
||||
func TestGeminiCachedContentsPatch_UsesChannelUsageRecordedAtCreate(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 500, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.repo.records["abc"].ChannelUsage = service.ChannelUsageFields{
|
||||
ChannelID: 5, OriginalModel: "gemini-alias", ChannelMappedModel: "gemini-2.5-pro",
|
||||
BillingModelSource: service.BillingModelSourceRequested,
|
||||
}
|
||||
|
||||
rec := env.patch()
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingAliasStoragePrice*23, totals[0], 0.01, "按创建时的计费模型(请求模型别名)计存储费")
|
||||
usageLog := env.usage.first()
|
||||
require.NotNil(t, usageLog.ChannelID)
|
||||
require.Equal(t, int64(5), *usageLog.ChannelID)
|
||||
require.True(t, strings.HasPrefix(usageLog.RequestID, "gcache:patch:abc:"), usageLog.RequestID)
|
||||
}
|
||||
|
||||
func cachedBillingCreateBody(expire time.Time) string {
|
||||
return fmt.Sprintf(`{"name":"cachedContents/up1","usageMetadata":{"totalTokenCount":%d},"expireTime":%q}`,
|
||||
cachedBillingTokens, expire.UTC().Format(time.RFC3339Nano))
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsCreate_ChargesInputAndStorage(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 500, 900)
|
||||
env.upstream.createBody = cachedBillingCreateBody(time.Now().Add(100 * time.Hour))
|
||||
|
||||
rec := env.create()
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingInputPrice+cachedBillingTokens*cachedBillingStoragePrice*24, totals[0], 1e-9)
|
||||
require.Len(t, env.repo.records, 1)
|
||||
for _, record := range env.repo.records {
|
||||
require.False(t, record.ExpireTime.After(time.Now().Add(24*time.Hour)), "本地到期时间不超过请求的有效期")
|
||||
require.Equal(t, "gcache:create:"+record.PublicID, env.usage.first().RequestID, "费用行按缓存 ID 可检索")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsCreate_UnaffordableStorageDeletesUpstreamAndBillsInputOnly(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1, 900)
|
||||
env.upstream.createBody = cachedBillingCreateBody(time.Now().Add(24 * time.Hour))
|
||||
|
||||
rec := env.create()
|
||||
require.NotEqual(t, http.StatusOK, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "error")
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete), "删除已创建的上游缓存,避免产生存储费")
|
||||
require.Empty(t, env.repo.records, "不登记缓存")
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingInputPrice, totals[0], 1e-9, "只收已发生的缓存 token 输入费")
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsCreate_UnusableUpstreamResponseCleansUpWithoutRetry(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 500, 900, 901)
|
||||
env.upstream.createBody = `{"name":"cachedContents/up1","expireTime":"2099-01-01T00:00:00Z"}`
|
||||
|
||||
rec := env.create()
|
||||
require.Equal(t, http.StatusBadGateway, rec.Code)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPost), "不换号重建")
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete), "删除已创建的上游缓存")
|
||||
require.Empty(t, env.repo.records)
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
func (e *cachedBillingEnv) patchTTL(id, ttl string) *httptest.ResponseRecorder {
|
||||
return e.serve(http.MethodPatch, "/v1beta/cachedContents/"+id+"?updateMask=ttl", fmt.Sprintf(`{"ttl":%q}`, ttl),
|
||||
gin.Params{{Key: "cacheID", Value: id}}, e.h.GeminiCachedContentsPatch)
|
||||
}
|
||||
|
||||
// requireTTLNear 断言发往上游的 ttl("Ns")与 want 相差不超过 5 秒。
|
||||
func requireTTLNear(t *testing.T, ttl string, want time.Duration) {
|
||||
t.Helper()
|
||||
seconds, err := strconv.ParseInt(strings.TrimSuffix(ttl, "s"), 10, 64)
|
||||
require.NoError(t, err, ttl)
|
||||
require.InDelta(t, want.Seconds(), float64(seconds), 5, ttl)
|
||||
}
|
||||
|
||||
// 同一缓存的修改串行执行:第一个 PATCH 的上游调用返回前,第二个 PATCH 直接得到 409 且不调用上游,
|
||||
// 本地到期时间与上游最后生效的修改一致,只按第一个 PATCH 实际推进的时长收费。
|
||||
func TestGeminiCachedContentsPatch_ConcurrentUpdateConflictsAndKeepsUpstreamInSync(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.beforePost = func(n int) {
|
||||
if n == 1 {
|
||||
rec := env.patchTTL("abc", "86400s")
|
||||
require.Equal(t, http.StatusConflict, rec.Code, rec.Body.String())
|
||||
require.Contains(t, rec.Body.String(), `"ABORTED"`)
|
||||
}
|
||||
}
|
||||
|
||||
rec := env.patchTTL("abc", "7200s")
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "冲突的修改不调用上游")
|
||||
require.WithinDuration(t, time.Now().Add(2*time.Hour), env.repo.records["abc"].ExpireTime, 5*time.Second, "本地到期时间与上游一致")
|
||||
require.False(t, env.repo.pendingMutation(100))
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice, totals[0], cachedBillingTokens*cachedBillingStoragePrice*time.Minute.Hours(), "只收推进的 1 小时")
|
||||
}
|
||||
|
||||
// 上游已修改但本地写回失败:不计费并保留修改意图;租约内的修改返回 409,
|
||||
// 租约过期后后台恢复把上游有效期回滚到修改前的到期时间并清除意图。
|
||||
func TestGeminiCachedContentsPatch_CommitFailureRollsBackUpstreamWithoutCharge(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
baseline := time.Now().Add(time.Hour)
|
||||
env.seed(baseline)
|
||||
env.repo.failCommit = true
|
||||
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch))
|
||||
require.True(t, env.repo.pendingMutation(100), "上游可能已修改:保留修改意图")
|
||||
require.Empty(t, env.usage.totals(), "未提交的修改不收费")
|
||||
|
||||
env.repo.failCommit = false
|
||||
require.Equal(t, http.StatusConflict, env.patch().Code, "租约内的修改被拒绝")
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch))
|
||||
|
||||
env.repo.expireLease(100)
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 2, env.upstream.methodCalls(http.MethodPatch))
|
||||
requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline))
|
||||
require.False(t, env.repo.pendingMutation(100))
|
||||
require.True(t, env.repo.records["abc"].ExpireTime.Equal(baseline), "本地到期时间保持修改前的值")
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
// 租约过期的遗留修改意图在下一次修改前先回滚,随后的修改按修改前的到期时间计费。
|
||||
func TestGeminiCachedContentsPatch_StaleMutationRolledBackBeforeNextUpdate(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
baseline := time.Now().Add(time.Hour)
|
||||
env.seed(baseline)
|
||||
env.repo.failCommit = true
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
env.repo.failCommit = false
|
||||
env.repo.expireLease(100)
|
||||
|
||||
rec := env.patch()
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.Equal(t, 3, env.upstream.methodCalls(http.MethodPatch), "失败的修改、回滚、本次修改")
|
||||
requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline))
|
||||
require.Equal(t, "86400s", env.upstream.patchTTLs[2])
|
||||
totals := env.waitForUsage(t, 1)
|
||||
require.Len(t, totals, 1)
|
||||
require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice*23, totals[0], 0.01)
|
||||
}
|
||||
|
||||
// 上游结果未知(传输层错误)时保留修改意图等待回滚;上游明确拒绝时清除修改意图。两种情况都不收费。
|
||||
func TestGeminiCachedContentsPatch_MutationKeptOnlyWhenUpstreamOutcomeUnknown(t *testing.T) {
|
||||
t.Run("transport error", func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.transportErrors = map[string]error{http.MethodPatch: errors.New("connection reset")}
|
||||
|
||||
require.Equal(t, http.StatusBadGateway, env.patch().Code)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.usage.totals())
|
||||
})
|
||||
t.Run("upstream rejected", func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodPatch: {status: http.StatusBadRequest,
|
||||
body: `{"error":{"code":400,"message":"invalid ttl","status":"INVALID_ARGUMENT"}}`}}
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, env.patch().Code)
|
||||
require.False(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.usage.totals())
|
||||
})
|
||||
}
|
||||
|
||||
// 修改前的到期时间已过:恢复时删除上游缓存与本地记录。
|
||||
func TestGeminiCachedContentMutationRecovery_DeletesUpstreamWhenBaselineExpired(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(-time.Minute))
|
||||
env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)}
|
||||
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete))
|
||||
require.Zero(t, env.upstream.methodCalls(http.MethodPatch))
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
}
|
||||
|
||||
// cachedBillingReservationCache 是启用在途预留时的计费缓存桩:余额与订阅用量取固定值,各维度预留按 key 记账。
|
||||
type cachedBillingReservationCache struct {
|
||||
service.BillingCache
|
||||
mu sync.Mutex
|
||||
balance float64
|
||||
sub *service.SubscriptionCacheData
|
||||
reserved map[string]map[string]float64
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) GetUserBalance(context.Context, int64) (float64, error) {
|
||||
return c.balance, nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) GetSubscriptionCache(context.Context, int64, int64) (*service.SubscriptionCacheData, error) {
|
||||
return c.sub, nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) reserveLocked(key, requestID string, amount float64) float64 {
|
||||
if c.reserved == nil {
|
||||
c.reserved = map[string]map[string]float64{}
|
||||
}
|
||||
if c.reserved[key] == nil {
|
||||
c.reserved[key] = map[string]float64{}
|
||||
}
|
||||
others := 0.0
|
||||
for _, v := range c.reserved[key] {
|
||||
others += v
|
||||
}
|
||||
c.reserved[key][requestID] = amount
|
||||
return others
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) ReserveInflightBalance(_ context.Context, userID int64, requestID string, amount, balance float64, _ time.Duration) (bool, float64, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
key := fmt.Sprintf("user:%d", userID)
|
||||
others := c.reserveLocked(key, requestID, amount)
|
||||
if len(c.reserved[key]) > 1 && balance-others < amount {
|
||||
delete(c.reserved[key], requestID)
|
||||
return false, others, nil
|
||||
}
|
||||
return true, others, nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) ReleaseInflightBalance(_ context.Context, userID int64, requestID string) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
delete(c.reserved[fmt.Sprintf("user:%d", userID)], requestID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) ReserveInflightScope(_ context.Context, scope, requestID string, amount float64, _ time.Duration) (float64, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.reserveLocked(scope, requestID, amount), nil
|
||||
}
|
||||
|
||||
func (c *cachedBillingReservationCache) ReleaseInflightScope(_ context.Context, scope, requestID string) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
delete(c.reserved[scope], requestID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func enableCachedBillingReservations(t *testing.T, env *cachedBillingEnv, cache *cachedBillingReservationCache) {
|
||||
t.Helper()
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation.Enabled = true
|
||||
bcs := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(bcs.Stop)
|
||||
env.h.billingCacheService = bcs
|
||||
}
|
||||
|
||||
// 确定金额的续期费按 API Key 与余额维度预留:第一笔续期(约 $103.5)尚未入账时,
|
||||
// 第二个缓存的续期超出 Key 剩余额度或余额,在调用上游前被拒绝。
|
||||
func TestGeminiCachedContentsPatch_UnsettledChargeHoldsKeyQuotaAndBalance(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
balance, quota float64
|
||||
}{
|
||||
{name: "api key quota", balance: 1000, quota: 110},
|
||||
{name: "balance", balance: 110},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, tc.balance, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
second := *env.repo.records["abc"]
|
||||
second.ID, second.PublicID, second.UpstreamName = 101, "def", "cachedContents/up2"
|
||||
env.repo.records["def"] = &second
|
||||
env.apiKey.Quota = tc.quota
|
||||
enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: tc.balance})
|
||||
env.upstream.beforePost = func(n int) {
|
||||
if n == 1 {
|
||||
rec := env.patchTTL("def", "86400s")
|
||||
require.NotEqual(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
require.Equal(t, http.StatusOK, env.patch().Code)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "第二笔续期没有到达上游")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 订阅模式的确定金额费用按订阅维度预留:第一笔未入账时第二笔超出日剩余额度被拒绝,释放后放行。
|
||||
func TestReserveGeminiCachedContentCharge_UnsettledChargeHoldsSubscriptionHeadroom(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
limit := 110.0
|
||||
env.group.SubscriptionType = service.SubscriptionTypeSubscription
|
||||
env.group.DailyLimitUSD = &limit
|
||||
enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: 1000,
|
||||
sub: &service.SubscriptionCacheData{Status: service.SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour)}})
|
||||
sub := &service.UserSubscription{ID: 1, UserID: env.apiKey.UserID, GroupID: env.group.ID, Status: service.SubscriptionStatusActive,
|
||||
ExpiresAt: time.Now().Add(time.Hour)}
|
||||
reserve := func() (func(), error) {
|
||||
var done func()
|
||||
var err error
|
||||
env.serve(http.MethodPatch, "/", `{}`, nil, func(c *gin.Context) {
|
||||
done, err = env.h.reserveGeminiCachedContentCharge(c, env.apiKey, sub, 103.5)
|
||||
})
|
||||
return done, err
|
||||
}
|
||||
|
||||
firstDone, err := reserve()
|
||||
require.NoError(t, err)
|
||||
_, err = reserve()
|
||||
require.ErrorIs(t, err, service.ErrDailyLimitExceeded)
|
||||
firstDone()
|
||||
secondDone, err := reserve()
|
||||
require.NoError(t, err)
|
||||
secondDone()
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsPatch_UnaffordableExtensionRejectedBeforeUpstream(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
|
||||
rec := env.patch()
|
||||
require.NotEqual(t, http.StatusOK, rec.Code)
|
||||
require.Zero(t, env.upstream.methodCalls(http.MethodPatch), "余额不足以支付延长部分时不调用上游")
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
// 上游对缓存不存在的响应:AI Studio 403;Vertex 缓存资源接口 404,文案含上游缓存数字 ID。
|
||||
var cachedBillingUpstreamNotFoundCases = []struct {
|
||||
name string
|
||||
failure cachedBillingUpstreamFailure
|
||||
}{
|
||||
{name: "ai studio 403", failure: cachedBillingUpstreamFailure{status: http.StatusForbidden,
|
||||
body: `{"error":{"code":403,"message":"CachedContent not found (or permission denied)","status":"PERMISSION_DENIED"}}`}},
|
||||
{name: "vertex 404", failure: cachedBillingUpstreamFailure{status: http.StatusNotFound,
|
||||
body: `{"error":{"code":404,"message":"Cached content 8079767076522164224 is not found.","status":"NOT_FOUND"}}`}},
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsPatch_UpstreamNotFoundForgetsAndReturnsStandardForbidden(t *testing.T) {
|
||||
for _, tc := range cachedBillingUpstreamNotFoundCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodPatch: tc.failure}
|
||||
|
||||
rec := env.patch()
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Equal(t, service.GeminiCachedContentNotFoundResponse, rec.Body.String())
|
||||
require.NotContains(t, rec.Body.String(), "8079767076522164224", "不应暴露上游缓存 ID")
|
||||
require.Equal(t, []int64{100}, env.repo.deleted, "上游已不存在的缓存应从本地移除")
|
||||
require.Empty(t, env.usage.totals(), "上游未延长不收费")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsDelete_UpstreamNotFoundStillForgets(t *testing.T) {
|
||||
for _, tc := range cachedBillingUpstreamNotFoundCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodDelete: tc.failure}
|
||||
|
||||
rec := env.serve(http.MethodDelete, "/v1beta/cachedContents/abc", "",
|
||||
gin.Params{{Key: "cacheID", Value: "abc"}}, env.h.GeminiCachedContentsDelete)
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.JSONEq(t, `{}`, rec.Body.String())
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 持有账号暂时不支持显式缓存(能力被关闭)时,恢复流程不能把修改意图当作已回滚而清除;
|
||||
// 能力恢复后再回滚上游并清除意图。
|
||||
func TestGeminiCachedContentMutationRecovery_KeepsIntentWhileAccountCannotCache(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.accounts[900].Credentials["base_url"] = "https://gw.example.com"
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true
|
||||
baseline := time.Now().Add(time.Hour)
|
||||
env.seed(baseline)
|
||||
env.repo.failCommit = true
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch))
|
||||
env.repo.failCommit = false
|
||||
env.repo.expireLease(100)
|
||||
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err, "账号暂不支持缓存时恢复失败")
|
||||
require.Zero(t, recovered)
|
||||
require.True(t, env.repo.pendingMutation(100), "未回滚上游前保留修改意图")
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "不改动上游")
|
||||
require.Empty(t, env.repo.deleted)
|
||||
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true
|
||||
env.repo.expireLease(100)
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 2, env.upstream.methodCalls(http.MethodPatch))
|
||||
requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline))
|
||||
require.False(t, env.repo.pendingMutation(100))
|
||||
require.True(t, env.repo.records["abc"].ExpireTime.Equal(baseline))
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
// 持有账号长期不支持显式缓存,且修改前到期时间加上单次最长有效期都已过去:上游缓存必已自行过期,
|
||||
// 恢复流程不再保留意图,直接移除本地记录。
|
||||
func TestGeminiCachedContentMutationRecovery_AbandonsIntentOnceUpstreamMustHaveExpired(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.accounts[900].Credentials["base_url"] = "https://gw.example.com"
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = true
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.repo.failCommit = true
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
env.repo.failCommit = false
|
||||
env.repo.expireLease(100)
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false
|
||||
|
||||
// 修改前到期时间 + 24h 上限 + 1h 容忍仍未过去:保留意图。
|
||||
env.repo.records["abc"].ExpireTime = time.Now().Add(-20 * time.Hour)
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err)
|
||||
require.Zero(t, recovered)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.repo.deleted)
|
||||
|
||||
env.repo.expireLease(100)
|
||||
env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour)
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "不改动上游")
|
||||
require.Zero(t, env.upstream.methodCalls(http.MethodDelete))
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
require.Empty(t, env.usage.totals())
|
||||
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, recovered, "作废的记录不再被选中")
|
||||
}
|
||||
|
||||
// 作废判定的边界:默认上限 24h 加 1h 容忍;配置上限低于 24h 时仍按 24h 兜底。
|
||||
func TestGeminiCachedContentMutationRecovery_AbandonBoundary(t *testing.T) {
|
||||
disableCache := func(env *cachedBillingEnv) {
|
||||
env.accounts[900].Credentials["base_url"] = "https://gw.example.com"
|
||||
env.accounts[900].Credentials[service.GeminiExplicitCacheUpstreamCredentialKey] = false
|
||||
}
|
||||
t.Run("within grace keeps intent", func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(-25*time.Hour + time.Minute))
|
||||
env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)}
|
||||
disableCache(env)
|
||||
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err)
|
||||
require.Zero(t, recovered)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.repo.deleted)
|
||||
})
|
||||
t.Run("configured max ttl below default still waits a day", func(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(-12 * time.Hour))
|
||||
env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)}
|
||||
disableCache(env)
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.GeminiCachedContentMaxTTLSeconds = 3600
|
||||
compat := service.NewGeminiMessagesCompatService(&cachedBillingAccountRepo{accounts: env.accounts}, nil, nil, nil, nil, nil, env.upstream, nil, cfg)
|
||||
svc := service.NewGeminiCachedContentService(env.repo, compat, cfg)
|
||||
|
||||
recovered, err := svc.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err)
|
||||
require.Zero(t, recovered)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
|
||||
env.repo.expireLease(100)
|
||||
env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour)
|
||||
recovered, err = svc.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
})
|
||||
}
|
||||
|
||||
// 持有账号可用但上游回滚持续失败(如 Key 已失效):上游缓存必已自行过期后不再重试,移除本地记录。
|
||||
func TestGeminiCachedContentMutationRecovery_AbandonsIntentWhenRollbackKeepsFailing(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(-time.Minute))
|
||||
env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)}
|
||||
env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodDelete: {status: http.StatusUnauthorized,
|
||||
body: `{"error":{"code":401,"message":"API key not valid","status":"UNAUTHENTICATED"}}`}}
|
||||
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.Error(t, err, "上游缓存可能仍存活:保留意图等待重试")
|
||||
require.Zero(t, recovered)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete))
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
require.Empty(t, env.repo.deleted)
|
||||
|
||||
env.repo.expireLease(100)
|
||||
env.repo.records["abc"].ExpireTime = time.Now().Add(-26 * time.Hour)
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 2, env.upstream.methodCalls(http.MethodDelete), "账号可用时仍先尝试回滚")
|
||||
require.Equal(t, []int64{100}, env.repo.deleted)
|
||||
require.Empty(t, env.usage.totals())
|
||||
}
|
||||
|
||||
// 持有账号已删除时没有凭据可回滚上游:恢复流程移除本地记录并清除意图,不再重试。
|
||||
func TestGeminiCachedContentMutationRecovery_ForgetsRecordWhenAccountDeleted(t *testing.T) {
|
||||
env := newCachedBillingEnv(t, 1000, 900)
|
||||
env.seed(time.Now().Add(time.Hour))
|
||||
env.repo.failCommit = true
|
||||
require.Equal(t, http.StatusInternalServerError, env.patch().Code)
|
||||
require.True(t, env.repo.pendingMutation(100))
|
||||
env.repo.failCommit = false
|
||||
env.repo.expireLease(100)
|
||||
|
||||
delete(env.accounts, 900)
|
||||
recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, recovered)
|
||||
require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "无凭据,不改动上游")
|
||||
require.Zero(t, env.upstream.methodCalls(http.MethodDelete))
|
||||
require.Equal(t, []int64{100}, env.repo.deleted, "本地记录作废")
|
||||
require.Empty(t, env.usage.totals())
|
||||
require.Equal(t, http.StatusForbidden, env.serve(http.MethodGet, "/v1beta/cachedContents/abc", "", gin.Params{{Key: "cacheID", Value: "abc"}}, env.h.GeminiCachedContentsGet).Code, "作废的记录对客户端不可见")
|
||||
|
||||
recovered, err = env.h.geminiCachedContentService.RecoverStaleMutations(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, recovered, "不再重复处理")
|
||||
}
|
||||
@@ -0,0 +1,848 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// geminiCachedContentCreateMaxAttempts 限制创建缓存时实际发往上游的账号数。
|
||||
const geminiCachedContentCreateMaxAttempts = 3
|
||||
|
||||
// geminiCachedContentPersistTimeout 是上游操作完成后落库的超时;落库与客户端连接脱钩。
|
||||
const geminiCachedContentPersistTimeout = 15 * time.Second
|
||||
|
||||
const geminiCachedContentUnavailableMessage = "The account holding this cached content is temporarily unavailable"
|
||||
|
||||
// geminiCachedContentsPreamble 校验显式缓存管理接口的公共前置条件:gemini 分组、非 composite 分组、功能已装配。
|
||||
func (h *GatewayHandler) geminiCachedContentsPreamble(c *gin.Context, component string) (*service.APIKey, middleware.AuthSubject, *zap.Logger, bool) {
|
||||
apiKey, ok := middleware.GetAPIKeyFromContext(c)
|
||||
if !ok || apiKey == nil {
|
||||
googleError(c, http.StatusUnauthorized, "Invalid API key")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
authSubject, ok := middleware.GetAuthSubjectFromContext(c)
|
||||
if !ok {
|
||||
googleError(c, http.StatusInternalServerError, "User context not found")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
if h.geminiCachedContentService == nil {
|
||||
googlePlatformError(c, http.StatusNotFound, "Explicit context caching is not enabled")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
if apiKey.Group == nil || apiKey.GroupID == nil || apiKey.Group.Platform == service.PlatformComposite {
|
||||
googleError(c, http.StatusBadRequest, "Explicit context caching is not supported for this API key group")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
if effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return nil, middleware.AuthSubject{}, nil, false
|
||||
}
|
||||
reqLog := requestLogger(
|
||||
c,
|
||||
component,
|
||||
zap.Int64("user_id", authSubject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
)
|
||||
return apiKey, authSubject, reqLog, true
|
||||
}
|
||||
|
||||
func readGeminiCachedContentBody(c *gin.Context) ([]byte, bool) {
|
||||
body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request)
|
||||
if err != nil {
|
||||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||||
googleError(c, http.StatusRequestEntityTooLarge, buildBodyTooLargeMessage(maxErr.Limit))
|
||||
return nil, false
|
||||
}
|
||||
googleError(c, http.StatusBadRequest, "Failed to read request body")
|
||||
return nil, false
|
||||
}
|
||||
if len(body) == 0 {
|
||||
googleError(c, http.StatusBadRequest, "Request body is empty")
|
||||
return nil, false
|
||||
}
|
||||
return body, true
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) checkGeminiCachedContentBilling(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subscription *service.UserSubscription) bool {
|
||||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
|
||||
reqLog.Info("gemini.cached_content.billing_eligibility_check_failed", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// GeminiCachedContentsCreate POST /v1beta/cachedContents
|
||||
// 创建请求内一次性收取缓存 token 的输入费与整个有效期的存储费。
|
||||
func (h *GatewayHandler) GeminiCachedContentsCreate(c *gin.Context) {
|
||||
apiKey, authSubject, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.create")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
svc := h.geminiCachedContentService
|
||||
|
||||
body, ok := readGeminiCachedContentBody(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
req, err := service.ParseGeminiCachedContentCreateRequest(body, svc.Now(), svc.MaxTTL())
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
modelName := req.Model
|
||||
reqLog = reqLog.With(zap.String("model", modelName))
|
||||
|
||||
setOpsRequestContext(c, modelName, false)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(false, false)))
|
||||
pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context())
|
||||
c.Request = c.Request.WithContext(pricingCtx)
|
||||
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, modelName, body); decision != nil && !decision.AllowNextStage {
|
||||
googleSecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, modelName)
|
||||
reqModel := modelName
|
||||
if channelMapping.Mapped {
|
||||
modelName = channelMapping.MappedModel
|
||||
}
|
||||
|
||||
subscription, _ := middleware.GetSubscriptionFromContext(c)
|
||||
geminiConcurrency := NewConcurrencyHelper(h.concurrencyHelper.concurrencyService, SSEPingFormatNone, 0)
|
||||
streamStarted := false
|
||||
userReleaseFunc, err := geminiConcurrency.AcquireUserSlotWithWait(c, authSubject.UserID, authSubject.Concurrency, false, &streamStarted)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.cached_content.user_slot_acquire_failed", zap.Error(err))
|
||||
googleConcurrencyError(c, err, "user")
|
||||
return
|
||||
}
|
||||
userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
|
||||
if userReleaseFunc != nil {
|
||||
defer userReleaseFunc()
|
||||
}
|
||||
|
||||
if !h.checkGeminiCachedContentBilling(c, reqLog, apiKey, subscription) {
|
||||
return
|
||||
}
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.cached_content.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
excluded, err := h.gatewayService.GeminiExplicitCacheExclusions(c.Request.Context(), apiKey.GroupID, 0)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.cached_content.list_accounts_failed", zap.Error(err))
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts support explicit context caching")
|
||||
return
|
||||
}
|
||||
|
||||
var lastUpstreamErr *service.GeminiCachedContentUpstreamError
|
||||
var ttlRejectedLimit time.Duration
|
||||
for attempts := 0; attempts < geminiCachedContentCreateMaxAttempts; {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, "", modelName, excluded, "", int64(0))
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
reqLog.Info("gemini.cached_content.account_select_exhausted", zap.Error(err))
|
||||
break
|
||||
}
|
||||
account := selection.Account
|
||||
if !service.IsGeminiExplicitCacheAccount(account) {
|
||||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
ttl, ttlAllowed := req.TTLForAccount(account)
|
||||
if !ttlAllowed {
|
||||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
excluded[account.ID] = struct{}{}
|
||||
ttlRejectedLimit = max(ttlRejectedLimit, account.GeminiExplicitCacheUpstreamMaxTTL())
|
||||
continue
|
||||
}
|
||||
accountReleaseFunc, ok := h.acquireGeminiCachedContentAccountSlot(c, reqLog, geminiConcurrency, selection)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
admissionCtx := service.ContextWithSelectionProfitGate(c.Request.Context(), selection)
|
||||
latest, vetoed, reason := h.gatewayService.GatewayProfitControlVetoLatest(admissionCtx, account)
|
||||
if vetoed {
|
||||
accountReleaseFunc()
|
||||
reqLog.Debug("gemini.cached_content.account_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason))
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
account = latest
|
||||
attempts++
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
|
||||
mappedModel := account.GetMappedModel(modelName)
|
||||
usageFields := clientRequestedUsageFields(c, channelMapping, reqModel, mappedModel)
|
||||
if !h.gatewayService.CacheStoragePriced(c.Request.Context(), &service.RecordUsageInput{
|
||||
Result: &service.ForwardResult{Model: modelName, UpstreamModel: mappedModel, CacheStorageTokenHours: 1},
|
||||
APIKey: apiKey,
|
||||
ChannelUsageFields: usageFields,
|
||||
}) {
|
||||
accountReleaseFunc()
|
||||
googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", modelName))
|
||||
return
|
||||
}
|
||||
upstreamModel, err := service.GeminiCachedContentUpstreamModelResource(account, mappedModel)
|
||||
if err != nil {
|
||||
accountReleaseFunc()
|
||||
reqLog.Warn("gemini.cached_content.upstream_model_invalid", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
publicID, err := service.NewGeminiCachedContentPublicID()
|
||||
if err != nil {
|
||||
accountReleaseFunc()
|
||||
googleError(c, http.StatusInternalServerError, "Failed to allocate cached content id")
|
||||
return
|
||||
}
|
||||
|
||||
upstream, err := svc.CreateUpstream(c.Request.Context(), account, req, mappedModel, publicID, ttl)
|
||||
accountReleaseFunc()
|
||||
if err != nil {
|
||||
var upErr *service.GeminiCachedContentUpstreamError
|
||||
if errors.As(err, &upErr) {
|
||||
reqLog.Info("gemini.cached_content.upstream_create_failed", zap.Int64("account_id", account.ID), zap.Int("status", upErr.StatusCode))
|
||||
if upErr.RetryableOnOtherAccount() {
|
||||
lastUpstreamErr = upErr
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
h.writeGeminiCachedContentUpstreamError(c, upErr)
|
||||
return
|
||||
}
|
||||
// 上游已返回成功但响应不可用:不换号重建,能识别出资源名时删除上游缓存。
|
||||
var malformed *service.GeminiCachedContentMalformedResponseError
|
||||
if errors.As(err, &malformed) {
|
||||
reqLog.Error("gemini.cached_content.upstream_create_malformed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
if malformed.Name != "" {
|
||||
h.cleanupGeminiCachedContentUpstream(reqLog, account, malformed.Name)
|
||||
}
|
||||
googleError(c, http.StatusBadGateway, "Upstream returned an invalid cached content response")
|
||||
return
|
||||
}
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
reqLog.Warn("gemini.cached_content.upstream_create_error", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
excluded[account.ID] = struct{}{}
|
||||
continue
|
||||
}
|
||||
|
||||
// 缓存 token 数在上游创建后才确定:按确定金额校验余额 / 额度并预留,不足时删除上游缓存,
|
||||
// 只收取已发生的缓存 token 输入费。
|
||||
tokens := upstream.TotalTokenCount
|
||||
storageTokenHours := service.GeminiCachedContentStorageTokenHours(tokens, ttl)
|
||||
amount := h.gatewayService.QuoteUsageCost(c.Request.Context(), &service.RecordUsageInput{
|
||||
Result: &service.ForwardResult{
|
||||
Model: modelName, UpstreamModel: mappedModel,
|
||||
Usage: service.ClaudeUsage{InputTokens: int(tokens)}, CacheStorageTokenHours: storageTokenHours,
|
||||
},
|
||||
APIKey: apiKey, User: apiKey.User, Subscription: subscription, PricingAt: pricingAt, ChannelUsageFields: usageFields,
|
||||
})
|
||||
inflightRelease()
|
||||
chargeDone, err := h.reserveGeminiCachedContentCharge(c, apiKey, subscription, amount)
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.cached_content.storage_charge_rejected", zap.Float64("amount", amount), zap.Error(err))
|
||||
h.cleanupGeminiCachedContentUpstream(reqLog, account, upstream.Name)
|
||||
h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt,
|
||||
service.GeminiCachedContentCreateUsageRequestID(publicID), modelName, mappedModel, int(tokens), 0, usageFields, body)
|
||||
writeGeminiCachedContentBillingError(c, err)
|
||||
return
|
||||
}
|
||||
defer chargeDone()
|
||||
|
||||
expireTime := service.ClampGeminiCachedContentExpire(upstream.ExpireTime, svc.Now().Add(ttl))
|
||||
record := &service.GeminiCachedContent{
|
||||
PublicID: publicID,
|
||||
UserID: authSubject.UserID,
|
||||
APIKeyID: apiKey.ID,
|
||||
GroupID: *apiKey.GroupID,
|
||||
AccountID: account.ID,
|
||||
UpstreamName: upstream.Name,
|
||||
Model: modelName,
|
||||
UpstreamModel: upstreamModel,
|
||||
DisplayName: req.DisplayName,
|
||||
TotalTokenCount: tokens,
|
||||
ExpireTime: expireTime,
|
||||
ChannelUsage: usageFields,
|
||||
}
|
||||
persistCtx, cancelPersist := geminiCachedContentPersistContext(c)
|
||||
err = svc.Save(persistCtx, record)
|
||||
cancelPersist()
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.save_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
h.cleanupGeminiCachedContentUpstream(reqLog, account, record.UpstreamName)
|
||||
googleError(c, http.StatusInternalServerError, "Failed to persist cached content")
|
||||
return
|
||||
}
|
||||
|
||||
h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt,
|
||||
service.GeminiCachedContentCreateUsageRequestID(publicID), modelName, mappedModel, int(tokens), storageTokenHours, usageFields, body)
|
||||
c.JSON(http.StatusOK, service.GeminiCachedContentView(record))
|
||||
return
|
||||
}
|
||||
|
||||
if lastUpstreamErr != nil {
|
||||
h.writeGeminiCachedContentUpstreamError(c, lastUpstreamErr)
|
||||
return
|
||||
}
|
||||
if ttlRejectedLimit > 0 {
|
||||
googleError(c, http.StatusBadRequest, fmt.Sprintf("ttl exceeds the maximum allowed %ds by the upstream accounts", int64(ttlRejectedLimit.Seconds())))
|
||||
return
|
||||
}
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts support explicit context caching")
|
||||
}
|
||||
|
||||
// GeminiCachedContentsList GET /v1beta/cachedContents
|
||||
// 只返回当前 API Key 在当前分组下仍有效的缓存,不透传上游 list(同一账号上其他租户的缓存不可见)。
|
||||
func (h *GatewayHandler) GeminiCachedContentsList(c *gin.Context) {
|
||||
apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.list")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
pageSize := service.GeminiCachedContentListDefaultPageSize
|
||||
if raw := strings.TrimSpace(c.Query("pageSize")); raw != "" {
|
||||
n, err := strconv.Atoi(raw)
|
||||
if err != nil || n < 0 {
|
||||
googleError(c, http.StatusBadRequest, "pageSize must be a non-negative integer")
|
||||
return
|
||||
}
|
||||
if n > 0 {
|
||||
pageSize = min(n, service.GeminiCachedContentListMaxPageSize)
|
||||
}
|
||||
}
|
||||
beforeID := int64(0)
|
||||
if raw := strings.TrimSpace(c.Query("pageToken")); raw != "" {
|
||||
id, ok := decodeGeminiCachedContentPageToken(raw)
|
||||
if !ok {
|
||||
googleError(c, http.StatusBadRequest, "Invalid pageToken")
|
||||
return
|
||||
}
|
||||
beforeID = id
|
||||
}
|
||||
records, err := h.geminiCachedContentService.List(c.Request.Context(), apiKey.ID, *apiKey.GroupID, beforeID, pageSize+1)
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.list_failed", zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to list cached contents")
|
||||
return
|
||||
}
|
||||
resp := gin.H{}
|
||||
if len(records) > pageSize {
|
||||
records = records[:pageSize]
|
||||
resp["nextPageToken"] = encodeGeminiCachedContentPageToken(records[len(records)-1].ID)
|
||||
}
|
||||
if len(records) > 0 {
|
||||
views := make([]map[string]any, 0, len(records))
|
||||
for _, record := range records {
|
||||
views = append(views, service.GeminiCachedContentView(record))
|
||||
}
|
||||
resp["cachedContents"] = views
|
||||
}
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GeminiCachedContentsGet GET /v1beta/cachedContents/:cacheID
|
||||
func (h *GatewayHandler) GeminiCachedContentsGet(c *gin.Context) {
|
||||
apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.get")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
record, ok := h.resolveGeminiCachedContentOrWrite(c, reqLog, apiKey, c.Param("cacheID"))
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, service.GeminiCachedContentView(record))
|
||||
}
|
||||
|
||||
// GeminiCachedContentsPatch PATCH /v1beta/cachedContents/:cacheID
|
||||
// 只允许修改有效期;延长部分在本次请求内按增量时长收取存储费,缩短不退费。
|
||||
// 同一缓存的修改串行执行:上游修改、本地写回与计费都在修改意图的租约内完成,并发的修改返回 409。
|
||||
// 上游结果未知或本地写回失败的修改不计费,租约过期后由恢复流程把上游有效期回滚到修改前的到期时间。
|
||||
func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) {
|
||||
apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.patch")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
svc := h.geminiCachedContentService
|
||||
body, ok := readGeminiCachedContentBody(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
record, ok := h.resolveGeminiCachedContentOrWrite(c, reqLog, apiKey, c.Param("cacheID"))
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
account, err := svc.BoundAccount(c.Request.Context(), record)
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.load_account_failed", zap.Int64("account_id", record.AccountID), zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to load cached content account")
|
||||
return
|
||||
}
|
||||
if account == nil {
|
||||
googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage)
|
||||
return
|
||||
}
|
||||
ttl, err := service.ParseGeminiCachedContentPatchRequest(body, c.Query("updateMask"), svc.Now(), svc.MaxTTLForAccount(account))
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
setOpsRequestContext(c, record.Model, false)
|
||||
pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context())
|
||||
c.Request = c.Request.WithContext(pricingCtx)
|
||||
|
||||
subscription, _ := middleware.GetSubscriptionFromContext(c)
|
||||
if extension := svc.Now().Add(ttl).Sub(record.ExpireTime); extension > 0 {
|
||||
usage := &service.RecordUsageInput{
|
||||
Result: &service.ForwardResult{
|
||||
Model: record.Model, UpstreamModel: record.UpstreamModelID(),
|
||||
CacheStorageTokenHours: service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extension),
|
||||
},
|
||||
APIKey: apiKey, User: apiKey.User, Subscription: subscription, PricingAt: pricingAt,
|
||||
ChannelUsageFields: record.ChannelUsage,
|
||||
}
|
||||
if !h.gatewayService.CacheStoragePriced(c.Request.Context(), usage) {
|
||||
googlePlatformError(c, http.StatusBadRequest, fmt.Sprintf("Explicit context caching is not available for model %s: storage pricing is not configured", record.Model))
|
||||
return
|
||||
}
|
||||
if !h.checkGeminiCachedContentBilling(c, reqLog, apiKey, subscription) {
|
||||
return
|
||||
}
|
||||
chargeDone, err := h.reserveGeminiCachedContentCharge(c, apiKey, subscription, h.gatewayService.QuoteUsageCost(c.Request.Context(), usage))
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.cached_content.storage_charge_rejected", zap.Error(err))
|
||||
writeGeminiCachedContentBillingError(c, err)
|
||||
return
|
||||
}
|
||||
defer chargeDone()
|
||||
}
|
||||
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
mutation, err := svc.BeginMutation(c.Request.Context(), record)
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, service.ErrGeminiCachedContentNotFound):
|
||||
googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage)
|
||||
case errors.Is(err, service.ErrGeminiCachedContentMutationConflict):
|
||||
reqLog.Info("gemini.cached_content.mutation_conflict", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
writeGeminiCachedContentMutationConflict(c)
|
||||
default:
|
||||
reqLog.Error("gemini.cached_content.begin_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to update cached content")
|
||||
}
|
||||
return
|
||||
}
|
||||
upstream, err := svc.PatchUpstream(c.Request.Context(), account, record, ttl)
|
||||
if err != nil {
|
||||
var upErr *service.GeminiCachedContentUpstreamError
|
||||
upstreamRejected := errors.As(err, &upErr)
|
||||
if upstreamRejected || errors.Is(err, service.ErrGeminiCachedContentRequestNotSent) {
|
||||
// 上游有效期未变:清除修改意图。
|
||||
persistCtx, cancelPersist := geminiCachedContentPersistContext(c)
|
||||
if releaseErr := svc.ReleaseMutation(persistCtx, mutation); releaseErr != nil {
|
||||
reqLog.Warn("gemini.cached_content.release_mutation_failed", zap.Error(releaseErr))
|
||||
}
|
||||
cancelPersist()
|
||||
}
|
||||
if upstreamRejected {
|
||||
if upErr.NotFound() {
|
||||
if forgetErr := svc.Forget(c.Request.Context(), record); forgetErr != nil {
|
||||
reqLog.Warn("gemini.cached_content.forget_failed", zap.Error(forgetErr))
|
||||
}
|
||||
}
|
||||
h.writeGeminiCachedContentUpstreamError(c, upErr)
|
||||
return
|
||||
}
|
||||
reqLog.Warn("gemini.cached_content.upstream_patch_error", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
googleError(c, http.StatusBadGateway, "Upstream request failed")
|
||||
return
|
||||
}
|
||||
expireTime := service.ClampGeminiCachedContentExpire(upstream.ExpireTime, svc.Now().Add(ttl))
|
||||
persistCtx, cancelPersist := geminiCachedContentPersistContext(c)
|
||||
err = svc.CommitMutation(persistCtx, mutation, expireTime)
|
||||
cancelPersist()
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.commit_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
if errors.Is(err, service.ErrGeminiCachedContentMutationConflict) {
|
||||
writeGeminiCachedContentMutationConflict(c)
|
||||
return
|
||||
}
|
||||
googleError(c, http.StatusInternalServerError, "Failed to persist cached content expiration")
|
||||
return
|
||||
}
|
||||
// 提交前库中的到期时间一直是修改前的值:按本次实际推进的时长收费。
|
||||
if extended := expireTime.Sub(mutation.Baseline()); extended > 0 {
|
||||
h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt,
|
||||
service.GeminiCachedContentPatchUsageRequestID(record.PublicID), record.Model, record.UpstreamModelID(), 0,
|
||||
service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), record.ChannelUsage, body)
|
||||
}
|
||||
c.JSON(http.StatusOK, service.GeminiCachedContentView(record))
|
||||
}
|
||||
|
||||
// GeminiCachedContentsDelete DELETE /v1beta/cachedContents/:cacheID
|
||||
// 删除不退还已收取的存储费。
|
||||
func (h *GatewayHandler) GeminiCachedContentsDelete(c *gin.Context) {
|
||||
apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.delete")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
svc := h.geminiCachedContentService
|
||||
record, ok := h.resolveGeminiCachedContentOrWrite(c, reqLog, apiKey, c.Param("cacheID"))
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
account, err := svc.BoundAccount(c.Request.Context(), record)
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.load_account_failed", zap.Int64("account_id", record.AccountID), zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to load cached content account")
|
||||
return
|
||||
}
|
||||
if account != nil {
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
if err := svc.DeleteUpstream(c.Request.Context(), account, record); err != nil {
|
||||
var upErr *service.GeminiCachedContentUpstreamError
|
||||
if !errors.As(err, &upErr) {
|
||||
reqLog.Warn("gemini.cached_content.upstream_delete_error", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
googleError(c, http.StatusBadGateway, "Upstream request failed")
|
||||
return
|
||||
}
|
||||
if !upErr.NotFound() {
|
||||
h.writeGeminiCachedContentUpstreamError(c, upErr)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
persistCtx, cancelPersist := geminiCachedContentPersistContext(c)
|
||||
err = svc.Forget(persistCtx, record)
|
||||
cancelPersist()
|
||||
if err != nil {
|
||||
reqLog.Error("gemini.cached_content.forget_failed", zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to delete cached content")
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{})
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) resolveGeminiCachedContentOrWrite(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, cacheID string) (*service.GeminiCachedContent, bool) {
|
||||
record, err := h.geminiCachedContentService.Resolve(c.Request.Context(), apiKey.ID, apiKey.GroupID, cacheID)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrGeminiCachedContentNotFound) {
|
||||
googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage)
|
||||
return nil, false
|
||||
}
|
||||
reqLog.Error("gemini.cached_content.lookup_failed", zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to load cached content")
|
||||
return nil, false
|
||||
}
|
||||
return record, true
|
||||
}
|
||||
|
||||
// resolveGeminiCachedContentBinding 处理 generateContent / countTokens 请求对显式缓存的引用:
|
||||
// 校验归属与模型后把引用改写为上游资源名,返回绑定记录;请求未引用缓存时返回 nil。
|
||||
// 返回 false 表示已写出错误响应。
|
||||
func (h *GatewayHandler) resolveGeminiCachedContentBinding(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, modelName string, body []byte) (*service.GeminiCachedContent, []byte, bool) {
|
||||
ref, err := service.FindGeminiCachedContentReference(body)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadRequest, err.Error())
|
||||
return nil, nil, false
|
||||
}
|
||||
if ref == nil {
|
||||
return nil, body, true
|
||||
}
|
||||
if middleware.HasForcePlatform(c) || h.geminiCachedContentService == nil ||
|
||||
apiKey.Group == nil || apiKey.Group.Platform == service.PlatformComposite {
|
||||
googleError(c, http.StatusBadRequest, "cachedContent is not supported on this route")
|
||||
return nil, nil, false
|
||||
}
|
||||
record, err := h.geminiCachedContentService.Resolve(c.Request.Context(), apiKey.ID, apiKey.GroupID, ref.Name)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrGeminiCachedContentNotFound) {
|
||||
googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage)
|
||||
return nil, nil, false
|
||||
}
|
||||
reqLog.Error("gemini.cached_content.lookup_failed", zap.Error(err))
|
||||
googleError(c, http.StatusInternalServerError, "Failed to load cached content")
|
||||
return nil, nil, false
|
||||
}
|
||||
if record.Model != modelName {
|
||||
googleError(c, http.StatusBadRequest, fmt.Sprintf(
|
||||
"Model used by GenerateContent request (models/%s) and CachedContent (models/%s) has to be the same.", modelName, record.Model))
|
||||
return nil, nil, false
|
||||
}
|
||||
rewritten, err := service.RewriteGeminiCachedContentReference(body, ref, record.UpstreamName)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadRequest, "Invalid cachedContent reference")
|
||||
return nil, nil, false
|
||||
}
|
||||
return record, rewritten, true
|
||||
}
|
||||
|
||||
// acquireGeminiCachedContentAccountSlot 获取选中账号的并发槽(含等待计划);返回 false 表示已写出错误响应。
|
||||
func (h *GatewayHandler) acquireGeminiCachedContentAccountSlot(c *gin.Context, reqLog *zap.Logger, helper *ConcurrencyHelper, selection *service.AccountSelectionResult) (func(), bool) {
|
||||
if selection.Acquired {
|
||||
release := selection.ReleaseFunc
|
||||
if release == nil {
|
||||
release = func() {}
|
||||
}
|
||||
return wrapReleaseOnDone(c.Request.Context(), release), true
|
||||
}
|
||||
account := selection.Account
|
||||
if selection.WaitPlan == nil {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts")
|
||||
return nil, false
|
||||
}
|
||||
canWait, err := helper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.cached_content.account_wait_counter_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
} else if !canWait {
|
||||
googleError(c, http.StatusTooManyRequests, "Too many pending requests, please retry later")
|
||||
return nil, false
|
||||
}
|
||||
streamStarted := false
|
||||
release, acquireErr := helper.AcquireAccountSlotWithWaitTimeout(c, account.ID, selection.WaitPlan.MaxConcurrency, selection.WaitPlan.Timeout, false, &streamStarted)
|
||||
if err == nil {
|
||||
helper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
||||
}
|
||||
if acquireErr != nil {
|
||||
reqLog.Warn("gemini.cached_content.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(acquireErr))
|
||||
googleConcurrencyError(c, acquireErr, "account")
|
||||
return nil, false
|
||||
}
|
||||
if release == nil {
|
||||
release = func() {}
|
||||
}
|
||||
return wrapReleaseOnDone(c.Request.Context(), release), true
|
||||
}
|
||||
|
||||
// submitGeminiCachedContentUsage 记录缓存创建 / 延长的计费:inputTokens 按输入价计,storageTokenHours 按计费模型的存储单价计;
|
||||
// requestID 是带固定前缀的费用行标识(见 GeminiCachedContentUsageRequestIDPrefix),落库时不被请求 id 覆盖。
|
||||
func (h *GatewayHandler) submitGeminiCachedContentUsage(
|
||||
c *gin.Context,
|
||||
apiKey *service.APIKey,
|
||||
account *service.Account,
|
||||
subscription *service.UserSubscription,
|
||||
pricingAt time.Time,
|
||||
requestID string,
|
||||
model string,
|
||||
upstreamModel string,
|
||||
inputTokens int,
|
||||
storageTokenHours float64,
|
||||
usageFields service.ChannelUsageFields,
|
||||
payload []byte,
|
||||
) {
|
||||
if inputTokens <= 0 && storageTokenHours <= 0 {
|
||||
return
|
||||
}
|
||||
if upstreamModel == model {
|
||||
upstreamModel = ""
|
||||
}
|
||||
result := &service.ForwardResult{
|
||||
RequestID: requestID,
|
||||
Model: model,
|
||||
UpstreamModel: upstreamModel,
|
||||
Usage: service.ClaudeUsage{InputTokens: inputTokens},
|
||||
CacheStorageTokenHours: storageTokenHours,
|
||||
}
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
requestPayloadHash := service.HashUsageRequestPayload(payload)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
PricingAt: pricingAt,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: usageFields,
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.gemini_v1beta.cached_contents"),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.String("request_id", result.RequestID),
|
||||
).Error("gemini.cached_content.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// geminiCachedContentPersistContext 返回与客户端连接脱钩的落库 context:上游操作已生效时,客户端断开不应让本地记录缺失。
|
||||
func geminiCachedContentPersistContext(c *gin.Context) (context.Context, context.CancelFunc) {
|
||||
return context.WithTimeout(context.WithoutCancel(c.Request.Context()), geminiCachedContentPersistTimeout)
|
||||
}
|
||||
|
||||
// cleanupGeminiCachedContentUpstream 尽力删除已在上游创建、但网关不会登记的缓存,避免其持续产生存储费。
|
||||
func (h *GatewayHandler) cleanupGeminiCachedContentUpstream(reqLog *zap.Logger, account *service.Account, upstreamName string) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), geminiCachedContentPersistTimeout)
|
||||
defer cancel()
|
||||
if err := h.geminiCachedContentService.DeleteUpstream(ctx, account, &service.GeminiCachedContent{UpstreamName: upstreamName}); err != nil {
|
||||
reqLog.Error("gemini.cached_content.orphan_cleanup_failed", zap.Int64("account_id", account.ID), zap.String("upstream_name", upstreamName), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// reserveGeminiCachedContentCharge 校验余额 / 订阅额度 / API Key 额度足以支付确定金额 amount(倍率后),
|
||||
// 并按该金额在用户、订阅、API Key 维度登记在途预留(不受预留上限封顶,见 ReserveFixedCharge)。
|
||||
// 返回的 done 须在 handler 结束时调用;预留随计费任务交接,落库后释放。
|
||||
func (h *GatewayHandler) reserveGeminiCachedContentCharge(c *gin.Context, apiKey *service.APIKey, subscription *service.UserSubscription, amount float64) (func(), error) {
|
||||
ctx := c.Request.Context()
|
||||
res, err := h.billingCacheService.ReserveFixedCharge(ctx, apiKey.User, apiKey, apiKey.Group, subscription, amount)
|
||||
if err != nil {
|
||||
return inflightNoop, err
|
||||
}
|
||||
if res == nil {
|
||||
return inflightNoop, nil
|
||||
}
|
||||
c.Request = c.Request.WithContext(service.WithInflightReservation(ctx, res))
|
||||
return res.HandlerDone, nil
|
||||
}
|
||||
|
||||
// writeGeminiCachedContentMutationConflict 回写同一缓存的修改冲突(Google 并发冲突语义 ABORTED),客户端可稍后重试。
|
||||
func writeGeminiCachedContentMutationConflict(c *gin.Context) {
|
||||
c.JSON(http.StatusConflict, gin.H{
|
||||
"error": gin.H{
|
||||
"code": http.StatusConflict,
|
||||
"message": "Another update to this cached content is in progress, please retry later",
|
||||
"status": "ABORTED",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func writeGeminiCachedContentBillingError(c *gin.Context, err error) {
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
}
|
||||
|
||||
// writeGeminiCachedContentUpstreamError 回写上游错误:缓存不存在统一返回标准 403(各上游形态不同且可能含上游资源 ID);
|
||||
// 账号凭据或权限类错误不是客户端问题,统一改为 502;限流 / 不可用按普通 Gemini 路径处理(透传规则、通用映射),
|
||||
// 只有上游网关账号的这类响应反映持有缓存的上游账号状态,原样回写。
|
||||
func (h *GatewayHandler) writeGeminiCachedContentUpstreamError(c *gin.Context, upErr *service.GeminiCachedContentUpstreamError) {
|
||||
status := upErr.StatusCode
|
||||
message := strings.TrimSpace(extractGoogleErrorMessage(upErr.Body))
|
||||
service.SetOpsUpstreamError(c, status, message, "")
|
||||
if upErr.NotFound() {
|
||||
service.MarkOpsRequestScopedError(c, "permission_error")
|
||||
c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse))
|
||||
return
|
||||
}
|
||||
if (status == http.StatusTooManyRequests || status >= http.StatusInternalServerError) && !upErr.UpstreamGateway {
|
||||
h.handleGeminiFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: status, ResponseBody: upErr.Body})
|
||||
return
|
||||
}
|
||||
if status == http.StatusUnauthorized || status == http.StatusForbidden {
|
||||
googleError(c, http.StatusBadGateway, "Upstream account rejected the cached content request")
|
||||
return
|
||||
}
|
||||
if len(upErr.Body) == 0 || !json.Valid(upErr.Body) {
|
||||
if message == "" {
|
||||
message = http.StatusText(status)
|
||||
}
|
||||
googleError(c, status, message)
|
||||
return
|
||||
}
|
||||
c.Data(status, "application/json", upErr.Body)
|
||||
}
|
||||
|
||||
func extractGoogleErrorMessage(body []byte) string {
|
||||
var envelope struct {
|
||||
Error struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &envelope); err != nil {
|
||||
return ""
|
||||
}
|
||||
return envelope.Error.Message
|
||||
}
|
||||
|
||||
func encodeGeminiCachedContentPageToken(id int64) string {
|
||||
return base64.RawURLEncoding.EncodeToString([]byte(strconv.FormatInt(id, 10)))
|
||||
}
|
||||
|
||||
func decodeGeminiCachedContentPageToken(raw string) (int64, bool) {
|
||||
decoded, err := base64.RawURLEncoding.DecodeString(raw)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
id, err := strconv.ParseInt(string(decoded), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
// mergeAccountIDSets 返回两个账号 ID 集合的并集(不修改入参)。
|
||||
func mergeAccountIDSets(a, b map[int64]struct{}) map[int64]struct{} {
|
||||
out := make(map[int64]struct{}, len(a)+len(b))
|
||||
for id := range a {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
for id := range b {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// handleGeminiCachedContentFailoverExhausted 处理绑定缓存的请求在持有账号上重试耗尽:没有其他账号可换。
|
||||
// 上游网关账号的限流 / 不可用响应反映持有缓存的上游账号状态,原样返回供客户端判断;
|
||||
// 其余(含官方账号)仍按普通 Gemini 路径处理:透传规则、通用映射。
|
||||
func (h *GatewayHandler) handleGeminiCachedContentFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError) {
|
||||
if failoverErr == nil || !failoverErr.BoundUpstreamPassthrough || len(failoverErr.ResponseBody) == 0 || !json.Valid(failoverErr.ResponseBody) {
|
||||
h.handleGeminiFailoverExhausted(c, failoverErr)
|
||||
return
|
||||
}
|
||||
service.SetOpsUpstreamError(c, failoverErr.StatusCode, service.ExtractUpstreamErrorMessage(failoverErr.ResponseBody), "")
|
||||
c.Data(failoverErr.StatusCode, "application/json", failoverErr.ResponseBody)
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
middleware "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"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type geminiCachedContentHandlerRepoStub struct {
|
||||
records []*service.GeminiCachedContent
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) Create(_ context.Context, record *service.GeminiCachedContent) error {
|
||||
record.ID = int64(len(r.records) + 1)
|
||||
r.records = append(r.records, record)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) GetForOwner(_ context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) {
|
||||
for _, record := range r.records {
|
||||
if record.PublicID == publicID && record.APIKeyID == apiKeyID {
|
||||
return record, nil
|
||||
}
|
||||
}
|
||||
return nil, service.ErrGeminiCachedContentNotFound
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) ListForOwner(_ context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*service.GeminiCachedContent, error) {
|
||||
var out []*service.GeminiCachedContent
|
||||
for i := len(r.records) - 1; i >= 0; i-- {
|
||||
record := r.records[i]
|
||||
if record.APIKeyID == apiKeyID && record.GroupID == groupID && record.ExpireTime.After(activeAt) && (beforeID <= 0 || record.ID < beforeID) {
|
||||
out = append(out, record)
|
||||
}
|
||||
if len(out) == limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) BeginMutation(context.Context, int64, time.Time, string, time.Duration) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) TakeOverStaleMutation(context.Context, int64, string, time.Duration) (*service.GeminiCachedContent, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) CommitMutation(context.Context, int64, string, time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) ReleaseMutation(context.Context, int64, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) ListStaleMutations(context.Context, int) ([]int64, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentHandlerRepoStub) SoftDelete(context.Context, int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func newGeminiCachedContentHandlerForTest(repo *geminiCachedContentHandlerRepoStub) *GatewayHandler {
|
||||
return &GatewayHandler{geminiCachedContentService: service.NewGeminiCachedContentService(repo, nil, &config.Config{})}
|
||||
}
|
||||
|
||||
func geminiCachedContentTestAPIKey(platform string) *service.APIKey {
|
||||
groupID := int64(42)
|
||||
return &service.APIKey{ID: 7, UserID: 3, GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: platform}}
|
||||
}
|
||||
|
||||
func newGeminiCachedContentTestContext(method, target string, apiKey *service.APIKey) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(method, target, nil)
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
return c, rec
|
||||
}
|
||||
|
||||
func seedGeminiCachedContent(repo *geminiCachedContentHandlerRepoStub, publicID string, apiKeyID int64, expire time.Time) *service.GeminiCachedContent {
|
||||
record := &service.GeminiCachedContent{
|
||||
PublicID: publicID,
|
||||
UserID: 3,
|
||||
APIKeyID: apiKeyID,
|
||||
GroupID: 42,
|
||||
AccountID: 900,
|
||||
UpstreamName: "cachedContents/upstream" + publicID,
|
||||
Model: "gemini-3.8-flash",
|
||||
UpstreamModel: "models/gemini-3.8-flash",
|
||||
TotalTokenCount: 6273,
|
||||
ExpireTime: expire,
|
||||
CreatedAt: time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC),
|
||||
UpdatedAt: time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC),
|
||||
}
|
||||
_ = repo.Create(context.Background(), record)
|
||||
return record
|
||||
}
|
||||
|
||||
func TestResolveGeminiCachedContentBinding(t *testing.T) {
|
||||
repo := &geminiCachedContentHandlerRepoStub{}
|
||||
seedGeminiCachedContent(repo, "abc", 7, time.Now().Add(time.Hour))
|
||||
seedGeminiCachedContent(repo, "foreign", 8, time.Now().Add(time.Hour))
|
||||
h := newGeminiCachedContentHandlerForTest(repo)
|
||||
apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
|
||||
|
||||
t.Run("request without cache passes through", func(t *testing.T) {
|
||||
c, _ := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
body := []byte(`{"contents":[]}`)
|
||||
record, out, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash", body)
|
||||
require.True(t, ok)
|
||||
require.Nil(t, record)
|
||||
require.Equal(t, body, out)
|
||||
})
|
||||
|
||||
t.Run("owned cache rewrites reference and binds account", func(t *testing.T) {
|
||||
c, _ := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
record, out, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash",
|
||||
[]byte(`{"cachedContent":"cachedContents/abc","contents":[]}`))
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(900), record.AccountID)
|
||||
require.JSONEq(t, `{"cachedContent":"cachedContents/upstreamabc","contents":[]}`, string(out))
|
||||
})
|
||||
|
||||
t.Run("cache owned by another key looks missing", func(t *testing.T) {
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
_, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash",
|
||||
[]byte(`{"cachedContent":"cachedContents/foreign","contents":[]}`))
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), service.GeminiCachedContentNotFoundMessage)
|
||||
})
|
||||
|
||||
t.Run("model must match the cache literally", func(t *testing.T) {
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-flash-latest:generateContent", apiKey)
|
||||
_, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-flash-latest",
|
||||
[]byte(`{"cachedContent":"cachedContents/abc","contents":[]}`))
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "has to be the same")
|
||||
})
|
||||
|
||||
t.Run("composite groups are rejected", func(t *testing.T) {
|
||||
composite := geminiCachedContentTestAPIKey(service.PlatformComposite)
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", composite)
|
||||
_, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), composite, "gemini-3.8-flash",
|
||||
[]byte(`{"cachedContent":"cachedContents/abc","contents":[]}`))
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentsListAndGet(t *testing.T) {
|
||||
repo := &geminiCachedContentHandlerRepoStub{}
|
||||
seedGeminiCachedContent(repo, "first", 7, time.Now().Add(time.Hour))
|
||||
seedGeminiCachedContent(repo, "second", 7, time.Now().Add(time.Hour))
|
||||
seedGeminiCachedContent(repo, "expired", 7, time.Now().Add(-time.Minute))
|
||||
seedGeminiCachedContent(repo, "foreign", 8, time.Now().Add(time.Hour))
|
||||
h := newGeminiCachedContentHandlerForTest(repo)
|
||||
apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
|
||||
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents?pageSize=1", apiKey)
|
||||
h.GeminiCachedContentsList(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var page struct {
|
||||
CachedContents []map[string]any `json:"cachedContents"`
|
||||
NextPageToken string `json:"nextPageToken"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &page))
|
||||
require.Len(t, page.CachedContents, 1)
|
||||
require.Equal(t, "cachedContents/second", page.CachedContents[0]["name"])
|
||||
require.NotEmpty(t, page.NextPageToken)
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents?pageSize=1&pageToken="+page.NextPageToken, apiKey)
|
||||
h.GeminiCachedContentsList(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
page.NextPageToken = ""
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &page))
|
||||
require.Len(t, page.CachedContents, 1)
|
||||
require.Equal(t, "cachedContents/first", page.CachedContents[0]["name"])
|
||||
require.Empty(t, page.NextPageToken, "过期与其他 Key 的缓存不出现在列表里")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents/first", apiKey)
|
||||
c.Params = gin.Params{{Key: "cacheID", Value: "first"}}
|
||||
h.GeminiCachedContentsGet(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var view map[string]any
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &view))
|
||||
require.Equal(t, "cachedContents/first", view["name"])
|
||||
require.Equal(t, "models/gemini-3.8-flash", view["model"])
|
||||
require.NotContains(t, rec.Body.String(), "upstreamfirst", "上游资源名不暴露给客户端")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents/foreign", apiKey)
|
||||
c.Params = gin.Params{{Key: "cacheID", Value: "foreign"}}
|
||||
h.GeminiCachedContentsGet(c)
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.True(t, strings.Contains(rec.Body.String(), service.GeminiCachedContentNotFoundMessage))
|
||||
}
|
||||
|
||||
func TestHandleGeminiCachedContentFailoverExhausted(t *testing.T) {
|
||||
h := &GatewayHandler{}
|
||||
apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
|
||||
unavailable := `{"error":{"code":503,"message":"The account holding this cached content is temporarily unavailable","status":"UNAVAILABLE"}}`
|
||||
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusServiceUnavailable, ResponseBody: []byte(unavailable), BoundUpstreamPassthrough: true})
|
||||
require.Equal(t, http.StatusServiceUnavailable, rec.Code, "绑定缓存的请求无账号可换,上游网关的 503 原样返回")
|
||||
require.JSONEq(t, unavailable, rec.Body.String())
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(`{"error":{"code":429,"message":"Too many pending requests, please retry later","status":"RESOURCE_EXHAUSTED"}}`), BoundUpstreamPassthrough: true})
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "Too many pending requests")
|
||||
|
||||
officialQuota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.QuotaFailure","violations":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}]}}`
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(officialQuota)})
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code, "官方账号的限流按普通 Gemini 路径映射")
|
||||
require.NotContains(t, rec.Body.String(), "FreeTier", "不向客户端暴露上游账号的配额明细")
|
||||
require.Contains(t, rec.Body.String(), "Upstream rate limit exceeded")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusUnauthorized, ResponseBody: []byte(`{"error":{"code":401,"message":"bad key"}}`)})
|
||||
require.Equal(t, http.StatusBadGateway, rec.Code, "账号凭据类错误仍按网关映射,不透传给客户端")
|
||||
require.NotContains(t, rec.Body.String(), "bad key")
|
||||
}
|
||||
|
||||
func TestWriteGeminiCachedContentUpstreamError_RateLimitMappedUnlessUpstreamGateway(t *testing.T) {
|
||||
h := &GatewayHandler{}
|
||||
apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
|
||||
quota := `{"error":{"code":429,"message":"You exceeded your current quota","status":"RESOURCE_EXHAUSTED","details":[{"quotaId":"GenerateRequestsPerDayPerProjectPerModel-FreeTier"}]}}`
|
||||
|
||||
c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey)
|
||||
h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota)})
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
||||
require.NotContains(t, rec.Body.String(), "FreeTier", "官方账号的限流不透传配额明细")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey)
|
||||
h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusTooManyRequests, Body: []byte(quota), UpstreamGateway: true})
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
||||
require.JSONEq(t, quota, rec.Body.String(), "上游网关账号的限流反映持有缓存的上游账号状态,原样回写")
|
||||
|
||||
c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/cachedContents", apiKey)
|
||||
h.writeGeminiCachedContentUpstreamError(c, &service.GeminiCachedContentUpstreamError{StatusCode: http.StatusBadRequest, Body: []byte(`{"error":{"code":400,"message":"Cached content is too small","status":"INVALID_ARGUMENT"}}`)})
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "too small", "请求校验类错误仍原样回写")
|
||||
}
|
||||
@@ -74,7 +74,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
account, err := h.geminiCompatService.SelectAccountForModelsEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
if err != nil {
|
||||
if len(agModels) > 0 {
|
||||
c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(agModels)})
|
||||
@@ -85,7 +85,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models")
|
||||
res, err := h.geminiCompatService.ForwardListModels(c.Request.Context(), account)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadGateway, err.Error())
|
||||
return
|
||||
@@ -251,7 +251,7 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
account, err := h.geminiCompatService.SelectAccountForModelsEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
if err != nil {
|
||||
// 没有 gemini 账户,检查是否有 antigravity 账户可用
|
||||
hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
|
||||
@@ -265,7 +265,7 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
res, err := h.geminiCompatService.ForwardAIStudioGET(c.Request.Context(), account, "/v1beta/models/"+modelName)
|
||||
res, err := h.geminiCompatService.ForwardGetModel(c.Request.Context(), account, modelName)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusBadGateway, err.Error())
|
||||
return
|
||||
@@ -356,6 +356,12 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
modelName = channelMapping.MappedModel
|
||||
}
|
||||
|
||||
// 引用显式缓存的请求只能发往持有该缓存的账号,且不参与会话粘性与换号。
|
||||
cachedContentBinding, body, ok := h.resolveGeminiCachedContentBinding(c, reqLog, apiKey, modelName, body)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
// Get subscription (may be nil)
|
||||
subscription, _ := middleware.GetSubscriptionFromContext(c)
|
||||
|
||||
@@ -382,11 +388,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
// 2) billing eligibility check (after wait)
|
||||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
|
||||
reqLog.Info("gemini.billing_eligibility_check_failed", zap.Error(err))
|
||||
status, _, message, retryAfter := billingErrorDetails(err)
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleError(c, status, message)
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -394,19 +400,22 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.inflight_reservation_rejected", zap.Error(err))
|
||||
status, _, message, retryAfter := billingErrorDetails(err)
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleError(c, status, message)
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// 3) select account (sticky session based on request body)
|
||||
// 优先使用 Gemini CLI 的会话标识(privileged-user-id + tmp 目录哈希)
|
||||
sessionHash := extractGeminiCLISessionHash(c, body)
|
||||
if sessionHash == "" {
|
||||
sessionHash := ""
|
||||
if cachedContentBinding == nil {
|
||||
sessionHash = extractGeminiCLISessionHash(c, body)
|
||||
}
|
||||
if sessionHash == "" && cachedContentBinding == nil {
|
||||
// Fallback: 使用通用的会话哈希生成逻辑(适用于其他客户端)
|
||||
parsedReq, _ := service.ParseGatewayRequest(service.NewRequestBodyRef(body), domain.PlatformGemini)
|
||||
if parsedReq != nil {
|
||||
@@ -443,7 +452,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
var geminiPrefixHash string
|
||||
var geminiSessionUUID string
|
||||
var matchedDigestChain string
|
||||
useDigestFallback := sessionBoundAccountID == 0
|
||||
useDigestFallback := sessionBoundAccountID == 0 && cachedContentBinding == nil
|
||||
|
||||
if useDigestFallback {
|
||||
// 解析 Gemini 请求体
|
||||
@@ -507,7 +516,21 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0
|
||||
cleanedForUnknownBinding := false
|
||||
|
||||
fs := NewFailoverState(h.maxAccountSwitchesGemini, hasBoundSession)
|
||||
maxAccountSwitches := h.maxAccountSwitchesGemini
|
||||
var cachedContentExclusions map[int64]struct{}
|
||||
if cachedContentBinding != nil {
|
||||
// 缓存只存在于绑定账号:排除分组内其余账号,失败不换号。
|
||||
maxAccountSwitches = 0
|
||||
sessionBoundAccountID = cachedContentBinding.AccountID
|
||||
c.Request = c.Request.WithContext(service.WithGeminiCachedContentBound(c.Request.Context()))
|
||||
cachedContentExclusions, err = h.gatewayService.GeminiExplicitCacheExclusions(c.Request.Context(), apiKey.GroupID, cachedContentBinding.AccountID)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.cached_content.list_accounts_failed", zap.Error(err))
|
||||
googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage)
|
||||
return
|
||||
}
|
||||
}
|
||||
fs := NewFailoverState(maxAccountSwitches, hasBoundSession)
|
||||
|
||||
// 单账号分组提前设置 SingleAccountRetry 标记,让 Service 层首次 503 就不设模型限流标记。
|
||||
// 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流,导致后续请求连续快速失败。
|
||||
@@ -517,12 +540,32 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
}
|
||||
|
||||
for {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
excludedAccountIDs := fs.FailedAccountIDs
|
||||
if cachedContentExclusions != nil {
|
||||
excludedAccountIDs = mergeAccountIDSets(cachedContentExclusions, fs.FailedAccountIDs)
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, excludedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
if err == nil && cachedContentBinding != nil && selection.Account.ID != cachedContentBinding.AccountID {
|
||||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
err = service.ErrNoAvailableAccounts
|
||||
}
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gemini.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if cachedContentBinding != nil {
|
||||
if fs.LastFailoverErr != nil {
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, fs.LastFailoverErr)
|
||||
return
|
||||
}
|
||||
reqLog.Info("gemini.cached_content.bound_account_unavailable", zap.Int64("account_id", cachedContentBinding.AccountID), zap.Error(err))
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage)
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, service.PlatformGemini)
|
||||
if !cls.ModelNotFound {
|
||||
@@ -631,6 +674,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
reqLog.Debug("gemini.account_slot_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason))
|
||||
if cachedContentBinding != nil {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
googleError(c, http.StatusServiceUnavailable, geminiCachedContentUnavailableMessage)
|
||||
return
|
||||
}
|
||||
if fs.RecordProfitVeto(account.ID) == FailoverExhausted {
|
||||
reqLog.Warn("gemini.profit_veto_attempts_exhausted", zap.Int("profit_veto_count", fs.ProfitVetoCount()))
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
@@ -684,6 +732,10 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
case FailoverContinue:
|
||||
continue
|
||||
case FailoverExhausted:
|
||||
if cachedContentBinding != nil {
|
||||
h.handleGeminiCachedContentFailoverExhausted(c, fs.LastFailoverErr)
|
||||
return
|
||||
}
|
||||
h.handleGeminiFailoverExhausted(c, fs.LastFailoverErr)
|
||||
return
|
||||
case FailoverCanceled:
|
||||
@@ -809,7 +861,8 @@ func (h *GatewayHandler) handleGeminiFailoverExhausted(c *gin.Context, failoverE
|
||||
c.Set(service.OpsSkipPassthroughKey, true)
|
||||
}
|
||||
|
||||
googleError(c, respCode, msg)
|
||||
// 按透传规则改写的是上游错误,不登记本地错误类型。
|
||||
googleErrorWithType(c, respCode, "", "", msg)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -844,7 +897,21 @@ type pathParseError struct{ msg string }
|
||||
|
||||
func (e *pathParseError) Error() string { return e.msg }
|
||||
|
||||
// googleError 以 Google 错误格式回写网关产生的错误,并按状态码向 ops 登记错误类型:
|
||||
// 4xx 为客户端请求问题,其余状态码不登记(按平台错误统计)。
|
||||
func googleError(c *gin.Context, status int, message string) {
|
||||
googleErrorWithType(c, status, googleClientErrorType(status), "", message)
|
||||
}
|
||||
|
||||
// googlePlatformError 以 Google 错误格式回写由平台侧原因(如管理员未完成配置)导致的错误,ops 按平台错误统计。
|
||||
func googlePlatformError(c *gin.Context, status int, message string) {
|
||||
googleErrorWithType(c, status, "", "", message)
|
||||
}
|
||||
|
||||
// googleErrorWithType 以 Google 错误格式回写错误,并向 ops 登记错误类型与错误码(Google 错误体没有 type 字段),
|
||||
// 与 Claude / OpenAI 格式错误同口径分类;errType 为空时不登记。
|
||||
func googleErrorWithType(c *gin.Context, status int, errType, code, message string) {
|
||||
service.SetOpsLocalErrorType(c, errType, code)
|
||||
c.JSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"code": status,
|
||||
@@ -854,11 +921,27 @@ func googleError(c *gin.Context, status int, message string) {
|
||||
})
|
||||
}
|
||||
|
||||
// googleClientErrorType 返回客户端请求类 4xx 状态码对应的 ops 错误类型;其余状态码返回空串。
|
||||
func googleClientErrorType(status int) string {
|
||||
switch status {
|
||||
case http.StatusBadRequest, http.StatusRequestEntityTooLarge:
|
||||
return "invalid_request_error"
|
||||
case http.StatusUnauthorized:
|
||||
return "authentication_error"
|
||||
case http.StatusForbidden:
|
||||
return "permission_error"
|
||||
case http.StatusNotFound:
|
||||
return "not_found_error"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// googleConcurrencyError 以 Google 错误格式回写并发槽获取失败,状态码与文案
|
||||
// 沿用 concurrencyErrorResponse 的统一映射(客户端断开为 499)。
|
||||
func googleConcurrencyError(c *gin.Context, err error, slotType string) {
|
||||
status, _, _, message := concurrencyErrorResponse(err, slotType)
|
||||
googleError(c, status, message)
|
||||
status, errType, code, message := concurrencyErrorResponse(err, slotType)
|
||||
googleErrorWithType(c, status, errType, code, message)
|
||||
}
|
||||
|
||||
func writeUpstreamResponse(c *gin.Context, res *service.UpstreamHTTPResult) {
|
||||
|
||||
@@ -109,7 +109,7 @@ func TestGeminiV1BetaHandler_ListModelsAntigravityFallback(t *testing.T) {
|
||||
expectedBehavior string
|
||||
}{
|
||||
{
|
||||
name: "有Gemini账户-调用ForwardAIStudioGET",
|
||||
name: "有Gemini账户-调用ForwardListModels",
|
||||
hasGeminiAccount: true,
|
||||
hasAntigravity: false,
|
||||
expectedBehavior: "forward_to_upstream",
|
||||
@@ -155,7 +155,7 @@ func TestGeminiV1BetaHandler_GetModelAntigravityFallback(t *testing.T) {
|
||||
expectedBehavior string
|
||||
}{
|
||||
{
|
||||
name: "有Gemini账户-调用ForwardAIStudioGET",
|
||||
name: "有Gemini账户-调用ForwardGetModel",
|
||||
hasGeminiAccount: true,
|
||||
hasAntigravity: false,
|
||||
expectedBehavior: "forward_to_upstream",
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"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 geminiVertexModelsTokenCache struct{}
|
||||
|
||||
func (geminiVertexModelsTokenCache) GetAccessToken(context.Context, string) (string, error) {
|
||||
return "ya29.vertex-token", nil
|
||||
}
|
||||
|
||||
func (geminiVertexModelsTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (geminiVertexModelsTokenCache) DeleteAccessToken(context.Context, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (geminiVertexModelsTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (geminiVertexModelsTokenCache) ReleaseRefreshLock(context.Context, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 分组里只有 Agent Platform (Vertex) 服务账号时,模型列表与模型详情接口返回 Vertex 上的 Gemini 模型。
|
||||
func TestGeminiV1BetaModelsEndpoints_VertexServiceAccountOnlyGroup(t *testing.T) {
|
||||
groupID := int64(8)
|
||||
repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{groupID: {{
|
||||
ID: 18,
|
||||
Platform: service.PlatformGemini,
|
||||
Type: service.AccountTypeServiceAccount,
|
||||
Credentials: map[string]any{
|
||||
"location": "global",
|
||||
"service_account_json": map[string]any{
|
||||
"type": "service_account",
|
||||
"project_id": "proj",
|
||||
"private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n",
|
||||
"client_email": "svc@proj.iam.gserviceaccount.com",
|
||||
},
|
||||
},
|
||||
}}}}}
|
||||
upstream := &geminiMixedModelsUpstream{status: http.StatusOK, body: `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}`}
|
||||
tokenProvider := service.NewGeminiTokenProvider(nil, geminiVertexModelsTokenCache{}, nil)
|
||||
h := &GatewayHandler{geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, tokenProvider, nil, upstream, nil, &config.Config{})}
|
||||
apiKey := &service.APIKey{GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformGemini}}
|
||||
|
||||
t.Run("list", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
|
||||
h.GeminiV1BetaListModels(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.JSONEq(t, `{"models":[{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}]}`, rec.Body.String())
|
||||
})
|
||||
|
||||
t.Run("get", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-3.8-flash", nil)
|
||||
c.Params = gin.Params{{Key: "model", Value: "gemini-3.8-flash"}}
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
|
||||
h.GeminiV1BetaGetModel(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.JSONEq(t, `{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}`, rec.Body.String())
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestToLocalizedModelPlazaGroups(t *testing.T) {
|
||||
visible := []service.PlazaGroup{
|
||||
{
|
||||
ID: 1, Name: "标准", Description: "号池", Platform: "anthropic",
|
||||
I18n: service.GroupI18n{"en": {Name: "Standard", Description: "Pro + Max pool"}},
|
||||
},
|
||||
{ID: 2, Name: "无译文", Description: "原描述", Platform: "openai"},
|
||||
}
|
||||
|
||||
english := toLocalizedModelPlazaGroups(visible, nil, "en")
|
||||
require.Equal(t, "Standard", english[0].Name)
|
||||
require.Equal(t, "Pro + Max pool", english[0].Description)
|
||||
require.Equal(t, "无译文", english[1].Name)
|
||||
require.Equal(t, "原描述", english[1].Description)
|
||||
|
||||
chinese := toLocalizedModelPlazaGroups(visible, nil, "zh")
|
||||
require.Equal(t, "标准", chinese[0].Name)
|
||||
require.Equal(t, "号池", chinese[0].Description)
|
||||
|
||||
// service 层对象不被改写,同一批分组可继续按别的语言输出
|
||||
require.Equal(t, "标准", visible[0].Name)
|
||||
|
||||
raw, err := json.Marshal(english[0])
|
||||
require.NoError(t, err)
|
||||
var decoded map[string]any
|
||||
require.NoError(t, json.Unmarshal(raw, &decoded))
|
||||
require.NotContains(t, decoded, "i18n")
|
||||
}
|
||||
|
||||
func TestLocalizeUserAvailableGroups(t *testing.T) {
|
||||
refs := []service.AvailableGroupRef{
|
||||
{ID: 1, Name: "标准", Platform: "anthropic", I18n: service.GroupI18n{"en": {Name: "Standard"}}},
|
||||
{ID: 2, Name: "无译文", Platform: "anthropic"},
|
||||
}
|
||||
allowed := map[int64]struct{}{1: {}, 2: {}}
|
||||
|
||||
english := filterUserVisibleGroups(refs, allowed)
|
||||
localizeUserAvailableGroups(english, "en")
|
||||
require.Equal(t, "Standard", english[0].Name)
|
||||
require.Equal(t, "无译文", english[1].Name)
|
||||
|
||||
chinese := filterUserVisibleGroups(refs, allowed)
|
||||
localizeUserAvailableGroups(chinese, "zh")
|
||||
require.Equal(t, "标准", chinese[0].Name)
|
||||
|
||||
raw, err := json.Marshal(english[0])
|
||||
require.NoError(t, err)
|
||||
var decoded map[string]any
|
||||
require.NoError(t, json.Unmarshal(raw, &decoded))
|
||||
require.NotContains(t, decoded, "i18n")
|
||||
}
|
||||
|
||||
// 用户侧 handler 输出带分组的 DTO 时必须经过 Localized,否则分组名称与描述会
|
||||
// 停留在原字段,并把译文表原样带给用户。管理端(handler/admin)不在此列。
|
||||
func TestUserFacingGroupDTOsAreLocalized(t *testing.T) {
|
||||
mapperCall := regexp.MustCompile(`dto\.(APIKey|UserSubscription|RedeemCode|UsageLog|Group)FromService[A-Za-z]*\(`)
|
||||
|
||||
files, err := filepath.Glob("*.go")
|
||||
require.NoError(t, err)
|
||||
|
||||
checked := 0
|
||||
for _, file := range files {
|
||||
if strings.HasSuffix(file, "_test.go") {
|
||||
continue
|
||||
}
|
||||
source, err := os.ReadFile(file)
|
||||
require.NoError(t, err)
|
||||
for i, line := range strings.Split(string(source), "\n") {
|
||||
if !mapperCall.MatchString(line) {
|
||||
continue
|
||||
}
|
||||
checked++
|
||||
require.Containsf(t, line, ".Localized(",
|
||||
"%s:%d outputs a group DTO without Localized(Accept-Language)", file, i+1)
|
||||
}
|
||||
}
|
||||
require.NotZero(t, checked, "expected to find user-facing group DTO call sites")
|
||||
}
|
||||
@@ -151,10 +151,10 @@ func (h *ModelPlazaHandler) Get(c *gin.Context) {
|
||||
|
||||
visible := filterPlazaVisibleGroups(groups, allowedGroups, restrictPublicGroups)
|
||||
|
||||
out := make([]modelPlazaGroup, 0, len(visible))
|
||||
for i := range visible {
|
||||
out = append(out, toModelPlazaGroupDTO(&visible[i], userRates))
|
||||
}
|
||||
// 广场可匿名访问,分组文案随请求语言变化,提示中间缓存按语言区分。
|
||||
// 追加而非覆盖:CORS 中间件已为具体来源写入 Vary: Origin。
|
||||
c.Writer.Header().Add("Vary", "Accept-Language")
|
||||
out := toLocalizedModelPlazaGroups(visible, userRates, c.GetHeader("Accept-Language"))
|
||||
response.Success(c, modelPlazaResponse{
|
||||
Description: rt.Description,
|
||||
Groups: out,
|
||||
@@ -185,6 +185,17 @@ func filterPlazaVisibleGroups(
|
||||
return visible
|
||||
}
|
||||
|
||||
// toLocalizedModelPlazaGroups 转换可见分组,并把名称与描述换成请求首选语言的版本。
|
||||
func toLocalizedModelPlazaGroups(visible []service.PlazaGroup, userRates map[int64]float64, acceptLanguage string) []modelPlazaGroup {
|
||||
out := make([]modelPlazaGroup, 0, len(visible))
|
||||
for i := range visible {
|
||||
group := toModelPlazaGroupDTO(&visible[i], userRates)
|
||||
group.Name, group.Description = visible[i].I18n.Localize(acceptLanguage, group.Name, group.Description)
|
||||
out = append(out, group)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// toModelPlazaGroupDTO 将 service 层广场分组映射为白名单 DTO,并合并用户专属倍率。
|
||||
func toModelPlazaGroupDTO(g *service.PlazaGroup, userRates map[int64]float64) modelPlazaGroup {
|
||||
models := make([]modelPlazaModel, 0, len(g.Models))
|
||||
|
||||
@@ -3,11 +3,14 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -262,3 +265,43 @@ func TestFilterPlazaVisibleGroups_SubscribedExclusiveGroup(t *testing.T) {
|
||||
require.Equal(t, int64(42), visible[0].ID)
|
||||
}
|
||||
}
|
||||
|
||||
type plazaVarySettingRepo struct{ service.SettingRepository }
|
||||
|
||||
func (plazaVarySettingRepo) GetMultiple(context.Context, []string) (map[string]string, error) {
|
||||
return map[string]string{service.SettingKeyModelPlazaEnabled: "true"}, nil
|
||||
}
|
||||
|
||||
type plazaVaryChannelRepo struct{ service.ChannelRepository }
|
||||
|
||||
func (plazaVaryChannelRepo) ListAll(context.Context) ([]service.Channel, error) { return nil, nil }
|
||||
|
||||
type plazaVaryGroupRepo struct{ service.GroupRepository }
|
||||
|
||||
func (plazaVaryGroupRepo) ListActive(context.Context) ([]service.Group, error) { return nil, nil }
|
||||
|
||||
func TestModelPlazaHandler_VaryKeepsCORSOrigin(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
h := NewModelPlazaHandler(
|
||||
service.NewModelPlazaService(plazaVaryChannelRepo{}, plazaVaryGroupRepo{}, nil, nil, nil),
|
||||
nil,
|
||||
service.NewSettingService(plazaVarySettingRepo{}, &config.Config{}),
|
||||
)
|
||||
r := gin.New()
|
||||
r.Use(middleware.CORS(config.CORSConfig{AllowedOrigins: []string{"https://a.example", "https://b.example"}}))
|
||||
r.GET("/api/v1/model-plaza", h.Get)
|
||||
|
||||
for _, origin := range []string{"https://a.example", "https://b.example"} {
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/model-plaza", nil)
|
||||
req.Header.Set("Origin", origin)
|
||||
req.Header.Set("Accept-Language", "en")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code, w.Body.String())
|
||||
require.Equal(t, origin, w.Header().Get("Access-Control-Allow-Origin"))
|
||||
// 共享缓存需同时按来源与语言区分,否则 A 来源的响应可能被复用给 B。
|
||||
require.Equal(t, []string{"Origin", "Accept-Language"}, w.Header().Values("Vary"))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1183,6 +1183,14 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
// Google 格式错误体没有 type 字段:按网关回写时登记的错误类型补齐。
|
||||
if errType, code := service.OpsLocalErrorType(c); errType != "" && !parsed.StreamFailure &&
|
||||
(parsed.ErrorType == "" || parsed.ErrorType == "api_error") {
|
||||
parsed.ErrorType = errType
|
||||
if code != "" {
|
||||
parsed.Code = code
|
||||
}
|
||||
}
|
||||
normalizedType := normalizeOpsErrorType(parsed.ErrorType, parsed.Code)
|
||||
|
||||
phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, parsed.Message, parsed.Code, status)
|
||||
@@ -2214,6 +2222,11 @@ func classifyOpsSeverity(errType string, status int) string {
|
||||
|
||||
func classifyOpsErrorLog(c *gin.Context, errType, message, code string, status int) (phase string, isBusinessLimited bool, errorOwner string, errorSource string) {
|
||||
phase = classifyOpsPhase(errType, message, code)
|
||||
if service.IsOpsRequestScopedError(c) {
|
||||
// 失败由请求内容决定:只按错误类型分类,上游错误上下文不参与归因。
|
||||
isBusinessLimited = classifyOpsIsBusinessLimited(errType, phase, code, status, message)
|
||||
return phase, isBusinessLimited, classifyOpsErrorOwner(phase, message), classifyOpsErrorSource(phase, message)
|
||||
}
|
||||
routingCapacityLimited := isOpsRoutingCapacityLimited(c)
|
||||
clientBusinessLimited := service.HasOpsClientBusinessLimited(c)
|
||||
localModelConfiguration := clientBusinessLimited && service.OpsClientBusinessLimitedReason(c) == service.OpsClientBusinessLimitedReasonLocalModelConfiguration
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// recordGoogleFormatOpsError 让 handle 以 Google 错误格式回写错误,返回 ops 中间件记录的错误日志。
|
||||
func recordGoogleFormatOpsError(t *testing.T, handle gin.HandlerFunc) *service.OpsInsertErrorLogInput {
|
||||
t.Helper()
|
||||
setupOpsErrorLogTestQueue(t, 1)
|
||||
gin.SetMode(gin.TestMode)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
router := gin.New()
|
||||
router.Use(OpsErrorLoggerMiddleware(ops))
|
||||
router.POST("/v1beta/cachedContents", handle)
|
||||
|
||||
router.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1beta/cachedContents", nil))
|
||||
|
||||
require.Equal(t, int64(1), OpsErrorLogQueueLength())
|
||||
return (<-opsErrorLogQueue).entry
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleClientErrorsClassifiedAsClient(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
status int
|
||||
errType string
|
||||
phase string
|
||||
}{
|
||||
{status: http.StatusBadRequest, errType: "invalid_request_error", phase: "request"},
|
||||
{status: http.StatusRequestEntityTooLarge, errType: "invalid_request_error", phase: "request"},
|
||||
{status: http.StatusUnauthorized, errType: "authentication_error", phase: "auth"},
|
||||
{status: http.StatusForbidden, errType: "permission_error", phase: "request"},
|
||||
{status: http.StatusNotFound, errType: "not_found_error", phase: "request"},
|
||||
} {
|
||||
t.Run(http.StatusText(tc.status), func(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
googleError(c, tc.status, "client request problem")
|
||||
})
|
||||
require.Equal(t, tc.status, entry.StatusCode)
|
||||
require.Equal(t, tc.errType, entry.ErrorType)
|
||||
require.Equal(t, tc.phase, entry.ErrorPhase)
|
||||
require.Equal(t, "client", entry.ErrorOwner)
|
||||
require.Equal(t, "client_request", entry.ErrorSource)
|
||||
require.Equal(t, "P3", entry.Severity)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GooglePlatformErrorStaysPlatform(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
googlePlatformError(c, http.StatusBadRequest, "Explicit context caching is not available for model m: storage pricing is not configured")
|
||||
})
|
||||
require.Equal(t, "api_error", entry.ErrorType)
|
||||
require.Equal(t, "internal", entry.ErrorPhase)
|
||||
require.Equal(t, "platform", entry.ErrorOwner)
|
||||
require.Equal(t, "gateway", entry.ErrorSource)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleServerErrorsNotTypedByStatus(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
googleError(c, http.StatusInternalServerError, "Failed to persist cached content")
|
||||
})
|
||||
require.Equal(t, "api_error", entry.ErrorType)
|
||||
require.Equal(t, "platform", entry.ErrorOwner)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleBillingErrorMatchesClaudeFormat(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
status, code, message, _ := billingErrorDetails(service.ErrInsufficientBalance)
|
||||
googleErrorWithType(c, status, code, "", message)
|
||||
})
|
||||
require.Equal(t, "billing_error", entry.ErrorType)
|
||||
require.Equal(t, "request", entry.ErrorPhase)
|
||||
require.Equal(t, "client", entry.ErrorOwner)
|
||||
require.True(t, entry.IsBusinessLimited)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleConcurrencyErrorClassifiedAsRequest(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
googleConcurrencyError(c, &ConcurrencyError{SlotType: "user"}, "user")
|
||||
})
|
||||
require.Equal(t, http.StatusTooManyRequests, entry.StatusCode)
|
||||
require.Equal(t, "rate_limit_error", entry.ErrorType)
|
||||
require.Equal(t, "request", entry.ErrorPhase)
|
||||
require.Equal(t, "client", entry.ErrorOwner)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_GoogleErrorAfterUpstreamFailureKeepsUpstreamAttribution(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
service.SetOpsUpstreamError(c, http.StatusBadRequest, "upstream rejected", "")
|
||||
googleError(c, http.StatusBadRequest, "upstream rejected")
|
||||
})
|
||||
require.Equal(t, "upstream", entry.ErrorPhase)
|
||||
require.Equal(t, "provider", entry.ErrorOwner)
|
||||
require.Equal(t, "upstream_http", entry.ErrorSource)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_RequestScopedUpstreamErrorClassifiedAsClient(t *testing.T) {
|
||||
entry := recordGoogleFormatOpsError(t, func(c *gin.Context) {
|
||||
service.SetOpsUpstreamError(c, http.StatusBadRequest, "Invalid resource state for cache content 8079767076522164224.", "")
|
||||
c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{{
|
||||
Platform: service.PlatformGemini, AccountID: 173, UpstreamStatusCode: http.StatusBadRequest, Kind: "http_error",
|
||||
Message: "Invalid resource state for cache content 8079767076522164224.",
|
||||
}})
|
||||
service.MarkOpsRequestScopedError(c, "permission_error")
|
||||
c.Data(http.StatusForbidden, "application/json", []byte(service.GeminiCachedContentNotFoundResponse))
|
||||
})
|
||||
require.Equal(t, http.StatusForbidden, entry.StatusCode)
|
||||
require.Equal(t, "permission_error", entry.ErrorType)
|
||||
require.Equal(t, "request", entry.ErrorPhase)
|
||||
require.Equal(t, "client", entry.ErrorOwner)
|
||||
require.Equal(t, "client_request", entry.ErrorSource)
|
||||
require.False(t, entry.IsBusinessLimited)
|
||||
require.NotNil(t, entry.UpstreamStatusCode, "上游原始错误保留为明细")
|
||||
require.Equal(t, http.StatusBadRequest, *entry.UpstreamStatusCode)
|
||||
}
|
||||
|
||||
func TestSetOpsLocalErrorTypeClearsRequestScopedMark(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
service.MarkOpsRequestScopedError(c, "permission_error")
|
||||
require.True(t, service.IsOpsRequestScopedError(c))
|
||||
|
||||
service.SetOpsLocalErrorType(c, "", "")
|
||||
require.False(t, service.IsOpsRequestScopedError(c))
|
||||
errType, code := service.OpsLocalErrorType(c)
|
||||
require.Empty(t, errType)
|
||||
require.Empty(t, code)
|
||||
}
|
||||
@@ -78,7 +78,7 @@ func (h *PaymentHandler) GetPlans(c *gin.Context) {
|
||||
gi := groupInfo[p.GroupID]
|
||||
result = append(result, planWithPlatform{
|
||||
ID: int64(p.ID), GroupID: p.GroupID,
|
||||
GroupPlatform: gi.Platform, GroupName: gi.Name,
|
||||
GroupPlatform: gi.Platform, GroupName: gi.I18n.LocalizeName(c.GetHeader("Accept-Language"), gi.Name),
|
||||
RateMultiplier: gi.RateMultiplier, PeakRateEnabled: gi.PeakRateEnabled,
|
||||
PeakStart: gi.PeakStart, PeakEnd: gi.PeakEnd, PeakRateMultiplier: gi.PeakRateMultiplier,
|
||||
Name: p.Name, Description: p.Description, Price: p.Price, OriginalPrice: p.OriginalPrice,
|
||||
@@ -126,7 +126,7 @@ func (h *PaymentHandler) GetCheckoutInfo(c *gin.Context) {
|
||||
gi := groupInfo[p.GroupID]
|
||||
planList = append(planList, checkoutPlan{
|
||||
ID: int64(p.ID), GroupID: p.GroupID,
|
||||
GroupPlatform: gi.Platform, GroupName: gi.Name,
|
||||
GroupPlatform: gi.Platform, GroupName: gi.I18n.LocalizeName(c.GetHeader("Accept-Language"), gi.Name),
|
||||
RateMultiplier: gi.RateMultiplier,
|
||||
PeakRateEnabled: gi.PeakRateEnabled, PeakStart: gi.PeakStart,
|
||||
PeakEnd: gi.PeakEnd, PeakRateMultiplier: gi.PeakRateMultiplier,
|
||||
|
||||
@@ -59,7 +59,7 @@ func (h *RedeemHandler) Redeem(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, dto.RedeemCodeFromService(result))
|
||||
response.Success(c, dto.RedeemCodeFromService(result).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
|
||||
// GetHistory returns the user's redemption history
|
||||
@@ -98,7 +98,7 @@ func (h *RedeemHandler) GetHistory(c *gin.Context) {
|
||||
}
|
||||
out := make([]dto.RedeemCode, 0, len(codes))
|
||||
for i := range codes {
|
||||
out = append(out, *dto.RedeemCodeFromService(&codes[i]))
|
||||
out = append(out, *dto.RedeemCodeFromService(&codes[i]).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
response.Paginated(c, out, result.Total, page, pageSize)
|
||||
return
|
||||
@@ -115,7 +115,7 @@ func (h *RedeemHandler) GetHistory(c *gin.Context) {
|
||||
|
||||
out := make([]dto.RedeemCode, 0, len(codes))
|
||||
for i := range codes {
|
||||
out = append(out, *dto.RedeemCodeFromService(&codes[i]))
|
||||
out = append(out, *dto.RedeemCodeFromService(&codes[i]).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
response.Success(c, out)
|
||||
}
|
||||
|
||||
@@ -100,9 +100,14 @@ func googleSecurityAuditError(c *gin.Context, decision *securityaudit.Decision)
|
||||
return
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
googleError(c, securityAuditStatus(decision), securityAuditMessage(decision))
|
||||
googleErrorWithType(c, securityAuditStatus(decision), securityAuditErrorCode(decision), "", securityAuditMessage(decision))
|
||||
return
|
||||
}
|
||||
errType := "api_error"
|
||||
if decision.Kind == securityaudit.DecisionBlock {
|
||||
errType = "permission_error"
|
||||
}
|
||||
service.SetOpsLocalErrorType(c, errType, securityAuditErrorCode(decision))
|
||||
status := securityAuditStatus(decision)
|
||||
googleStatus := googleapi.HTTPStatusToGoogleStatus(status)
|
||||
if status == http.StatusServiceUnavailable {
|
||||
|
||||
@@ -24,6 +24,7 @@ func TestPromptAuditGatePrecedesAccountBillingAndUpstreamSideEffects(t *testing.
|
||||
{file: "gateway_handler_chat_completions.go", function: "ChatCompletions", auditToken: "checkSecurityAudit"},
|
||||
{file: "gateway_handler_responses.go", function: "Responses", auditToken: "checkSecurityAudit"},
|
||||
{file: "gemini_v1beta_handler.go", function: "GeminiV1BetaModels", auditToken: "checkSecurityAudit"},
|
||||
{file: "gemini_cached_contents_handler.go", function: "GeminiCachedContentsCreate", auditToken: "checkSecurityAudit"},
|
||||
{file: "openai_gateway_handler.go", function: "Responses", auditToken: "checkSecurityAudit"},
|
||||
{file: "openai_gateway_handler.go", function: "Messages", auditToken: "checkSecurityAudit"},
|
||||
{file: "openai_chat_completions.go", function: "ChatCompletions", auditToken: "checkSecurityAudit"},
|
||||
|
||||
@@ -59,7 +59,7 @@ func (h *SubscriptionHandler) List(c *gin.Context) {
|
||||
|
||||
out := make([]dto.UserSubscription, 0, len(subscriptions))
|
||||
for i := range subscriptions {
|
||||
out = append(out, *dto.UserSubscriptionFromService(&subscriptions[i]))
|
||||
out = append(out, *dto.UserSubscriptionFromService(&subscriptions[i]).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
response.Success(c, out)
|
||||
}
|
||||
@@ -81,7 +81,7 @@ func (h *SubscriptionHandler) GetActive(c *gin.Context) {
|
||||
|
||||
out := make([]dto.UserSubscription, 0, len(subscriptions))
|
||||
for i := range subscriptions {
|
||||
out = append(out, *dto.UserSubscriptionFromService(&subscriptions[i]))
|
||||
out = append(out, *dto.UserSubscriptionFromService(&subscriptions[i]).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
response.Success(c, out)
|
||||
}
|
||||
@@ -110,8 +110,11 @@ func (h *SubscriptionHandler) GetProgress(c *gin.Context) {
|
||||
// Skip subscriptions with errors
|
||||
continue
|
||||
}
|
||||
if sub.Group != nil {
|
||||
progress.GroupName = sub.Group.I18n.LocalizeName(c.GetHeader("Accept-Language"), progress.GroupName)
|
||||
}
|
||||
result = append(result, SubscriptionProgressInfo{
|
||||
Subscription: dto.UserSubscriptionFromService(sub),
|
||||
Subscription: dto.UserSubscriptionFromService(sub).Localized(c.GetHeader("Accept-Language")),
|
||||
Progress: progress,
|
||||
})
|
||||
}
|
||||
@@ -150,7 +153,7 @@ func (h *SubscriptionHandler) GetSummary(c *gin.Context) {
|
||||
|
||||
// Add group info if preloaded
|
||||
if sub.Group != nil {
|
||||
item.GroupName = sub.Group.Name
|
||||
item.GroupName = sub.Group.I18n.LocalizeName(c.GetHeader("Accept-Language"), sub.Group.Name)
|
||||
if sub.Group.DailyLimitUSD != nil {
|
||||
item.DailyLimitUSD = *sub.Group.DailyLimitUSD
|
||||
}
|
||||
|
||||
@@ -254,7 +254,7 @@ func (h *UsageHandler) List(c *gin.Context) {
|
||||
|
||||
out := make([]dto.UsageLog, 0, len(records))
|
||||
for i := range records {
|
||||
out = append(out, *dto.UsageLogFromService(&records[i]))
|
||||
out = append(out, *dto.UsageLogFromService(&records[i]).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
response.Paginated(c, out, result.Total, page, pageSize)
|
||||
}
|
||||
@@ -400,7 +400,7 @@ func (h *UsageHandler) GetByID(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, dto.UsageLogFromService(record))
|
||||
response.Success(c, dto.UsageLogFromService(record).Localized(c.GetHeader("Accept-Language")))
|
||||
}
|
||||
|
||||
// Stats handles getting usage statistics
|
||||
|
||||
@@ -113,11 +113,13 @@ func ProvideGatewayHandler(
|
||||
cfg *config.Config,
|
||||
settingService *service.SettingService,
|
||||
coordinator *securityaudit.Coordinator,
|
||||
geminiCachedContentService *service.GeminiCachedContentService,
|
||||
) *GatewayHandler {
|
||||
h := NewGatewayHandler(gatewayService, openAIGatewayService, geminiCompatService, antigravityGatewayService,
|
||||
userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool,
|
||||
errorPassthroughService, contentModerationService, userMsgQueueService, cfg, settingService)
|
||||
h.securityAuditCoordinator = coordinator
|
||||
h.geminiCachedContentService = geminiCachedContentService
|
||||
return h
|
||||
}
|
||||
|
||||
|
||||
@@ -1940,6 +1940,72 @@ func TestSonnet55SignedThinkingResponsesRoundTrip(t *testing.T) {
|
||||
require.Equal(t, "tool_use", blocks[2].Type)
|
||||
}
|
||||
|
||||
func TestHaiku55ResponsesThinkingAndSampling(t *testing.T) {
|
||||
for _, effort := range []string{"", "low", "medium", "high", "xhigh", "max", "none"} {
|
||||
req := &ResponsesRequest{Model: "claude-haiku-5-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: effort}}
|
||||
out, err := ResponsesToAnthropicRequest(req)
|
||||
require.NoError(t, err, effort)
|
||||
require.Zero(t, out.Thinking.BudgetTokens)
|
||||
if effort == "none" {
|
||||
require.Equal(t, "disabled", out.Thinking.Type)
|
||||
require.Equal(t, "low", out.OutputConfig.Effort)
|
||||
} else {
|
||||
require.Equal(t, "adaptive", out.Thinking.Type)
|
||||
if effort == "" {
|
||||
effort = "medium"
|
||||
}
|
||||
require.Equal(t, effort, out.OutputConfig.Effort)
|
||||
}
|
||||
}
|
||||
// Haiku 5.5 accepts forced tool_choice, unlike opus/sonnet 5.5.
|
||||
for _, choice := range []string{`"required"`, `{"type":"function","name":"lookup"}`} {
|
||||
out, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-haiku-5-5", Input: json.RawMessage(`"hello"`), ToolChoice: json.RawMessage(choice)})
|
||||
require.NoError(t, err, choice)
|
||||
require.NotEmpty(t, out.ToolChoice)
|
||||
}
|
||||
temperature := 0.7
|
||||
_, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-haiku-5-5", Input: json.RawMessage(`"hello"`), Temperature: &temperature})
|
||||
require.ErrorContains(t, err, "temperature")
|
||||
for _, topP := range []float64{0.5, 0.995, 1} {
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-haiku-5-5", Input: json.RawMessage(`"hello"`), TopP: &topP})
|
||||
require.ErrorContains(t, err, "top_p", "top_p=%v", topP)
|
||||
}
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-haiku-5-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: "minimal"}})
|
||||
require.ErrorContains(t, err, "reasoning effort")
|
||||
|
||||
temperature = 1
|
||||
topP := 0.99
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-haiku-5-5", Input: json.RawMessage(`"hello"`), Temperature: &temperature})
|
||||
require.NoError(t, err)
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-haiku-5-5", Input: json.RawMessage(`"hello"`), TopP: &topP})
|
||||
require.NoError(t, err)
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-haiku-5-5", Input: json.RawMessage(`"hello"`), Temperature: &temperature, TopP: &topP})
|
||||
require.ErrorContains(t, err, "together")
|
||||
}
|
||||
|
||||
func TestHaiku55SignedThinkingResponsesRoundTrip(t *testing.T) {
|
||||
block := AnthropicContentBlock{Type: "thinking", Thinking: "", Signature: "signed-haiku-block"}
|
||||
response := AnthropicToResponsesResponse(&AnthropicResponse{Model: "claude-haiku-5-5", Content: []AnthropicContentBlock{block, {Type: "text", Text: "progress"}, {Type: "tool_use", ID: "toolu_1", Name: "lookup", Input: json.RawMessage(`{}`)}}})
|
||||
require.Len(t, response.Output, 3)
|
||||
require.NotEmpty(t, response.Output[0].EncryptedContent)
|
||||
require.Equal(t, "message", response.Output[1].Type)
|
||||
raw, err := json.Marshal(response.Output)
|
||||
require.NoError(t, err)
|
||||
var items []ResponsesInputItem
|
||||
require.NoError(t, json.Unmarshal(raw, &items))
|
||||
items = append(items, ResponsesInputItem{Type: "function_call_output", CallID: response.Output[2].CallID, Output: "ok"})
|
||||
raw, err = json.Marshal(items)
|
||||
require.NoError(t, err)
|
||||
converted, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-haiku-5-5", Input: raw})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, converted.Messages, 2)
|
||||
var blocks []AnthropicContentBlock
|
||||
require.NoError(t, json.Unmarshal(converted.Messages[0].Content, &blocks))
|
||||
require.Equal(t, block, blocks[0])
|
||||
require.Equal(t, "text", blocks[1].Type)
|
||||
require.Equal(t, "tool_use", blocks[2].Type)
|
||||
}
|
||||
|
||||
func TestGPT6ChatSamplingAndCacheFields(t *testing.T) {
|
||||
temperature := 0.7
|
||||
for _, model := range []string{"gpt-6-sol", "gpt-6-luna"} {
|
||||
|
||||
@@ -53,7 +53,7 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
|
||||
for _, block := range resp.Content {
|
||||
switch block.Type {
|
||||
case "thinking", "redacted_thinking":
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(resp.Model)) && (block.Signature != "" || block.Data != "") {
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(resp.Model) || claude.IsHaiku55(resp.Model)) && (block.Signature != "" || block.Data != "") {
|
||||
item := ResponsesOutput{Type: "reasoning", ID: generateItemID(), EncryptedContent: encodeAnthropicThinking(block)}
|
||||
if block.Thinking != "" {
|
||||
item.Summary = []ResponsesSummary{{Type: "summary_text", Text: block.Thinking}}
|
||||
@@ -72,7 +72,7 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
|
||||
})
|
||||
}
|
||||
case "text":
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(resp.Model)) && block.Text != "" {
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(resp.Model) || claude.IsHaiku55(resp.Model)) && block.Text != "" {
|
||||
outputs = append(outputs, ResponsesOutput{Type: "message", ID: generateItemID(), Role: "assistant", Status: "completed", Content: []ResponsesContentPart{{Type: "output_text", Text: block.Text}}})
|
||||
continue
|
||||
}
|
||||
@@ -285,7 +285,7 @@ func ResponsesEventToSSE(evt ResponsesStreamEvent) (string, error) {
|
||||
func anthToResHandleMessageStart(evt *AnthropicStreamEvent, state *AnthropicEventToResponsesState) []ResponsesStreamEvent {
|
||||
if evt.Message != nil {
|
||||
state.ResponseID = evt.Message.ID
|
||||
state.PreserveThinkingSignatures = state.PreserveThinkingSignatures || claude.IsOpus55(evt.Message.Model) || claude.IsSonnet55(evt.Message.Model)
|
||||
state.PreserveThinkingSignatures = state.PreserveThinkingSignatures || claude.IsOpus55(evt.Message.Model) || claude.IsSonnet55(evt.Message.Model) || claude.IsHaiku55(evt.Message.Model)
|
||||
if state.Model == "" {
|
||||
state.Model = evt.Message.Model
|
||||
}
|
||||
|
||||
@@ -16,7 +16,8 @@ import (
|
||||
func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, error) {
|
||||
isOpus55 := claude.IsOpus55(req.Model)
|
||||
isSonnet55 := claude.IsSonnet55(req.Model)
|
||||
system, messages, err := convertResponsesInputToAnthropic(req.Instructions, req.Input, isOpus55 || isSonnet55)
|
||||
isHaiku55 := claude.IsHaiku55(req.Model)
|
||||
system, messages, err := convertResponsesInputToAnthropic(req.Instructions, req.Input, isOpus55 || isSonnet55 || isHaiku55)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -56,11 +57,13 @@ func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, erro
|
||||
out.ToolChoice = tc
|
||||
}
|
||||
|
||||
// The 5.5 models reject manual thinking and forced tool use. Sonnet 5.5
|
||||
// additionally supports between_tools to disable up-front thinking.
|
||||
// The 5.5 models reject manual thinking. Opus/Sonnet 5.5 additionally reject
|
||||
// forced tool use; Haiku 5.5 accepts it (starts with the tool call, no thinking
|
||||
// block). Sonnet 5.5 supports between_tools to disable up-front thinking;
|
||||
// Haiku 5.5 uses disabled (at high effort or below) instead.
|
||||
// Resolve the upstream model before conversion: client aliases need not
|
||||
// identify a Claude model.
|
||||
if isOpus55 || isSonnet55 {
|
||||
if isOpus55 || isSonnet55 || isHaiku55 {
|
||||
var choice struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
@@ -69,7 +72,7 @@ func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, erro
|
||||
return nil, fmt.Errorf("invalid tool_choice: %w", err)
|
||||
}
|
||||
}
|
||||
if choice.Type == "any" || choice.Type == "tool" {
|
||||
if (choice.Type == "any" || choice.Type == "tool") && !isHaiku55 {
|
||||
return nil, fmt.Errorf("%s does not support forced tool_choice; use auto or none", req.Model)
|
||||
}
|
||||
effort := "medium"
|
||||
@@ -82,6 +85,17 @@ func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, erro
|
||||
return nil, fmt.Errorf("claude-sonnet-5-5 does not support non-default top_p")
|
||||
}
|
||||
}
|
||||
if isHaiku55 {
|
||||
if req.Temperature != nil && *req.Temperature != 1 {
|
||||
return nil, fmt.Errorf("claude-haiku-5-5 does not support non-default temperature")
|
||||
}
|
||||
if req.TopP != nil && *req.TopP != 0.99 {
|
||||
return nil, fmt.Errorf("claude-haiku-5-5 does not support non-default top_p")
|
||||
}
|
||||
if req.Temperature != nil && req.TopP != nil {
|
||||
return nil, fmt.Errorf("claude-haiku-5-5 does not allow temperature and top_p together; omit both")
|
||||
}
|
||||
}
|
||||
if req.Reasoning != nil && req.Reasoning.Effort != "" {
|
||||
effort = req.Reasoning.Effort
|
||||
}
|
||||
@@ -96,6 +110,16 @@ func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, erro
|
||||
out.OutputConfig.Effort = "low"
|
||||
return out, nil
|
||||
}
|
||||
if isHaiku55 && effort == "none" {
|
||||
// Haiku 5.5's lowest thinking mode is disabled (valid at high
|
||||
// effort or below); low keeps token spend minimal like Sonnet above.
|
||||
out.Thinking = &AnthropicThinking{Type: "disabled"}
|
||||
if out.OutputConfig == nil {
|
||||
out.OutputConfig = &AnthropicOutputConfig{}
|
||||
}
|
||||
out.OutputConfig.Effort = "low"
|
||||
return out, nil
|
||||
}
|
||||
switch effort {
|
||||
case "low", "medium", "high", "xhigh", "max":
|
||||
default:
|
||||
|
||||
@@ -223,6 +223,12 @@ var DefaultModels = []Model{
|
||||
DisplayName: "Claude Haiku 4.5",
|
||||
CreatedAt: "2025-10-01T00:00:00Z",
|
||||
},
|
||||
{
|
||||
ID: "claude-haiku-5-5",
|
||||
Type: "model",
|
||||
DisplayName: "Claude Haiku 5.5",
|
||||
CreatedAt: "2026-10-07T00:00:00Z",
|
||||
},
|
||||
}
|
||||
|
||||
// DefaultModelIDs 返回默认模型的 ID 列表
|
||||
|
||||
@@ -43,3 +43,16 @@ func TestDefaultModelsContainsSonnet55(t *testing.T) {
|
||||
}
|
||||
t.Fatal("claude-sonnet-5-5 missing from DefaultModels")
|
||||
}
|
||||
|
||||
func TestDefaultModelsContainsHaiku55(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
for _, model := range DefaultModels {
|
||||
if model.ID == "claude-haiku-5-5" {
|
||||
require.Equal(t, "Claude Haiku 5.5", model.DisplayName)
|
||||
require.Equal(t, "2026-10-07T00:00:00Z", model.CreatedAt)
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatal("claude-haiku-5-5 missing from DefaultModels")
|
||||
}
|
||||
|
||||
@@ -27,6 +27,7 @@ var effortFamilies = []struct {
|
||||
{family: "claude-opus-4-5", levels: effortLowMediumHigh},
|
||||
{family: "claude-opus-5-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-opus-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-haiku-5-5", levels: effortLowMediumHighXHighMax},
|
||||
}
|
||||
|
||||
// EffortLevelsForModel returns the output_config.effort values accepted by a
|
||||
@@ -51,6 +52,11 @@ func IsSonnet55(model string) bool {
|
||||
return normalizeEffortModelID(model) == "claude-sonnet-5-5"
|
||||
}
|
||||
|
||||
// IsHaiku55 identifies the fixed Haiku 5.5 ID after provider/local suffix normalization.
|
||||
func IsHaiku55(model string) bool {
|
||||
return normalizeEffortModelID(model) == "claude-haiku-5-5"
|
||||
}
|
||||
|
||||
func normalizeEffortModelID(model string) string {
|
||||
id := strings.ToLower(strings.TrimSpace(model))
|
||||
id = strings.TrimPrefix(id, "models/")
|
||||
@@ -70,6 +76,9 @@ func normalizeEffortModelID(model string) string {
|
||||
if id == "claude-sonnet-5.5" {
|
||||
id = "claude-sonnet-5-5"
|
||||
}
|
||||
if id == "claude-haiku-5.5" {
|
||||
id = "claude-haiku-5-5"
|
||||
}
|
||||
if mapped, ok := ModelIDReverseOverrides[id]; ok {
|
||||
id = mapped
|
||||
}
|
||||
|
||||
@@ -18,6 +18,9 @@ func TestEffortLevelsForModel(t *testing.T) {
|
||||
{model: "claude-opus-5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "anthropic/claude-opus-5.5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "claude-sonnet-5-5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "claude-haiku-5-5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "anthropic/claude-haiku-5.5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "us.anthropic.claude-haiku-5-5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "us.anthropic.claude-sonnet-5-5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "claude-opus-4-5-20251101", want: []string{"low", "medium", "high"}},
|
||||
{model: "claude-haiku-4-5-20251001", want: nil},
|
||||
@@ -57,3 +60,19 @@ func TestIsSonnet55(t *testing.T) {
|
||||
require.False(t, IsSonnet55(model), model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsHaiku55(t *testing.T) {
|
||||
t.Parallel()
|
||||
for _, model := range []string{
|
||||
"claude-haiku-5-5",
|
||||
"anthropic/claude-haiku-5.5",
|
||||
"anthropic.claude-haiku-5-5",
|
||||
"us.anthropic.claude-haiku-5-5",
|
||||
"global.anthropic.claude-haiku-5-5-thinking",
|
||||
} {
|
||||
require.True(t, IsHaiku55(model), model)
|
||||
}
|
||||
for _, model := range []string{"claude-haiku-5", "claude-haiku-4-5", "claude-haiku-4-5-20251001", "claude-sonnet-5-5", "claude-opus-5-5"} {
|
||||
require.False(t, IsHaiku55(model), model)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -857,7 +857,7 @@ func (r *apiKeyRepository) ResetRateLimitWindows(ctx context.Context, id int64)
|
||||
// GetRateLimitData returns the current rate limit usage and window start times for an API key.
|
||||
func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (result *service.APIKeyRateLimitData, err error) {
|
||||
rows, err := r.sql.QueryContext(ctx, `
|
||||
SELECT usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start
|
||||
SELECT quota_used, usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start
|
||||
FROM api_keys
|
||||
WHERE id = $1 AND deleted_at IS NULL`,
|
||||
id)
|
||||
@@ -873,7 +873,7 @@ func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (resu
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
data := &service.APIKeyRateLimitData{}
|
||||
if err := rows.Scan(&data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil {
|
||||
if err := rows.Scan(&data.QuotaUsed, &data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return data, rows.Err()
|
||||
@@ -979,6 +979,7 @@ func groupEntityToService(g *dbent.Group) *service.Group {
|
||||
ID: g.ID,
|
||||
Name: g.Name,
|
||||
Description: derefString(g.Description),
|
||||
I18n: g.I18n,
|
||||
Platform: g.Platform,
|
||||
RateMultiplier: g.RateMultiplier,
|
||||
IsExclusive: g.IsExclusive,
|
||||
|
||||
@@ -464,6 +464,10 @@ func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_Basic() {
|
||||
newQuota, err = s.repo.IncrementQuotaUsed(s.ctx, key.ID, 2.5)
|
||||
s.Require().NoError(err, "IncrementQuotaUsed second")
|
||||
s.Require().Equal(4.0, newQuota, "第二次递增后应为 4.0")
|
||||
|
||||
data, err := s.repo.GetRateLimitData(s.ctx, key.ID)
|
||||
s.Require().NoError(err, "GetRateLimitData")
|
||||
s.Require().Equal(4.0, data.QuotaUsed, "用量读取带回总额度已用量")
|
||||
}
|
||||
|
||||
func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_NotFound() {
|
||||
|
||||
@@ -25,6 +25,20 @@ func billingInflightKeys(userID int64) (string, string) {
|
||||
fmt.Sprintf("%s{%d}", billingInflightAmtKeyPrefix, userID)
|
||||
}
|
||||
|
||||
// 订阅 / API Key 维度在途预留(确定金额费用):
|
||||
// - billing:inflight_scope:{scope} ZSET member=requestID score=过期时间(ms)
|
||||
// - billing:inflight_scope_amt:{scope} HASH field=requestID value=预留金额(USD)
|
||||
//
|
||||
// 过期清理与释放语义与余额维度相同,不续期。
|
||||
const (
|
||||
billingInflightScopeKeyPrefix = "billing:inflight_scope:"
|
||||
billingInflightScopeAmtKeyPrefix = "billing:inflight_scope_amt:"
|
||||
)
|
||||
|
||||
func billingInflightScopeKeys(scope string) (string, string) {
|
||||
return billingInflightScopeKeyPrefix + "{" + scope + "}", billingInflightScopeAmtKeyPrefix + "{" + scope + "}"
|
||||
}
|
||||
|
||||
var (
|
||||
// KEYS: [1]=zset [2]=hash
|
||||
// ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=balance [6]=key_ttl_ms
|
||||
@@ -85,6 +99,28 @@ var (
|
||||
redis.call('HDEL', KEYS[2], ARGV[1])
|
||||
return 1
|
||||
`)
|
||||
|
||||
// KEYS: [1]=zset [2]=hash
|
||||
// ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=key_ttl_ms
|
||||
// 先清理已过期成员,再登记本次预留;返回登记前其他在途预留的合计(string),是否放行由调用方按权威用量判定。
|
||||
reserveInflightScopeScript = redis.NewScript(`
|
||||
local now = tonumber(ARGV[1])
|
||||
local expired = redis.call('ZRANGEBYSCORE', KEYS[1], '-inf', now)
|
||||
if #expired > 0 then
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now)
|
||||
redis.call('HDEL', KEYS[2], unpack(expired))
|
||||
end
|
||||
local vals = redis.call('HVALS', KEYS[2])
|
||||
local sum = 0
|
||||
for i = 1, #vals do
|
||||
sum = sum + (tonumber(vals[i]) or 0)
|
||||
end
|
||||
redis.call('ZADD', KEYS[1], tonumber(ARGV[2]), ARGV[3])
|
||||
redis.call('HSET', KEYS[2], ARGV[3], ARGV[4])
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[5])
|
||||
redis.call('PEXPIRE', KEYS[2], ARGV[5])
|
||||
return tostring(sum)
|
||||
`)
|
||||
)
|
||||
|
||||
// ReserveInflightBalance 实现 service.InflightBalanceReservationCache。
|
||||
@@ -138,7 +174,35 @@ func (c *billingCache) RenewInflightBalance(ctx context.Context, userID int64, r
|
||||
return n == 1, nil
|
||||
}
|
||||
|
||||
// ReserveInflightScope 实现 service.InflightScopeReservationCache。
|
||||
func (c *billingCache) ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error) {
|
||||
zkey, hkey := billingInflightScopeKeys(scope)
|
||||
now := time.Now().UnixMilli()
|
||||
ttlMs := ttl.Milliseconds()
|
||||
if ttlMs <= 0 {
|
||||
ttlMs = 1
|
||||
}
|
||||
res, err := reserveInflightScopeScript.Run(ctx, c.rdb, []string{zkey, hkey},
|
||||
now,
|
||||
now+ttlMs,
|
||||
requestID,
|
||||
strconv.FormatFloat(amount, 'f', -1, 64),
|
||||
ttlMs,
|
||||
).Text()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return strconv.ParseFloat(res, 64)
|
||||
}
|
||||
|
||||
// ReleaseInflightScope 实现 service.InflightScopeReservationCache。
|
||||
func (c *billingCache) ReleaseInflightScope(ctx context.Context, scope, requestID string) error {
|
||||
zkey, hkey := billingInflightScopeKeys(scope)
|
||||
return releaseInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID).Err()
|
||||
}
|
||||
|
||||
var (
|
||||
_ service.InflightBalanceReservationCache = (*billingCache)(nil)
|
||||
_ service.InflightBalanceReservationRenewer = (*billingCache)(nil)
|
||||
_ service.InflightScopeReservationCache = (*billingCache)(nil)
|
||||
)
|
||||
|
||||
@@ -285,3 +285,72 @@ func TestInflightReservation_RenewDoesNotResurrectExpiredMember(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.False(t, exists)
|
||||
}
|
||||
|
||||
func TestInflightScopeReservation_ReturnsOthersAndReleases(t *testing.T) {
|
||||
mr, cache, _ := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
|
||||
others, err := cache.ReserveInflightScope(ctx, "apikey:7", "r1", 3, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, others)
|
||||
others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r2", 2, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 3, others, 1e-9, "返回登记前其他在途预留的合计")
|
||||
others, err = cache.ReserveInflightScope(ctx, "sub:7", "r3", 5, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, others, "不同维度互不影响")
|
||||
|
||||
require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1"))
|
||||
require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1"))
|
||||
others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r4", 1, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 2, others, 1e-9)
|
||||
|
||||
mr.FastForward(2 * time.Minute)
|
||||
others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r5", 1, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, others, "过期预留不再计入")
|
||||
}
|
||||
|
||||
func TestReserveFixedCharge_ConcurrentChargesAdmitOnlyWhatAPIKeyQuotaCovers(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 42}
|
||||
apiKey := &service.APIKey{ID: 9, UserID: user.ID, Quota: 1.0}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 100))
|
||||
|
||||
const n = 20
|
||||
var admitted atomic.Int32
|
||||
reservations := make(chan *service.InflightReservation, n)
|
||||
var wg sync.WaitGroup
|
||||
start := make(chan struct{})
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
res, err := svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 0.3)
|
||||
if err != nil {
|
||||
if !errors.Is(err, service.ErrAPIKeyQuotaExhausted) {
|
||||
t.Errorf("unexpected error: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
admitted.Add(1)
|
||||
reservations <- res
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(reservations)
|
||||
|
||||
require.Equal(t, int32(3), admitted.Load(), "Key 额度 1.0、每笔 0.3:同时在途的不超过 3 笔")
|
||||
for res := range reservations {
|
||||
res.Release()
|
||||
}
|
||||
zkey, _ := billingInflightScopeKeys("apikey:9")
|
||||
remaining, err := cache.rdb.ZCard(ctx, zkey).Result()
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, remaining)
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,205 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
type geminiCachedContentSQLExecutor interface {
|
||||
ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
|
||||
QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
|
||||
QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
|
||||
}
|
||||
|
||||
type geminiCachedContentRepository struct {
|
||||
sql geminiCachedContentSQLExecutor
|
||||
}
|
||||
|
||||
func NewGeminiCachedContentRepository(db *sql.DB) service.GeminiCachedContentRepository {
|
||||
return &geminiCachedContentRepository{sql: db}
|
||||
}
|
||||
|
||||
const geminiCachedContentColumns = `id, public_id, user_id, api_key_id, group_id, account_id, upstream_name, model,
|
||||
upstream_model, display_name, total_token_count, expire_time, channel_usage, created_at, updated_at`
|
||||
|
||||
func (r *geminiCachedContentRepository) Create(ctx context.Context, record *service.GeminiCachedContent) error {
|
||||
channelUsage, err := json.Marshal(record.ChannelUsage)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
row := r.sql.QueryRowContext(ctx, `
|
||||
INSERT INTO gemini_cached_contents (public_id, user_id, api_key_id, group_id, account_id, upstream_name, model,
|
||||
upstream_model, display_name, total_token_count, expire_time, channel_usage)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
|
||||
RETURNING id, created_at, updated_at`,
|
||||
record.PublicID, record.UserID, record.APIKeyID, record.GroupID, record.AccountID, record.UpstreamName,
|
||||
record.Model, record.UpstreamModel, record.DisplayName, record.TotalTokenCount, record.ExpireTime, channelUsage)
|
||||
if err := row.Scan(&record.ID, &record.CreatedAt, &record.UpdatedAt); err != nil {
|
||||
return translatePersistenceError(err, nil, service.ErrGeminiCachedContentExists)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) GetForOwner(ctx context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) {
|
||||
record, err := scanGeminiCachedContent(r.sql.QueryRowContext(ctx, `SELECT `+geminiCachedContentColumns+`
|
||||
FROM gemini_cached_contents
|
||||
WHERE public_id = $1 AND api_key_id = $2 AND deleted_at IS NULL`, publicID, apiKeyID))
|
||||
if err != nil {
|
||||
return nil, translatePersistenceError(err, service.ErrGeminiCachedContentNotFound, nil)
|
||||
}
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) ListForOwner(ctx context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*service.GeminiCachedContent, error) {
|
||||
if limit <= 0 {
|
||||
limit = service.GeminiCachedContentListDefaultPageSize
|
||||
}
|
||||
rows, err := r.sql.QueryContext(ctx, `SELECT `+geminiCachedContentColumns+`
|
||||
FROM gemini_cached_contents
|
||||
WHERE api_key_id = $1 AND group_id = $2 AND deleted_at IS NULL AND expire_time > $3 AND ($4 <= 0 OR id < $4)
|
||||
ORDER BY id DESC
|
||||
LIMIT $5`, apiKeyID, groupID, activeAt, beforeID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
var out []*service.GeminiCachedContent
|
||||
for rows.Next() {
|
||||
record, err := scanGeminiCachedContent(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, record)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) BeginMutation(ctx context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) {
|
||||
result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision)
|
||||
WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NULL AND expire_time = $4`,
|
||||
id, token, lease.Seconds(), baseline)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if affected > 0 {
|
||||
return true, nil
|
||||
}
|
||||
var exists bool
|
||||
if err := r.sql.QueryRowContext(ctx, `SELECT EXISTS (
|
||||
SELECT 1 FROM gemini_cached_contents WHERE id = $1 AND deleted_at IS NULL)`, id).Scan(&exists); err != nil {
|
||||
return false, err
|
||||
}
|
||||
if !exists {
|
||||
return false, service.ErrGeminiCachedContentNotFound
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) TakeOverStaleMutation(ctx context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) {
|
||||
record, err := scanGeminiCachedContent(r.sql.QueryRowContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision)
|
||||
WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NOT NULL AND mutation_lease_until < NOW()
|
||||
RETURNING `+geminiCachedContentColumns, id, token, lease.Seconds()))
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) CommitMutation(ctx context.Context, id int64, token string, expireTime time.Time) error {
|
||||
result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET expire_time = $3, updated_at = NOW(), mutation_token = NULL, mutation_lease_until = NULL
|
||||
WHERE id = $1 AND deleted_at IS NULL AND mutation_token = $2`, id, token, expireTime)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected == 0 {
|
||||
return service.ErrGeminiCachedContentMutationConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) ReleaseMutation(ctx context.Context, id int64, token string) error {
|
||||
_, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET mutation_token = NULL, mutation_lease_until = NULL
|
||||
WHERE id = $1 AND mutation_token = $2`, id, token)
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) ListStaleMutations(ctx context.Context, limit int) ([]int64, error) {
|
||||
rows, err := r.sql.QueryContext(ctx, `SELECT id FROM gemini_cached_contents
|
||||
WHERE mutation_token IS NOT NULL AND deleted_at IS NULL AND mutation_lease_until < NOW()
|
||||
ORDER BY mutation_lease_until
|
||||
LIMIT $1`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
var ids []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids, rows.Err()
|
||||
}
|
||||
|
||||
func (r *geminiCachedContentRepository) SoftDelete(ctx context.Context, id int64) error {
|
||||
_, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents
|
||||
SET deleted_at = NOW(), updated_at = NOW()
|
||||
WHERE id = $1 AND deleted_at IS NULL`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
type geminiCachedContentScanner interface {
|
||||
Scan(dest ...any) error
|
||||
}
|
||||
|
||||
func scanGeminiCachedContent(row geminiCachedContentScanner) (*service.GeminiCachedContent, error) {
|
||||
var record service.GeminiCachedContent
|
||||
var channelUsage []byte
|
||||
if err := row.Scan(
|
||||
&record.ID,
|
||||
&record.PublicID,
|
||||
&record.UserID,
|
||||
&record.APIKeyID,
|
||||
&record.GroupID,
|
||||
&record.AccountID,
|
||||
&record.UpstreamName,
|
||||
&record.Model,
|
||||
&record.UpstreamModel,
|
||||
&record.DisplayName,
|
||||
&record.TotalTokenCount,
|
||||
&record.ExpireTime,
|
||||
&channelUsage,
|
||||
&record.CreatedAt,
|
||||
&record.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(channelUsage) > 0 {
|
||||
if err := json.Unmarshal(channelUsage, &record.ChannelUsage); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return &record, nil
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGeminiCachedContentRepository_Lifecycle(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := &geminiCachedContentRepository{sql: tx}
|
||||
now := time.Now().UTC().Truncate(time.Microsecond)
|
||||
suffix := time.Now().Format("150405.000000")
|
||||
apiKeyID := int64(910001)
|
||||
|
||||
newRecord := func(publicID string, keyID, groupID int64, expire time.Time) *service.GeminiCachedContent {
|
||||
return &service.GeminiCachedContent{
|
||||
PublicID: publicID + suffix, UserID: 1, APIKeyID: keyID, GroupID: groupID, AccountID: 99,
|
||||
UpstreamName: "cachedContents/up" + publicID, Model: "gemini-3.8-flash", UpstreamModel: "models/gemini-3.8-flash",
|
||||
DisplayName: "label", TotalTokenCount: 2730, ExpireTime: expire,
|
||||
ChannelUsage: service.ChannelUsageFields{ChannelID: 7, OriginalModel: "gemini-pro", ChannelMappedModel: "gemini-3.8-flash", BillingModelSource: "requested"},
|
||||
}
|
||||
}
|
||||
first := newRecord("a", apiKeyID, 42, now.Add(time.Hour))
|
||||
require.NoError(t, repo.Create(ctx, first))
|
||||
require.Positive(t, first.ID)
|
||||
require.False(t, first.CreatedAt.IsZero())
|
||||
_, err := tx.ExecContext(ctx, `SAVEPOINT gemini_cached_content_duplicate`)
|
||||
require.NoError(t, err)
|
||||
require.ErrorIs(t, repo.Create(ctx, newRecord("a", apiKeyID, 42, now.Add(time.Hour))), service.ErrGeminiCachedContentExists)
|
||||
_, err = tx.ExecContext(ctx, `ROLLBACK TO SAVEPOINT gemini_cached_content_duplicate`)
|
||||
require.NoError(t, err)
|
||||
|
||||
second := newRecord("b", apiKeyID, 42, now.Add(time.Hour))
|
||||
require.NoError(t, repo.Create(ctx, second))
|
||||
require.NoError(t, repo.Create(ctx, newRecord("expired", apiKeyID, 42, now.Add(-time.Minute))))
|
||||
require.NoError(t, repo.Create(ctx, newRecord("othergroup", apiKeyID, 43, now.Add(time.Hour))))
|
||||
require.NoError(t, repo.Create(ctx, newRecord("otherkey", apiKeyID+1, 42, now.Add(time.Hour))))
|
||||
|
||||
got, err := repo.GetForOwner(ctx, apiKeyID, first.PublicID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "cachedContents/upa", got.UpstreamName)
|
||||
require.Equal(t, int64(2730), got.TotalTokenCount)
|
||||
require.True(t, got.ExpireTime.Equal(first.ExpireTime))
|
||||
require.Equal(t, first.ChannelUsage, got.ChannelUsage, "创建时的渠道计费口径随记录持久化")
|
||||
_, err = repo.GetForOwner(ctx, apiKeyID+1, first.PublicID)
|
||||
require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound)
|
||||
|
||||
list, err := repo.ListForOwner(ctx, apiKeyID, 42, now, 0, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, list, 2)
|
||||
require.Equal(t, second.PublicID, list[0].PublicID)
|
||||
page, err := repo.ListForOwner(ctx, apiKeyID, 42, now, list[0].ID, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, page, 1)
|
||||
require.Equal(t, first.PublicID, page[0].PublicID)
|
||||
|
||||
lease := 2 * time.Minute
|
||||
acquired, err := repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t1", lease)
|
||||
require.NoError(t, err)
|
||||
require.True(t, acquired)
|
||||
acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t2", lease)
|
||||
require.NoError(t, err)
|
||||
require.False(t, acquired, "有进行中的修改时不能再登记")
|
||||
taken, err := repo.TakeOverStaleMutation(ctx, first.ID, "t2", lease)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, taken, "租约未过期不能接管")
|
||||
require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t2", now), service.ErrGeminiCachedContentMutationConflict)
|
||||
|
||||
newExpire := now.Add(3 * time.Hour)
|
||||
require.NoError(t, repo.CommitMutation(ctx, first.ID, "t1", newExpire))
|
||||
got, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, got.ExpireTime.Equal(newExpire))
|
||||
require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t1", now), service.ErrGeminiCachedContentMutationConflict, "提交后修改意图已清除")
|
||||
|
||||
acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t3", lease)
|
||||
require.NoError(t, err)
|
||||
require.False(t, acquired, "到期时间已被其他修改改变")
|
||||
acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t3", lease)
|
||||
require.NoError(t, err)
|
||||
require.True(t, acquired)
|
||||
stale, err := repo.ListStaleMutations(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, stale, first.ID)
|
||||
|
||||
_, err = tx.ExecContext(ctx, `UPDATE gemini_cached_contents SET mutation_lease_until = NOW() - INTERVAL '1 second' WHERE id = $1`, first.ID)
|
||||
require.NoError(t, err)
|
||||
stale, err = repo.ListStaleMutations(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, stale, first.ID)
|
||||
taken, err = repo.TakeOverStaleMutation(ctx, first.ID, "t4", lease)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, taken)
|
||||
require.True(t, taken.ExpireTime.Equal(newExpire), "接管返回修改前的到期时间")
|
||||
require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t3", now), service.ErrGeminiCachedContentMutationConflict, "被接管后原持有者不能提交")
|
||||
require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t3"))
|
||||
acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t5", lease)
|
||||
require.NoError(t, err)
|
||||
require.False(t, acquired, "原持有者释放不影响接管者")
|
||||
require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t4"))
|
||||
stale, err = repo.ListStaleMutations(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, stale, first.ID)
|
||||
|
||||
require.NoError(t, repo.SoftDelete(ctx, first.ID))
|
||||
_, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID)
|
||||
require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound)
|
||||
_, err = repo.BeginMutation(ctx, first.ID, newExpire, "t6", lease)
|
||||
require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound)
|
||||
list, err = repo.ListForOwner(ctx, apiKeyID, 42, now, 0, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, list, 1)
|
||||
}
|
||||
@@ -167,6 +167,10 @@ func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *servi
|
||||
// 设置支持的模型系列(始终设置,空数组表示不限制)
|
||||
builder = builder.SetSupportedModelScopes(groupIn.SupportedModelScopes)
|
||||
|
||||
if groupIn.I18n != nil {
|
||||
builder = builder.SetI18n(groupIn.I18n)
|
||||
}
|
||||
|
||||
created, err := builder.Save(ctx)
|
||||
if err != nil {
|
||||
return translatePersistenceError(err, nil, service.ErrGroupExists)
|
||||
@@ -431,6 +435,12 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
|
||||
// 处理 SupportedModelScopes(始终设置,空数组表示不限制)
|
||||
builder = builder.SetSupportedModelScopes(groupIn.SupportedModelScopes)
|
||||
|
||||
i18n := groupIn.I18n
|
||||
if i18n == nil {
|
||||
i18n = service.GroupI18n{}
|
||||
}
|
||||
builder = builder.SetI18n(i18n)
|
||||
|
||||
updated, err := builder.Save(ctx)
|
||||
if err != nil {
|
||||
return translatePersistenceError(err, service.ErrGroupNotFound, service.ErrGroupExists)
|
||||
@@ -957,6 +967,29 @@ func (r *groupRepository) deleteCascade(ctx context.Context, id int64, requireEm
|
||||
if _, err := txClient.Group.Delete().Where(group.IDEQ(id)).Exec(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Lock and validate the text value before using JSONB operators. A legacy
|
||||
// malformed setting must not prevent deletion or be silently replaced.
|
||||
var moderationValue string
|
||||
err = scanSingleRow(ctx, exec, `SELECT value FROM settings WHERE key = $1 FOR UPDATE`,
|
||||
[]any{service.SettingKeyContentModerationConfig}, &moderationValue)
|
||||
if err != nil && !errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, err
|
||||
}
|
||||
// Keep the moderation scope in sync in the same transaction. Updating the
|
||||
// setting row in SQL also serializes concurrent group deletions across nodes.
|
||||
if err == nil && json.Valid([]byte(moderationValue)) {
|
||||
if _, err := exec.ExecContext(ctx, `UPDATE settings
|
||||
SET value = jsonb_set(value::jsonb, '{group_ids}',
|
||||
(SELECT COALESCE(jsonb_agg(group_id ORDER BY ord), '[]'::jsonb)
|
||||
FROM jsonb_array_elements(value::jsonb->'group_ids') WITH ORDINALITY AS ids(group_id, ord)
|
||||
WHERE group_id <> to_jsonb($2::bigint)))::text,
|
||||
updated_at = NOW()
|
||||
WHERE key = $1 AND jsonb_typeof(value::jsonb->'group_ids') = 'array'
|
||||
AND value::jsonb->'group_ids' @> jsonb_build_array($2::bigint)`,
|
||||
service.SettingKeyContentModerationConfig, id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
if tx != nil {
|
||||
if err := tx.Commit(); err != nil {
|
||||
|
||||
@@ -5,7 +5,9 @@ package repository
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -49,6 +51,55 @@ func TestGroupRepoSuite(t *testing.T) {
|
||||
suite.Run(t, new(GroupRepoSuite))
|
||||
}
|
||||
|
||||
func TestModerationConfigWriteRejectsGroupDeletedAfterValidation(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
repo := newGroupRepositoryWithSQL(client, integrationDB)
|
||||
settings := &settingRepository{client: client}
|
||||
group := &service.Group{Name: "moderation-concurrent-delete", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
|
||||
require.NoError(t, repo.Create(ctx, group))
|
||||
t.Cleanup(func() { _, _ = integrationDB.ExecContext(ctx, "DELETE FROM groups WHERE id=$1", group.ID) })
|
||||
key := service.SettingKeyContentModerationConfig
|
||||
previous, previousErr := settings.GetValue(ctx, key)
|
||||
t.Cleanup(func() {
|
||||
if previousErr == nil {
|
||||
_ = settings.Set(ctx, key, previous)
|
||||
} else {
|
||||
_ = settings.Delete(ctx, key)
|
||||
}
|
||||
})
|
||||
require.NoError(t, settings.Set(ctx, key, `{"all_groups":false,"group_ids":[]}`))
|
||||
|
||||
// Hold the deleter's group lock while the config writer starts with an ID
|
||||
// that was valid before deletion. It must recheck after this transaction commits.
|
||||
tx, err := integrationDB.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
var id int64
|
||||
require.NoError(t, tx.QueryRowContext(ctx, "SELECT id FROM groups WHERE id=$1 FOR UPDATE", group.ID).Scan(&id))
|
||||
finished := make(chan error, 1)
|
||||
go func() {
|
||||
finished <- settings.SetContentModerationConfig(ctx, `{"all_groups":false,"group_ids":[`+fmt.Sprint(group.ID)+`]}`, []int64{group.ID})
|
||||
}()
|
||||
select {
|
||||
case err := <-finished:
|
||||
t.Fatalf("config write passed a locked group before deletion: %v", err)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
}
|
||||
_, err = tx.ExecContext(ctx, "UPDATE groups SET deleted_at=NOW() WHERE id=$1", group.ID)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, tx.Commit())
|
||||
select {
|
||||
case err := <-finished:
|
||||
require.ErrorIs(t, err, service.ErrGroupNotFound)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("config write did not finish after deletion")
|
||||
}
|
||||
stored, err := settings.GetValue(ctx, key)
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{"all_groups":false,"group_ids":[]}`, stored)
|
||||
}
|
||||
|
||||
// --- Create / GetByID / Update / Delete ---
|
||||
|
||||
func (s *GroupRepoSuite) TestCreate() {
|
||||
@@ -201,6 +252,48 @@ func (s *GroupRepoSuite) TestUpdate() {
|
||||
s.Require().Equal("updated", got.Name)
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestI18n_RoundTrip() {
|
||||
group := &service.Group{
|
||||
Name: "i18n-roundtrip",
|
||||
Platform: service.PlatformAnthropic,
|
||||
RateMultiplier: 1.0,
|
||||
Status: service.StatusActive,
|
||||
SubscriptionType: service.SubscriptionTypeStandard,
|
||||
I18n: service.GroupI18n{
|
||||
"en": {Name: "Standard", Description: "Pro + Max pool"},
|
||||
},
|
||||
}
|
||||
s.Require().NoError(s.repo.Create(s.ctx, group))
|
||||
|
||||
got, err := s.repo.GetByID(s.ctx, group.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().Equal(group.I18n, got.I18n)
|
||||
|
||||
got.I18n = service.GroupI18n{"zh": {Description: "号池"}}
|
||||
s.Require().NoError(s.repo.Update(s.ctx, got))
|
||||
got, err = s.repo.GetByID(s.ctx, group.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().Equal(service.GroupI18n{"zh": {Description: "号池"}}, got.I18n)
|
||||
|
||||
// 未携带译文的写入落库为空对象,而不是 JSON null。
|
||||
got.I18n = nil
|
||||
s.Require().NoError(s.repo.Update(s.ctx, got))
|
||||
var stored string
|
||||
s.Require().NoError(scanSingleRow(s.ctx, s.tx, "SELECT i18n::text FROM groups WHERE id = $1", []any{group.ID}, &stored))
|
||||
s.Require().Equal("{}", stored)
|
||||
|
||||
plain := &service.Group{
|
||||
Name: "i18n-default",
|
||||
Platform: service.PlatformAnthropic,
|
||||
RateMultiplier: 1.0,
|
||||
Status: service.StatusActive,
|
||||
SubscriptionType: service.SubscriptionTypeStandard,
|
||||
}
|
||||
s.Require().NoError(s.repo.Create(s.ctx, plain))
|
||||
s.Require().NoError(scanSingleRow(s.ctx, s.tx, "SELECT i18n::text FROM groups WHERE id = $1", []any{plain.ID}, &stored))
|
||||
s.Require().Equal("{}", stored)
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestGetByID_PreservesMessagesDispatchModelConfig() {
|
||||
group := &service.Group{
|
||||
Name: "openai-dispatch",
|
||||
@@ -257,8 +350,14 @@ func (s *GroupRepoSuite) TestDeleteCascadeIfEmptyRejectsGroupWithNonDeletedAccou
|
||||
_, err := s.tx.ExecContext(s.ctx, "INSERT INTO account_groups (account_id, group_id, priority, created_at) VALUES ($1, $2, 1, NOW())", accountID, group.ID)
|
||||
s.Require().NoError(err)
|
||||
|
||||
config, marshalErr := json.Marshal(map[string]any{"group_ids": []int64{group.ID}, "all_groups": false})
|
||||
s.Require().NoError(marshalErr)
|
||||
s.Require().NoError(NewSettingRepository(s.tx.Client()).Set(s.ctx, service.SettingKeyContentModerationConfig, string(config)))
|
||||
_, err = s.repo.DeleteCascadeIfEmpty(s.ctx, group.ID)
|
||||
s.Require().ErrorIs(err, service.ErrGroupNotEmpty)
|
||||
stored, readErr := NewSettingRepository(s.tx.Client()).GetValue(s.ctx, service.SettingKeyContentModerationConfig)
|
||||
s.Require().NoError(readErr)
|
||||
s.Require().JSONEq(string(config), stored)
|
||||
_, err = s.repo.GetByID(s.ctx, group.ID)
|
||||
s.Require().NoError(err)
|
||||
var bindings int
|
||||
@@ -276,6 +375,76 @@ func (s *GroupRepoSuite) TestDeleteCascadeIfEmptyDeletesEmptyGroup() {
|
||||
s.Require().ErrorIs(err, service.ErrGroupNotFound)
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestDeleteCascadeRemovesOnlySelectedModerationGroup() {
|
||||
for _, guarded := range []bool{false, true} {
|
||||
name := map[bool]string{false: "normal", true: "guarded"}[guarded]
|
||||
s.Run(name, func() {
|
||||
deleted := &service.Group{Name: "moderation-deleted-" + name, Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
|
||||
kept := &service.Group{Name: "moderation-kept-" + name, Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
|
||||
s.Require().NoError(s.repo.Create(s.ctx, deleted))
|
||||
s.Require().NoError(s.repo.Create(s.ctx, kept))
|
||||
config := map[string]any{"group_ids": []int64{kept.ID, deleted.ID}, "all_groups": false, "mode": "observe"}
|
||||
raw, err := json.Marshal(config)
|
||||
s.Require().NoError(err)
|
||||
s.Require().NoError(NewSettingRepository(s.tx.Client()).Set(s.ctx, service.SettingKeyContentModerationConfig, string(raw)))
|
||||
if guarded {
|
||||
_, err = s.repo.DeleteCascadeIfEmpty(s.ctx, deleted.ID)
|
||||
} else {
|
||||
_, err = s.repo.DeleteCascade(s.ctx, deleted.ID)
|
||||
}
|
||||
s.Require().NoError(err)
|
||||
stored, err := NewSettingRepository(s.tx.Client()).GetValue(s.ctx, service.SettingKeyContentModerationConfig)
|
||||
s.Require().NoError(err)
|
||||
var got struct {
|
||||
GroupIDs []int64 `json:"group_ids"`
|
||||
Mode string `json:"mode"`
|
||||
}
|
||||
s.Require().NoError(json.Unmarshal([]byte(stored), &got))
|
||||
s.Require().Equal([]int64{kept.ID}, got.GroupIDs)
|
||||
s.Require().Equal("observe", got.Mode)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestDeleteCascadeUnknownGroupPreservesModerationConfig() {
|
||||
const original = `{"all_groups":false,"group_ids":[123,456],"mode":"observe"}`
|
||||
s.Require().NoError(NewSettingRepository(s.tx.Client()).Set(s.ctx, service.SettingKeyContentModerationConfig, original))
|
||||
for _, guarded := range []bool{false, true} {
|
||||
var err error
|
||||
if guarded {
|
||||
_, err = s.repo.DeleteCascadeIfEmpty(s.ctx, -1)
|
||||
} else {
|
||||
_, err = s.repo.DeleteCascade(s.ctx, -1)
|
||||
}
|
||||
s.Require().ErrorIs(err, service.ErrGroupNotFound)
|
||||
stored, readErr := NewSettingRepository(s.tx.Client()).GetValue(s.ctx, service.SettingKeyContentModerationConfig)
|
||||
s.Require().NoError(readErr)
|
||||
s.Require().JSONEq(original, stored)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestDeleteCascadePreservesMalformedModerationConfig() {
|
||||
for _, guarded := range []bool{false, true} {
|
||||
name := map[bool]string{false: "normal", true: "guarded"}[guarded]
|
||||
s.Run(name, func() {
|
||||
group := &service.Group{Name: "malformed-moderation-" + name, Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
|
||||
s.Require().NoError(s.repo.Create(s.ctx, group))
|
||||
const malformed = `{"group_ids":[oops]}`
|
||||
s.Require().NoError(NewSettingRepository(s.tx.Client()).Set(s.ctx, service.SettingKeyContentModerationConfig, malformed))
|
||||
var err error
|
||||
if guarded {
|
||||
_, err = s.repo.DeleteCascadeIfEmpty(s.ctx, group.ID)
|
||||
} else {
|
||||
_, err = s.repo.DeleteCascade(s.ctx, group.ID)
|
||||
}
|
||||
s.Require().NoError(err)
|
||||
stored, err := NewSettingRepository(s.tx.Client()).GetValue(s.ctx, service.SettingKeyContentModerationConfig)
|
||||
s.Require().NoError(err)
|
||||
s.Require().Equal(malformed, stored)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestDeleteCascadeIfEmptyIgnoresBindingsToSoftDeletedAccounts() {
|
||||
group := &service.Group{Name: "guarded-soft-deleted-account", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
|
||||
s.Require().NoError(s.repo.Create(s.ctx, group))
|
||||
|
||||
@@ -956,7 +956,11 @@ func filterSchedulerCredentials(credentials map[string]any) map[string]any {
|
||||
}
|
||||
// Candidate-list admission evaluates the account override before hydrating
|
||||
// the full account. Dropping it silently falls back to the platform threshold.
|
||||
keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold"}
|
||||
keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold",
|
||||
// Gemini 显式缓存的候选过滤(GeminiExplicitCacheExclusions)读本投影判定账号是否支持显式缓存:
|
||||
// 官方地址、Vertex 认证方式与上游网关的显式缓存声明。缺这几个键时只能按类型粗筛,
|
||||
// 不支持的账号要逐个选中后才被剔除。
|
||||
"base_url", "auth_mode", service.GeminiExplicitCacheUpstreamCredentialKey, service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey}
|
||||
filtered := make(map[string]any)
|
||||
for _, key := range keys {
|
||||
if value, ok := credentials[key]; ok && value != nil {
|
||||
|
||||
@@ -20,6 +20,23 @@ func TestFilterSchedulerCredentialsKeepsSubscriptionPlanType(t *testing.T) {
|
||||
require.NotContains(t, filtered, "refresh_token")
|
||||
}
|
||||
|
||||
func TestFilterSchedulerCredentialsKeepsGeminiExplicitCacheKeys(t *testing.T) {
|
||||
filtered := filterSchedulerCredentials(map[string]any{
|
||||
"api_key": "k",
|
||||
"base_url": "https://gw.example.com",
|
||||
"auth_mode": "apikey",
|
||||
"service_account_json": "secret",
|
||||
service.GeminiExplicitCacheUpstreamCredentialKey: true,
|
||||
service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey: 3600,
|
||||
})
|
||||
|
||||
require.Equal(t, "https://gw.example.com", filtered["base_url"])
|
||||
require.Equal(t, "apikey", filtered["auth_mode"])
|
||||
require.Equal(t, true, filtered[service.GeminiExplicitCacheUpstreamCredentialKey])
|
||||
require.Equal(t, 3600, filtered[service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey])
|
||||
require.NotContains(t, filtered, "service_account_json")
|
||||
}
|
||||
|
||||
func TestSchedulerMetadataAccountKeepsOpenAISubscriptionIdentity(t *testing.T) {
|
||||
account := service.Account{
|
||||
ID: 24,
|
||||
|
||||
@@ -2,6 +2,7 @@ package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/ent"
|
||||
@@ -9,6 +10,34 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
// SetContentModerationConfig holds live group rows until the setting is written.
|
||||
// Deletion takes FOR UPDATE on the group before cleaning this setting, so both
|
||||
// operations follow the same lock order: groups, then settings.
|
||||
func (r *settingRepository) SetContentModerationConfig(ctx context.Context, value string, groupIDs []int64) error {
|
||||
tx, err := r.client.Tx(ctx)
|
||||
if err != nil && !errors.Is(err, ent.ErrTxStarted) {
|
||||
return err
|
||||
}
|
||||
client := r.client
|
||||
if err == nil {
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
client = tx.Client()
|
||||
}
|
||||
if err := lockLiveGroups(ctx, client, groupIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now()
|
||||
if err := client.Setting.Create().SetKey(service.SettingKeyContentModerationConfig).
|
||||
SetValue(value).SetUpdatedAt(now).OnConflictColumns(setting.FieldKey).
|
||||
UpdateNewValues().Exec(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
if tx != nil {
|
||||
return tx.Commit()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type settingRepository struct {
|
||||
client *ent.Client
|
||||
}
|
||||
|
||||
@@ -82,6 +82,7 @@ var ProviderSet = wire.NewSet(
|
||||
NewUsageLogRepository,
|
||||
NewUsageBillingRepository,
|
||||
NewBatchImageRepository,
|
||||
NewGeminiCachedContentRepository,
|
||||
NewIdempotencyRepository,
|
||||
NewUsageCleanupRepository,
|
||||
NewDashboardAggregationRepository,
|
||||
|
||||
@@ -349,6 +349,12 @@ func RegisterGatewayRoutes(
|
||||
gemini.GET("/models/:model", h.Gateway.GeminiV1BetaGetModel)
|
||||
// Gin treats ":" as a param marker, but Gemini uses "{model}:{action}" in the same segment.
|
||||
gemini.POST("/models/*modelAction", h.Gateway.GeminiV1BetaModels)
|
||||
// 显式上下文缓存:网关自管资源 ID 与账号绑定,list/get 只读本地记录。
|
||||
gemini.POST("/cachedContents", h.Gateway.GeminiCachedContentsCreate)
|
||||
gemini.GET("/cachedContents", h.Gateway.GeminiCachedContentsList)
|
||||
gemini.GET("/cachedContents/:cacheID", h.Gateway.GeminiCachedContentsGet)
|
||||
gemini.PATCH("/cachedContents/:cacheID", h.Gateway.GeminiCachedContentsPatch)
|
||||
gemini.DELETE("/cachedContents/:cacheID", h.Gateway.GeminiCachedContentsDelete)
|
||||
}
|
||||
|
||||
// OpenAI Responses API(不带v1前缀的别名)— auto-route based on group platform
|
||||
|
||||
@@ -47,6 +47,7 @@ func TestEveryGatewayPOSTRouteIsClassifiedForPromptAuditCoverage(t *testing.T) {
|
||||
"/videos/edits": {"grok_media.go"},
|
||||
"/videos/extensions": {"grok_media.go"},
|
||||
"/models/*modelAction": {"gemini_v1beta_handler.go"},
|
||||
"/cachedContents": {"gemini_cached_contents_handler.go"},
|
||||
"/tts": {"grok_audio.go"},
|
||||
"/web_search": {"gateway_web_search.go"},
|
||||
"/x_search": {"gateway_web_search.go"},
|
||||
|
||||
@@ -129,11 +129,20 @@ func isOpenAIPersonalAccessTokenAuthMode(value string) bool {
|
||||
}
|
||||
}
|
||||
|
||||
// TempUnschedulableRule.Scope 的取值:命中后隔离的粒度。
|
||||
const (
|
||||
TempUnschedScopeAccount = "account"
|
||||
TempUnschedScopeModel = "model"
|
||||
)
|
||||
|
||||
type TempUnschedulableRule struct {
|
||||
ErrorCode int `json:"error_code"`
|
||||
Keywords []string `json:"keywords"`
|
||||
DurationMinutes int `json:"duration_minutes"`
|
||||
Description string `json:"description"`
|
||||
// Scope 决定命中后的隔离粒度:account 暂停整个账号,model 只冷却本次请求的模型。
|
||||
// 为空时模型已知且非 401 走模型级,否则账号级。
|
||||
Scope string `json:"scope,omitempty"`
|
||||
}
|
||||
|
||||
func (a *Account) IsActive() bool {
|
||||
@@ -339,6 +348,13 @@ func (a *Account) IsGeminiGoogleOne() bool {
|
||||
return a.Platform == PlatformGemini && a.Type == AccountTypeOAuth && a.GeminiOAuthType() == "google_one"
|
||||
}
|
||||
|
||||
// IsGeminiThirdPartyAPIKey reports whether a Gemini API Key account points at a
|
||||
// non-Google base URL. Such upstreams carry no AI Studio tier or daily quota.
|
||||
func (a *Account) IsGeminiThirdPartyAPIKey() bool {
|
||||
return a != nil && a.Platform == PlatformGemini && a.Type == AccountTypeAPIKey &&
|
||||
!isOfficialGeminiAPIBaseURL(a.GetCredential("base_url"))
|
||||
}
|
||||
|
||||
func (a *Account) CanGetUsage() bool {
|
||||
return a.Type == AccountTypeOAuth
|
||||
}
|
||||
@@ -460,6 +476,7 @@ func (a *Account) GetTempUnschedulableRules() []TempUnschedulableRule {
|
||||
Keywords: parseTempUnschedStrings(entry["keywords"]),
|
||||
DurationMinutes: parseTempUnschedInt(entry["duration_minutes"]),
|
||||
Description: parseTempUnschedString(entry["description"]),
|
||||
Scope: normalizeTempUnschedScope(entry["scope"]),
|
||||
}
|
||||
|
||||
if rule.ErrorCode <= 0 || rule.DurationMinutes <= 0 || len(rule.Keywords) == 0 {
|
||||
@@ -472,6 +489,18 @@ func (a *Account) GetTempUnschedulableRules() []TempUnschedulableRule {
|
||||
return rules
|
||||
}
|
||||
|
||||
// normalizeTempUnschedScope 只认 account / model 两个值,其余一律视为未设置。
|
||||
func normalizeTempUnschedScope(raw any) string {
|
||||
switch strings.ToLower(strings.TrimSpace(parseTempUnschedString(raw))) {
|
||||
case TempUnschedScopeAccount:
|
||||
return TempUnschedScopeAccount
|
||||
case TempUnschedScopeModel:
|
||||
return TempUnschedScopeModel
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func parseTempUnschedString(value any) string {
|
||||
s, ok := value.(string)
|
||||
if !ok {
|
||||
@@ -1575,7 +1604,7 @@ func (a *Account) GetCodingPlanProvider() string {
|
||||
}
|
||||
baseURL := strings.ToLower(a.GetOpenAIBaseURL())
|
||||
switch {
|
||||
case strings.Contains(baseURL, "api.kimi.com/coding"):
|
||||
case isKimiCodingBaseURL(a.GetOpenAIBaseURL()):
|
||||
return PlatformKimi
|
||||
case strings.Contains(baseURL, "bigmodel.cn"), strings.Contains(baseURL, "api.z.ai"):
|
||||
return PlatformZhipu
|
||||
@@ -1588,6 +1617,24 @@ func (a *Account) GetCodingPlanProvider() string {
|
||||
}
|
||||
}
|
||||
|
||||
func isKimiCodingBaseURL(baseURL string) bool {
|
||||
parsed, err := url.Parse(baseURL)
|
||||
if err != nil || parsed.Scheme != "https" || parsed.User != nil || parsed.Opaque != "" ||
|
||||
parsed.RawQuery != "" || parsed.ForceQuery || strings.Contains(baseURL, "#") {
|
||||
return false
|
||||
}
|
||||
host := strings.ToLower(parsed.Host)
|
||||
if host != "api.kimi.com" && host != "api.kimi.ai" && host != "api.kimi.com:443" && host != "api.kimi.ai:443" {
|
||||
return false
|
||||
}
|
||||
switch parsed.EscapedPath() {
|
||||
case "/coding", "/coding/", "/coding/v1", "/coding/v1/":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Account) GetOpenAIAccessToken() string {
|
||||
if !a.IsOpenAI() {
|
||||
return ""
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetCodingPlanProvider_KimiOfficialURLs(t *testing.T) {
|
||||
for _, host := range []string{"api.kimi.com", "api.kimi.ai", "API.KIMI.AI", "api.kimi.ai:443"} {
|
||||
for _, path := range []string{"/coding", "/coding/", "/coding/v1", "/coding/v1/"} {
|
||||
t.Run(host+path, func(t *testing.T) {
|
||||
account := &Account{Platform: PlatformKimi, Type: AccountTypeAPIKey, Credentials: map[string]any{
|
||||
"account_mode": AccountModeCoding,
|
||||
"base_url": "https://" + host + path,
|
||||
}}
|
||||
require.Equal(t, PlatformKimi, account.GetCodingPlanProvider())
|
||||
require.Equal(t, "https://"+host+"/coding/v1/usages", kimiQuotaURL(account.GetOpenAIBaseURL()))
|
||||
account.Credentials["account_mode"] = AccountModePayG
|
||||
require.Empty(t, account.GetCodingPlanProvider())
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetCodingPlanProvider_KimiRejectsUnsafeURLs(t *testing.T) {
|
||||
for _, baseURL := range []string{
|
||||
"https://api.kimi.ai.attacker.example/coding/v1",
|
||||
"https://api.kimi.com.attacker.example/coding/v1",
|
||||
"https://fake-api.kimi.ai/coding/v1",
|
||||
"https://relay.attacker.example/api.kimi.com/coding",
|
||||
"https://relay.attacker.example/?target=api.kimi.ai/coding/v1",
|
||||
"https://relay.attacker.example/#api.kimi.com/coding",
|
||||
"https://api.kimi.ai@relay.attacker.example/coding/v1",
|
||||
"https://user@api.kimi.ai/coding/v1",
|
||||
"https://api.kimi.ai/coding-other",
|
||||
"https://api.kimi.com/coding/v10",
|
||||
"https://api.kimi.ai/coding/v1/usages",
|
||||
"https://api.kimi.ai/v1?target=/coding/v1",
|
||||
"https://api.kimi.ai/coding/v1?target=attacker",
|
||||
"https://api.kimi.ai/coding/v1?",
|
||||
"https://api.kimi.ai/coding/v1#attacker",
|
||||
"https://api.kimi.ai/coding/v1#",
|
||||
"https://api.kimi.ai/%63oding/v1",
|
||||
"https://api.kimi.ai/coding%2fv1",
|
||||
"https://api.kimi.ai/coding/../coding/v1",
|
||||
"https://api.kimi.ai/CODING/v1",
|
||||
"https://api.kimi.ai:8443/coding/v1",
|
||||
"https://api.kimi.ai:/coding/v1",
|
||||
"http://api.kimi.ai/coding/v1",
|
||||
"//api.kimi.ai/coding/v1",
|
||||
"https://api.kimi.ai/%zz",
|
||||
} {
|
||||
t.Run(baseURL, func(t *testing.T) {
|
||||
account := &Account{Platform: PlatformKimi, Type: AccountTypeAPIKey, Credentials: map[string]any{
|
||||
"account_mode": AccountModeCoding,
|
||||
"base_url": baseURL,
|
||||
}}
|
||||
require.Empty(t, account.GetCodingPlanProvider())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetCodingPlanProvider_UnrelatedProvidersUnchanged(t *testing.T) {
|
||||
for _, testCase := range []struct {
|
||||
platform string
|
||||
baseURL string
|
||||
}{
|
||||
{PlatformZhipu, "https://open.bigmodel.cn/api/coding/paas/v4"},
|
||||
{PlatformZhipu, "https://api.z.ai/api/coding/paas/v4"},
|
||||
{PlatformMiniMax, "https://api.minimax.io/v1"},
|
||||
{PlatformMiniMax, "https://api.minimaxi.com/v1"},
|
||||
{PlatformMiniMax, "https://api.minimax.com/v1"},
|
||||
} {
|
||||
t.Run(testCase.baseURL, func(t *testing.T) {
|
||||
account := &Account{Platform: testCase.platform, Type: AccountTypeAPIKey, Credentials: map[string]any{
|
||||
"account_mode": AccountModeCoding,
|
||||
"base_url": testCase.baseURL,
|
||||
}}
|
||||
require.Equal(t, testCase.platform, account.GetCodingPlanProvider())
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -933,7 +933,7 @@ func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "", time.Time{}, true,
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 9.15, *result, 1e-12)
|
||||
require.InDelta(t, 4.575, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) {
|
||||
|
||||
@@ -2585,14 +2585,7 @@ func (s *AccountTestService) buildGeminiOAuthRequest(ctx context.Context, accoun
|
||||
}
|
||||
|
||||
func (s *AccountTestService) buildGeminiServiceAccountRequest(ctx context.Context, account *Account, modelID string, payload []byte) (*http.Request, error) {
|
||||
if s.geminiTokenProvider == nil {
|
||||
return nil, fmt.Errorf("gemini token provider not configured")
|
||||
}
|
||||
accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get service account access token: %w", err)
|
||||
}
|
||||
fullURL, err := buildVertexGeminiURL(account.VertexProjectID(), account.VertexLocation(modelID), modelID, "streamGenerateContent", true)
|
||||
fullURL, err := buildVertexGeminiAccountURL(account, modelID, "streamGenerateContent", true)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -2601,7 +2594,9 @@ func (s *AccountTestService) buildGeminiServiceAccountRequest(ctx context.Contex
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
if err := setVertexGeminiAuth(ctx, req, account, s.geminiTokenProvider); err != nil {
|
||||
return nil, fmt.Errorf("failed to authenticate vertex request: %w", err)
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -372,7 +372,7 @@ func normalizeCreateGroupInputForSimpleMode(input *CreateGroupInput) {
|
||||
return
|
||||
}
|
||||
*input = CreateGroupInput{
|
||||
Name: input.Name, Description: input.Description, Platform: input.Platform,
|
||||
Name: input.Name, Description: input.Description, I18n: input.I18n, Platform: input.Platform,
|
||||
RateMultiplier: 1, SubscriptionType: SubscriptionTypeStandard,
|
||||
}
|
||||
}
|
||||
@@ -381,7 +381,7 @@ func normalizeUpdateGroupInputForSimpleMode(input *UpdateGroupInput) {
|
||||
if input == nil {
|
||||
return
|
||||
}
|
||||
*input = UpdateGroupInput{Name: input.Name, Description: input.Description}
|
||||
*input = UpdateGroupInput{Name: input.Name, Description: input.Description, I18n: input.I18n}
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupInput) (*Group, error) {
|
||||
@@ -394,6 +394,10 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
|
||||
if input.RateMultiplier <= 0 {
|
||||
return nil, errors.New("rate_multiplier must be > 0")
|
||||
}
|
||||
i18n, err := normalizeGroupI18n(input.I18n)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
platform := NormalizeGroupPlatform(input.Platform)
|
||||
// 固定账号 manifest 配置:账号绑定发生在创建之后,创建时无法校验成员关系,
|
||||
@@ -565,6 +569,7 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
|
||||
group := &Group{
|
||||
Name: input.Name,
|
||||
Description: input.Description,
|
||||
I18n: i18n,
|
||||
Platform: platform,
|
||||
RateMultiplier: input.RateMultiplier,
|
||||
IsExclusive: input.IsExclusive,
|
||||
@@ -776,6 +781,13 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
|
||||
if input.Description != nil {
|
||||
group.Description = *input.Description
|
||||
}
|
||||
if input.I18n != nil {
|
||||
i18n, err := normalizeGroupI18n(*input.I18n)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
group.I18n = i18n
|
||||
}
|
||||
if input.Platform != "" {
|
||||
group.Platform = input.Platform
|
||||
}
|
||||
|
||||
@@ -97,6 +97,7 @@ func cloneGroupForDuplicate(source *Group, operationID string) *Group {
|
||||
return &Group{
|
||||
Name: duplicateGroupName(source.Name, 1),
|
||||
Description: source.Description,
|
||||
I18n: cloneGroupI18nForDuplicate(source.I18n),
|
||||
Platform: source.Platform,
|
||||
RateMultiplier: source.RateMultiplier,
|
||||
PeakRateEnabled: source.PeakRateEnabled,
|
||||
|
||||
@@ -235,6 +235,7 @@ type AdminBoundAuthIdentityChannel struct {
|
||||
type CreateGroupInput struct {
|
||||
Name string
|
||||
Description string
|
||||
I18n GroupI18n // 名称与描述的界面语言译文
|
||||
Platform string
|
||||
RateMultiplier float64
|
||||
IsExclusive bool
|
||||
@@ -315,6 +316,7 @@ type CreateGroupInput struct {
|
||||
type UpdateGroupInput struct {
|
||||
Name string
|
||||
Description *string
|
||||
I18n *GroupI18n // nil 表示不修改;非 nil 整体替换
|
||||
Platform string
|
||||
RateMultiplier *float64 // 使用指针以支持设置为0
|
||||
IsExclusive *bool
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAdminService_CreateGroup_NormalizesI18n(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
|
||||
Name: "Claude Code - 标准",
|
||||
Platform: PlatformAnthropic,
|
||||
RateMultiplier: 1.0,
|
||||
I18n: GroupI18n{
|
||||
"en": {Name: " Claude Code - Standard ", Description: "Pro + Max pool"},
|
||||
"zh": {},
|
||||
},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, repo.created)
|
||||
require.Equal(t, GroupI18n{
|
||||
"en": {Name: "Claude Code - Standard", Description: "Pro + Max pool"},
|
||||
}, repo.created.I18n)
|
||||
}
|
||||
|
||||
func TestAdminService_CreateGroup_RejectsInvalidI18nLocale(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
|
||||
Name: "g",
|
||||
Platform: PlatformAnthropic,
|
||||
RateMultiplier: 1.0,
|
||||
I18n: GroupI18n{"en_US": {Name: "x"}},
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
appErr := infraerrors.FromError(err)
|
||||
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
|
||||
require.Equal(t, "INVALID_GROUP_I18N", appErr.Reason)
|
||||
require.Nil(t, repo.created)
|
||||
}
|
||||
|
||||
func TestAdminService_CreateGroup_RejectsCaseCollidingI18nLocales(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
|
||||
Name: "默认",
|
||||
Platform: PlatformAnthropic,
|
||||
RateMultiplier: 1.0,
|
||||
I18n: GroupI18n{"en": {Name: "English A"}, "EN": {Name: "English B"}},
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
appErr := infraerrors.FromError(err)
|
||||
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
|
||||
require.Equal(t, "INVALID_GROUP_I18N", appErr.Reason)
|
||||
require.Nil(t, repo.created)
|
||||
}
|
||||
|
||||
func TestAdminService_UpdateGroup_I18n(t *testing.T) {
|
||||
newExisting := func() *Group {
|
||||
return &Group{
|
||||
ID: 1, Name: "标准", Platform: PlatformAnthropic, Status: StatusActive, RateMultiplier: 1,
|
||||
I18n: GroupI18n{"en": {Name: "Standard"}},
|
||||
}
|
||||
}
|
||||
|
||||
t.Run("omitted i18n keeps stored translations", func(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{getByID: newExisting()}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
_, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{Name: "标准版"})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "标准版", repo.updated.Name)
|
||||
require.Equal(t, GroupI18n{"en": {Name: "Standard"}}, repo.updated.I18n)
|
||||
})
|
||||
|
||||
t.Run("provided i18n replaces stored translations as a whole", func(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{getByID: newExisting()}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
next := GroupI18n{"zh": {Description: " 号池 "}}
|
||||
|
||||
_, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{I18n: &next})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, GroupI18n{"zh": {Description: "号池"}}, repo.updated.I18n)
|
||||
})
|
||||
|
||||
t.Run("empty i18n clears stored translations", func(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{getByID: newExisting()}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
next := GroupI18n{}
|
||||
|
||||
_, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{I18n: &next})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, repo.updated.I18n)
|
||||
require.Empty(t, repo.updated.I18n)
|
||||
})
|
||||
|
||||
t.Run("invalid i18n is rejected before persisting", func(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{getByID: newExisting()}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
next := GroupI18n{"not a locale": {Name: "x"}}
|
||||
|
||||
_, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{I18n: &next})
|
||||
|
||||
require.Error(t, err)
|
||||
require.Equal(t, "INVALID_GROUP_I18N", infraerrors.FromError(err).Reason)
|
||||
require.Nil(t, repo.updated)
|
||||
})
|
||||
|
||||
// 简易模式只放行名称与描述,译文属于同一类展示文案,同样放行。
|
||||
t.Run("simple mode keeps i18n editable", func(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{getByID: newExisting()}
|
||||
svc := &adminServiceImpl{cfg: &config.Config{RunMode: config.RunModeSimple}, groupRepo: repo}
|
||||
next := GroupI18n{"en": {Name: "Standard v2"}}
|
||||
rate := 9.0
|
||||
|
||||
_, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{I18n: &next, RateMultiplier: &rate})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, GroupI18n{"en": {Name: "Standard v2"}}, repo.updated.I18n)
|
||||
require.Equal(t, 1.0, repo.updated.RateMultiplier)
|
||||
})
|
||||
}
|
||||
|
||||
func TestCloneGroupForDuplicate_DropsNameTranslations(t *testing.T) {
|
||||
source := &Group{
|
||||
ID: 7, Name: "标准", Platform: PlatformAnthropic,
|
||||
I18n: GroupI18n{"en": {Name: "Standard", Description: "Pro + Max pool"}},
|
||||
}
|
||||
|
||||
clone := cloneGroupForDuplicate(source, "op-1")
|
||||
|
||||
require.Equal(t, GroupI18n{"en": {Description: "Pro + Max pool"}}, clone.I18n)
|
||||
}
|
||||
@@ -844,6 +844,28 @@ func TestSetAntigravityModelRateLimits_GeminiWritesFamilyScope(t *testing.T) {
|
||||
require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
|
||||
}
|
||||
|
||||
func TestSetAntigravityModelRateLimits_ImageDoesNotWriteFamilyScope(t *testing.T) {
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
svc := &AntigravityGatewayService{}
|
||||
account := &Account{ID: 789, Platform: PlatformAntigravity}
|
||||
resetAt := time.Now().Add(30 * time.Second)
|
||||
|
||||
success := svc.setAntigravityModelRateLimits(
|
||||
context.Background(),
|
||||
repo,
|
||||
account,
|
||||
"gemini-3.1-flash-image-preview",
|
||||
"[test]",
|
||||
429,
|
||||
resetAt,
|
||||
false,
|
||||
)
|
||||
|
||||
require.True(t, success)
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "gemini-3.1-flash-image-preview", repo.modelRateLimitCalls[0].modelKey)
|
||||
}
|
||||
|
||||
func TestSetAntigravityModelRateLimits_DoesNotDoubleMapCustomChain(t *testing.T) {
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
svc := &AntigravityGatewayService{}
|
||||
|
||||
@@ -130,6 +130,8 @@ type apiKeyAllByUserIDLister interface {
|
||||
|
||||
// APIKeyRateLimitData holds rate limit usage and window state for an API key.
|
||||
type APIKeyRateLimitData struct {
|
||||
// QuotaUsed 是总额度已用量(quota_used)。
|
||||
QuotaUsed float64
|
||||
Usage5h float64
|
||||
Usage1d float64
|
||||
Usage7d float64
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -573,6 +574,11 @@ func (s *BillingCacheService) InvalidateAPIKeyRateLimit(ctx context.Context, key
|
||||
// resets expired windows in-memory and triggers async DB reset,
|
||||
// and returns an error if any window limit is exceeded.
|
||||
func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey *APIKey) error {
|
||||
return s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, 0)
|
||||
}
|
||||
|
||||
// checkAPIKeyRateLimitsWithPending 在窗口用量上叠加 pending(即将扣除的金额)后检查限速窗口。
|
||||
func (s *BillingCacheService) checkAPIKeyRateLimitsWithPending(ctx context.Context, apiKey *APIKey, pending float64) error {
|
||||
if s.cache == nil {
|
||||
// No cache: fall back to reading from DB directly
|
||||
if s.apiKeyRateLimitLoader == nil {
|
||||
@@ -583,7 +589,7 @@ func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey
|
||||
return nil // Don't block requests on DB errors
|
||||
}
|
||||
return s.evaluateRateLimits(ctx, apiKey, data.Usage5h, data.Usage1d, data.Usage7d,
|
||||
data.Window5hStart, data.Window1dStart, data.Window7dStart)
|
||||
data.Window5hStart, data.Window1dStart, data.Window7dStart, pending)
|
||||
}
|
||||
|
||||
cacheData, err := s.cache.GetAPIKeyRateLimit(ctx, apiKey.ID)
|
||||
@@ -628,11 +634,11 @@ func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey
|
||||
t := time.Unix(cacheData.Window7d, 0)
|
||||
w7d = &t
|
||||
}
|
||||
return s.evaluateRateLimits(ctx, apiKey, cacheData.Usage5h, cacheData.Usage1d, cacheData.Usage7d, w5h, w1d, w7d)
|
||||
return s.evaluateRateLimits(ctx, apiKey, cacheData.Usage5h, cacheData.Usage1d, cacheData.Usage7d, w5h, w1d, w7d, pending)
|
||||
}
|
||||
|
||||
// evaluateRateLimits checks usage against limits, triggering async resets for expired windows.
|
||||
func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *APIKey, usage5h, usage1d, usage7d float64, w5h, w1d, w7d *time.Time) error {
|
||||
// evaluateRateLimits checks usage (plus pending) against limits, triggering async resets for expired windows.
|
||||
func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *APIKey, usage5h, usage1d, usage7d float64, w5h, w1d, w7d *time.Time, pending float64) error {
|
||||
needsReset := false
|
||||
|
||||
// Reset expired windows in-memory for check purposes
|
||||
@@ -675,13 +681,16 @@ func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *AP
|
||||
}
|
||||
|
||||
// Check limits
|
||||
if apiKey.RateLimit5h > 0 && usage5h >= apiKey.RateLimit5h {
|
||||
exceeds := func(usage, limit float64) bool {
|
||||
return limit > 0 && (usage >= limit || usage+pending > limit)
|
||||
}
|
||||
if exceeds(usage5h, apiKey.RateLimit5h) {
|
||||
return ErrAPIKeyRateLimit5hExceeded
|
||||
}
|
||||
if apiKey.RateLimit1d > 0 && usage1d >= apiKey.RateLimit1d {
|
||||
if exceeds(usage1d, apiKey.RateLimit1d) {
|
||||
return ErrAPIKeyRateLimit1dExceeded
|
||||
}
|
||||
if apiKey.RateLimit7d > 0 && usage7d >= apiKey.RateLimit7d {
|
||||
if exceeds(usage7d, apiKey.RateLimit7d) {
|
||||
return ErrAPIKeyRateLimit7dExceeded
|
||||
}
|
||||
return nil
|
||||
@@ -781,6 +790,110 @@ func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReserveFixedCharge 为扣费前即可确定金额的费用(倍率后 USD,如显式缓存存储费)做准入并登记在途预留。
|
||||
//
|
||||
// 余额模式要求缓存余额覆盖 amount 及用户的其他在途预留;订阅模式要求日 / 周 / 月剩余额度覆盖 amount;
|
||||
// API Key 总额度与限速窗口同样叠加 amount。订阅与 API Key 的用量取库中权威值,并叠加同一订阅 / Key 上其他
|
||||
// 在途预留:先登记本次预留再读用量,并发的两笔费用中后读用量的一方必然看到另一方(预留或已落库的扣费)。
|
||||
// 返回的句柄可能为 nil(未登记任何预留);调用方须在 handler 结束时调用 HandlerDone,预留在计费任务落库后释放。
|
||||
// 须在 CheckBillingEligibility 通过后调用;简易模式不校验。
|
||||
func (s *BillingCacheService) ReserveFixedCharge(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, amount float64) (*InflightReservation, error) {
|
||||
if user == nil || amount <= 0 || math.IsNaN(amount) || math.IsInf(amount, 0) {
|
||||
return nil, nil
|
||||
}
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
||||
return nil, nil
|
||||
}
|
||||
subscriptionBill := group != nil && group.IsSubscriptionType() && subscription != nil
|
||||
if !subscriptionBill {
|
||||
balance, err := s.GetUserBalance(ctx, user.ID)
|
||||
if err != nil {
|
||||
return nil, ErrBillingServiceUnavailable.WithCause(err)
|
||||
}
|
||||
if balance < amount {
|
||||
return nil, ErrInsufficientBalance
|
||||
}
|
||||
}
|
||||
res, err := s.reserveInflight(ctx, user, group, subscription, amount, false, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var subscriptionPending, apiKeyPending float64
|
||||
if subscriptionBill && (group.HasDailyLimit() || group.HasWeeklyLimit() || group.HasMonthlyLimit()) {
|
||||
res, subscriptionPending = s.reserveInflightScope(ctx, res, user.ID, inflightSubscriptionScope(subscription.ID), amount)
|
||||
}
|
||||
if apiKey != nil && (apiKey.Quota > 0 || apiKey.HasRateLimits()) {
|
||||
res, apiKeyPending = s.reserveInflightScope(ctx, res, user.ID, inflightAPIKeyScope(apiKey.ID), amount)
|
||||
}
|
||||
if subscriptionBill {
|
||||
err = s.checkFixedChargeSubscription(ctx, user.ID, group, amount+subscriptionPending)
|
||||
}
|
||||
if err == nil && apiKey != nil {
|
||||
err = s.checkFixedChargeAPIKey(ctx, apiKey, amount+apiKeyPending)
|
||||
}
|
||||
if err != nil {
|
||||
res.Release()
|
||||
return nil, err
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// checkFixedChargeSubscription 按库中订阅用量校验日 / 周 / 月剩余额度覆盖 pending。
|
||||
func (s *BillingCacheService) checkFixedChargeSubscription(ctx context.Context, userID int64, group *Group, pending float64) error {
|
||||
read := s.GetSubscriptionStatus
|
||||
if s.subRepo != nil {
|
||||
read = s.getSubscriptionFromDB
|
||||
}
|
||||
subData, err := read(ctx, userID, group.ID)
|
||||
if err != nil {
|
||||
return ErrBillingServiceUnavailable.WithCause(err)
|
||||
}
|
||||
if group.HasDailyLimit() && subData.DailyUsage+pending > *group.DailyLimitUSD {
|
||||
return ErrDailyLimitExceeded
|
||||
}
|
||||
if group.HasWeeklyLimit() && subData.WeeklyUsage+pending > *group.WeeklyLimitUSD {
|
||||
return ErrWeeklyLimitExceeded
|
||||
}
|
||||
if group.HasMonthlyLimit() && subData.MonthlyUsage+pending > *group.MonthlyLimitUSD {
|
||||
return ErrMonthlyLimitExceeded
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkFixedChargeAPIKey 按库中 Key 用量校验总额度与限速窗口覆盖 pending;未接入仓储时使用请求上下文与限速缓存中的用量。
|
||||
func (s *BillingCacheService) checkFixedChargeAPIKey(ctx context.Context, apiKey *APIKey, pending float64) error {
|
||||
if apiKey.Quota <= 0 && !apiKey.HasRateLimits() {
|
||||
return nil
|
||||
}
|
||||
if s.apiKeyRateLimitLoader == nil {
|
||||
if apiKey.Quota > 0 && apiKey.QuotaUsed+pending > apiKey.Quota {
|
||||
return ErrAPIKeyQuotaExhausted
|
||||
}
|
||||
if apiKey.HasRateLimits() {
|
||||
return s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, pending)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
data, err := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID)
|
||||
if err != nil {
|
||||
return ErrBillingServiceUnavailable.WithCause(err)
|
||||
}
|
||||
exceeds := func(usage, limit float64) bool {
|
||||
return limit > 0 && usage+pending > limit
|
||||
}
|
||||
switch {
|
||||
case exceeds(data.QuotaUsed, apiKey.Quota):
|
||||
return ErrAPIKeyQuotaExhausted
|
||||
case exceeds(data.EffectiveUsage5h(), apiKey.RateLimit5h):
|
||||
return ErrAPIKeyRateLimit5hExceeded
|
||||
case exceeds(data.EffectiveUsage1d(), apiKey.RateLimit1d):
|
||||
return ErrAPIKeyRateLimit1dExceeded
|
||||
case exceeds(data.EffectiveUsage7d(), apiKey.RateLimit7d):
|
||||
return ErrAPIKeyRateLimit7dExceeded
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkSimpleModeAPIKeyRateLimits is deliberately DB-authoritative. Redis
|
||||
// updates are asynchronous and can be dropped or missed after a committed
|
||||
// transaction, so using the cache here could let a limited simple-mode key
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -25,6 +26,16 @@ type InflightBalanceReservationCache interface {
|
||||
ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error
|
||||
}
|
||||
|
||||
// InflightScopeReservationCache 订阅 / API Key 维度在途预留的缓存能力(可选),用于扣费前即可确定金额的费用;
|
||||
// 未实现时这两个维度不登记预留(fail-open),准入仍按库中用量校验。
|
||||
type InflightScopeReservationCache interface {
|
||||
// ReserveInflightScope 原子地清理 scope 下的过期预留并登记 requestID 的预留(ttl 后自动失效),
|
||||
// 返回登记前该 scope 下其他在途预留的合计。
|
||||
ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error)
|
||||
// ReleaseInflightScope 释放 scope 下 requestID 的预留(幂等)。
|
||||
ReleaseInflightScope(ctx context.Context, scope, requestID string) error
|
||||
}
|
||||
|
||||
// InflightBalanceReservationRenewer 可选:续期仍存活的预留(流式长请求期间防止 TTL 过期)。
|
||||
type InflightBalanceReservationRenewer interface {
|
||||
// RenewInflightBalance 若 requestID 仍存在则把其过期时间推迟到 now+ttl;返回是否仍存在。
|
||||
@@ -54,11 +65,15 @@ const (
|
||||
// 不再续期,所以最长持有时间受 TTL 约束(计费任务卡死时预留自动过期)。
|
||||
// 所有方法对 nil 接收者安全。
|
||||
type InflightReservation struct {
|
||||
// cache 为 nil 表示未登记余额维度的预留。
|
||||
cache InflightBalanceReservationCache
|
||||
userID int64
|
||||
requestID string
|
||||
amount float64
|
||||
ttl time.Duration
|
||||
// scopes 是同一金额在订阅 / API Key 维度登记的预留(与余额维度共用 requestID),不续期。
|
||||
scopeCache InflightScopeReservationCache
|
||||
scopes []string
|
||||
|
||||
refs atomic.Int64
|
||||
releaseOnce sync.Once
|
||||
@@ -116,8 +131,15 @@ func (r *InflightReservation) Release() {
|
||||
r.refs.Store(0)
|
||||
relCtx, relCancel := context.WithTimeout(context.Background(), inflightReservationReleaseTimeout)
|
||||
defer relCancel()
|
||||
if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil {
|
||||
logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err)
|
||||
if r.cache != nil {
|
||||
if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil {
|
||||
logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err)
|
||||
}
|
||||
}
|
||||
for _, scope := range r.scopes {
|
||||
if err := r.scopeCache.ReleaseInflightScope(relCtx, scope, r.requestID); err != nil {
|
||||
logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for %s (expires by ttl): %v", scope, err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -213,7 +235,7 @@ func (s *BillingCacheService) InflightReservationFailClosedOnUnpriced() bool {
|
||||
|
||||
// ReserveInflightBalance 简化封装:返回释放函数,不续期、不做计费交接(预留最长存活 TTL)。
|
||||
func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (func(), error) {
|
||||
r, err := s.reserveInflight(ctx, user, group, subscription, estimate, false)
|
||||
r, err := s.reserveInflight(ctx, user, group, subscription, estimate, false, true)
|
||||
if err != nil {
|
||||
return noopRelease, err
|
||||
}
|
||||
@@ -232,10 +254,50 @@ func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user *
|
||||
// 开关关闭 / 简易模式 / 订阅模式 / estimate <= 0 / 缓存不支持 / 余额读取失败 / Redis 执行失败。
|
||||
// 仅当 Redis 明确判定 缓存余额 - 在途合计 < estimate(且已有在途请求)时返回 ErrInsufficientBalance。
|
||||
func (s *BillingCacheService) ReserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (*InflightReservation, error) {
|
||||
return s.reserveInflight(ctx, user, group, subscription, estimate, true)
|
||||
return s.reserveInflight(ctx, user, group, subscription, estimate, true, true)
|
||||
}
|
||||
|
||||
func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew bool) (*InflightReservation, error) {
|
||||
// inflightSubscriptionScope / inflightAPIKeyScope 是订阅 / API Key 维度在途预留的 scope 名。
|
||||
func inflightSubscriptionScope(subscriptionID int64) string {
|
||||
return "sub:" + strconv.FormatInt(subscriptionID, 10)
|
||||
}
|
||||
|
||||
func inflightAPIKeyScope(apiKeyID int64) string {
|
||||
return "apikey:" + strconv.FormatInt(apiKeyID, 10)
|
||||
}
|
||||
|
||||
// reserveInflightScope 在 scope 维度为 amount 登记在途预留并挂到 res 上(res 为 nil 时新建句柄),
|
||||
// 返回句柄与登记前该 scope 下其他在途预留的合计。开关关闭 / 缓存不支持 / Redis 执行失败时不登记(fail-open),合计按 0。
|
||||
func (s *BillingCacheService) reserveInflightScope(ctx context.Context, res *InflightReservation, userID int64, scope string, amount float64) (*InflightReservation, float64) {
|
||||
cfg, ok := s.inflightReservationConfig()
|
||||
if !ok {
|
||||
return res, 0
|
||||
}
|
||||
sc, ok := s.cache.(InflightScopeReservationCache)
|
||||
if !ok || sc == nil {
|
||||
return res, 0
|
||||
}
|
||||
if res == nil {
|
||||
ttl := defaultInflightReservationTTL
|
||||
if cfg.TTLSeconds > 0 {
|
||||
ttl = time.Duration(cfg.TTLSeconds) * time.Second
|
||||
}
|
||||
res = &InflightReservation{userID: userID, requestID: uuid.NewString(), amount: amount, ttl: ttl}
|
||||
res.refs.Store(1)
|
||||
}
|
||||
reserveCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), inflightReservationReserveTimeout)
|
||||
pending, err := sc.ReserveInflightScope(reserveCtx, scope, res.requestID, amount, res.ttl)
|
||||
cancel()
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation failed for %s (fail-open): %v", scope, err)
|
||||
return res, 0
|
||||
}
|
||||
res.scopeCache = sc
|
||||
res.scopes = append(res.scopes, scope)
|
||||
return res, pending
|
||||
}
|
||||
|
||||
func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew, capped bool) (*InflightReservation, error) {
|
||||
cfg, ok := s.inflightReservationConfig()
|
||||
if !ok || user == nil {
|
||||
return nil, nil
|
||||
@@ -243,7 +305,7 @@ func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, g
|
||||
if group != nil && group.IsSubscriptionType() && subscription != nil {
|
||||
return nil, nil
|
||||
}
|
||||
if cfg.MaxReservationUSD > 0 && estimate > cfg.MaxReservationUSD {
|
||||
if capped && cfg.MaxReservationUSD > 0 && estimate > cfg.MaxReservationUSD {
|
||||
estimate = cfg.MaxReservationUSD
|
||||
}
|
||||
if estimate <= 0 || math.IsNaN(estimate) || math.IsInf(estimate, 0) {
|
||||
|
||||
@@ -444,6 +444,20 @@ func (s *BillingService) initFallbackPricing() {
|
||||
CacheCreation1hPrice: 4e-6,
|
||||
SupportsCacheBreakdown: true,
|
||||
}
|
||||
// Claude Haiku 5.5(2026-10-07,$0.10/$0.50,prompt 超 100k tokens 整单 5x)。
|
||||
// 长文本阶梯与目录 *_above_100k_tokens 绝对价同口径,缺目录时回退也不丢阶梯。
|
||||
s.fallbackPrices["claude-haiku-5-5"] = &ModelPricing{
|
||||
InputPricePerToken: 0.1e-6,
|
||||
OutputPricePerToken: 0.5e-6,
|
||||
CacheCreationPricePerToken: 0.125e-6,
|
||||
CacheReadPricePerToken: 0.01e-6,
|
||||
CacheCreation5mPrice: 0.125e-6,
|
||||
CacheCreation1hPrice: 0.2e-6,
|
||||
SupportsCacheBreakdown: true,
|
||||
LongContextInputThreshold: 100000,
|
||||
LongContextInputMultiplier: 5,
|
||||
LongContextOutputMultiplier: 5,
|
||||
}
|
||||
// Claude Fable 5.x uses the same input/output and cache-write prices, while
|
||||
// Fable 5.1 reduces cache reads from $1 to $0.25 per MTok.
|
||||
s.fallbackPrices["claude-fable-5"] = &ModelPricing{
|
||||
@@ -474,14 +488,15 @@ func (s *BillingService) initFallbackPricing() {
|
||||
SupportsCacheBreakdown: false,
|
||||
}
|
||||
|
||||
// Gemini 3.6 Flash (Google AI pricing: $1.50 input / $7.50 output /
|
||||
// $0.15 cached input per MTok). Antigravity's -high/-low/-medium/-tiered
|
||||
// aliases are matched below so unavailable remote pricing never records
|
||||
// token-bearing requests at $0.
|
||||
// Gemini 3.6 Flash (Google AI pricing: $0.75 input / $3.75 output /
|
||||
// $0.075 cached input per MTok, promotional through 2026-12-31; official
|
||||
// rates double to $1.50/$7.50/$0.15 from 2027-01-01). Antigravity's
|
||||
// -high/-low/-medium/-tiered aliases are matched below so unavailable
|
||||
// remote pricing never records token-bearing requests at $0.
|
||||
s.fallbackPrices["gemini-3.6-flash"] = &ModelPricing{
|
||||
InputPricePerToken: 1.5e-6,
|
||||
OutputPricePerToken: 7.5e-6,
|
||||
CacheReadPricePerToken: 0.15e-6,
|
||||
InputPricePerToken: 0.75e-6,
|
||||
OutputPricePerToken: 3.75e-6,
|
||||
CacheReadPricePerToken: 0.075e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
}
|
||||
|
||||
@@ -1017,6 +1032,9 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing {
|
||||
if claude.IsSonnet55(modelLower) {
|
||||
return s.fallbackPrices["claude-sonnet-5-5"]
|
||||
}
|
||||
if claude.IsHaiku55(modelLower) {
|
||||
return s.fallbackPrices["claude-haiku-5-5"]
|
||||
}
|
||||
if strings.Contains(modelLower, "opus") {
|
||||
// "opus-5" 必须先判:不能用裸 "5" 匹配,否则 claude-opus-4-5 会被误判。
|
||||
if strings.Contains(modelLower, "opus-5") || strings.Contains(modelLower, "opus5") {
|
||||
@@ -1325,6 +1343,18 @@ func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) {
|
||||
return s.getModelPricingAt(model, timezone.Now())
|
||||
}
|
||||
|
||||
// CacheStoragePricePerTokenHour 返回模型的显式上下文缓存存储单价(每 token 每小时,USD)。
|
||||
// 优先级:价格表(override_file 等)按确定性识别配置的值 > 内置型号价 > 内置 Pro/Flash 系列兜底价;
|
||||
// 均未命中返回 0。
|
||||
func (s *BillingService) CacheStoragePricePerTokenHour(model string) float64 {
|
||||
if s != nil && s.pricingService != nil {
|
||||
if pricing := s.pricingService.GetIdentifiedModelPricing(model); pricing != nil && pricing.CacheStorageCostPerTokenPerHour > 0 {
|
||||
return pricing.CacheStorageCostPerTokenPerHour
|
||||
}
|
||||
}
|
||||
return builtinGeminiCacheStoragePrice(model)
|
||||
}
|
||||
|
||||
// getModelPricingAt 是 GetModelPricing 的带计费时点内部变体:pricingAt 显式
|
||||
// 驱动 DeepSeek pro→Flash 切换判定(切换点前 Pro 价、之后 Flash 价),使
|
||||
// 展示/估算路径可与历史补账同刻复算,测试也能用固定时点钉住断言。
|
||||
@@ -1476,6 +1506,8 @@ type CostInput struct {
|
||||
Resolver *ModelPricingResolver // 定价解析器
|
||||
Resolved *ResolvedPricing // 可选:预解析的定价结果(避免重复 Resolve 调用)
|
||||
LongContextBillingEnabled *bool
|
||||
// DisableModelLongContextPricing 豁免模型价卡的长上下文附加费,不改变自定义区间价格。
|
||||
DisableModelLongContextPricing bool
|
||||
}
|
||||
|
||||
// CalculateCostUnified 统一计费入口,支持三种计费模式。
|
||||
@@ -1491,7 +1523,7 @@ func (s *BillingService) CalculateCostUnified(input CostInput) (*CostBreakdown,
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
breakdown := s.computeTokenBreakdown(pricing, input.Tokens, input.RateMultiplier, input.ServiceTier, applyLongContextBilling)
|
||||
breakdown := s.computeTokenBreakdown(pricing, input.Tokens, input.RateMultiplier, input.ServiceTier, applyLongContextBilling && !input.DisableModelLongContextPricing)
|
||||
applyCostBreakdownMultiplier(breakdown, reasoningEffortBillingMultiplier(input.ReasoningEffort, pricing.ReasoningEffortMultipliers))
|
||||
return breakdown, nil
|
||||
}
|
||||
@@ -1578,7 +1610,7 @@ func (s *BillingService) calculateTokenCost(resolved *ResolvedPricing, input Cos
|
||||
}
|
||||
|
||||
// 官方长上下文阶梯仅在无区间定价时应用(区间定价已包含上下文分层)。
|
||||
applyLongCtx := len(resolved.Intervals) == 0 && contextTierPricingEnabled
|
||||
applyLongCtx := len(resolved.Intervals) == 0 && contextTierPricingEnabled && !input.DisableModelLongContextPricing
|
||||
|
||||
breakdown := s.computeTokenBreakdown(pricing, input.Tokens, input.RateMultiplier, input.ServiceTier, applyLongCtx)
|
||||
applyCostBreakdownMultiplier(breakdown, resolvedChannelTimeMultiplier(resolved, input.PricingAt))
|
||||
|
||||
@@ -2055,6 +2055,33 @@ func TestNewModelPricingCatalogFallbackAndContext(t *testing.T) {
|
||||
require.False(t, cost.LongContextBillingApplied)
|
||||
})
|
||||
}
|
||||
// claude-haiku-5-5:prompt 超 100k tokens 整单 5x(官方定价),
|
||||
// billing/catalog/fallback 三源口径一致(issue #7919)。
|
||||
for _, model := range []string{
|
||||
"claude-haiku-5-5",
|
||||
"anthropic/claude-haiku-5.5",
|
||||
} {
|
||||
t.Run(source+"/"+model, func(t *testing.T) {
|
||||
for _, n := range []int{99999, 100000, 100001} {
|
||||
tokens := UsageTokens{
|
||||
InputTokens: n - 3000, OutputTokens: 500,
|
||||
CacheReadTokens: 2000, CacheCreationTokens: 1000,
|
||||
}
|
||||
cost, err := svc.CalculateCost(model, tokens, 1)
|
||||
require.NoError(t, err)
|
||||
im, om := 1.0, 1.0
|
||||
if n > 100000 {
|
||||
im, om = 5.0, 5.0
|
||||
}
|
||||
require.InDelta(t, float64(tokens.InputTokens)*1e-7*im, cost.InputCost, 1e-10)
|
||||
require.InDelta(t, 1000*1.25e-7*im, cost.CacheCreationCost, 1e-10)
|
||||
require.InDelta(t, 2000*1e-8*im, cost.CacheReadCost, 1e-10)
|
||||
require.InDelta(t, 500*5e-7*om, cost.OutputCost, 1e-10)
|
||||
require.InDelta(t, cost.InputCost+cost.OutputCost+cost.CacheCreationCost+cost.CacheReadCost, cost.TotalCost, 1e-10)
|
||||
require.Equal(t, n > 100000, cost.LongContextBillingApplied)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
type AvailableGroupRef struct {
|
||||
ID int64
|
||||
Name string
|
||||
I18n GroupI18n
|
||||
Platform string
|
||||
SubscriptionType string
|
||||
RateMultiplier float64
|
||||
@@ -66,6 +67,7 @@ func (s *ChannelService) ListAvailable(ctx context.Context) ([]AvailableChannel,
|
||||
groupByID[g.ID] = AvailableGroupRef{
|
||||
ID: g.ID,
|
||||
Name: g.Name,
|
||||
I18n: g.I18n,
|
||||
Platform: g.Platform,
|
||||
SubscriptionType: g.SubscriptionType,
|
||||
RateMultiplier: g.RateMultiplier,
|
||||
|
||||
@@ -636,6 +636,8 @@ var monitorAPIKeyPatterns = []struct {
|
||||
{regexp.MustCompile(`xai-[A-Za-z0-9_-]{6,}`), "xai-***REDACTED***"},
|
||||
// Gemini / Google API Key:固定前缀 + 35 位
|
||||
{regexp.MustCompile(`AIza[A-Za-z0-9_-]{35}`), "AIza***REDACTED***"},
|
||||
// Gemini auth key:AQ.Ab 前缀 + 48 位,长度放宽以覆盖格式变体
|
||||
{regexp.MustCompile(`AQ\.Ab[A-Za-z0-9_-]{40,}`), "AQ.***REDACTED***"},
|
||||
// JWT 三段式(Bearer 后常出现):eyJxxx.eyJxxx.signature
|
||||
{regexp.MustCompile(`eyJ[A-Za-z0-9_-]{8,}\.eyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}`), "eyJ***REDACTED.JWT***"},
|
||||
}
|
||||
@@ -643,7 +645,7 @@ var monitorAPIKeyPatterns = []struct {
|
||||
// sanitizeErrorMessage 擦除错误/响应文本中可能泄露的 API key。
|
||||
// 处理两类来源:
|
||||
// 1. URL query 中的 ?key= / ?api_key= 等(Go *url.Error 会回填完整 URL)
|
||||
// 2. 上游 HTTP body 文本里直接出现的 sk-* / xai-* / AIza* / JWT 等密钥碎片
|
||||
// 2. 上游 HTTP body 文本里直接出现的 sk-* / xai-* / AIza* / AQ.* / JWT 等密钥碎片
|
||||
//
|
||||
// 注意:与 gemini_messages_compat_service.go 的 sanitizeUpstreamErrorMessage 关注点类似但参数集更广,
|
||||
// 监控模块独立维护,避免互相耦合。
|
||||
|
||||
@@ -279,6 +279,17 @@ func TestRunCheckForModel_Grok_RedactsXAIKeyFromUpstreamBody(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeErrorMessage_RedactsGeminiAuthKey(t *testing.T) {
|
||||
authKey := "AQ." + "Ab" + strings.Repeat("Xy9_-z", 8)
|
||||
got := sanitizeErrorMessage(`{"error":{"message":"API key not valid: ` + authKey + `"}}`)
|
||||
if strings.Contains(got, authKey) {
|
||||
t.Fatalf("Gemini auth key leaked: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "AQ.***REDACTED***") {
|
||||
t.Fatalf("expected redaction marker, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunCheckForModel_OpenAIResponses_DefaultRequest(t *testing.T) {
|
||||
h := &openAICaptureHandler{}
|
||||
endpoint := setupFakeOpenAI(t, h)
|
||||
|
||||
@@ -6,7 +6,9 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
@@ -72,15 +74,13 @@ func (r *fakeCNProbeAccountRepo) GetByID(ctx context.Context, id int64) (*Accoun
|
||||
return r.account, nil
|
||||
}
|
||||
|
||||
// kimi coding 账号的 base_url 指向中转(含 api.kimi.com/coding 路径段即可被识别
|
||||
// 为 kimi coding plan)→ 衍生额度端点落在中转主机上,白名单未列名必须拒绝。
|
||||
func TestCNProviderQuotaService_RejectsURLBlockedByPolicy(t *testing.T) {
|
||||
repo := &fakeCNProbeAccountRepo{account: &Account{
|
||||
ID: 1, Platform: PlatformKimi, Type: AccountTypeAPIKey, Status: StatusActive,
|
||||
Credentials: map[string]any{
|
||||
"account_mode": "coding",
|
||||
"api_key": "sk-test",
|
||||
"base_url": "https://relay.attacker.example/api.kimi.com/coding",
|
||||
"base_url": "https://api.kimi.ai/coding/v1",
|
||||
},
|
||||
}}
|
||||
upstream := &recordingHTTPUpstream{}
|
||||
@@ -92,6 +92,79 @@ func TestCNProviderQuotaService_RejectsURLBlockedByPolicy(t *testing.T) {
|
||||
require.Zero(t, upstream.calls, "probe must not issue any upstream request when URL policy rejects the target")
|
||||
}
|
||||
|
||||
func TestCNProviderQuotaService_RejectsSpoofedKimiURL(t *testing.T) {
|
||||
repo := &fakeCNProbeAccountRepo{account: &Account{
|
||||
ID: 1, Platform: PlatformKimi, Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"account_mode": AccountModeCoding,
|
||||
"api_key": "test-only-placeholder",
|
||||
"base_url": "https://relay.attacker.example/api.kimi.com/coding",
|
||||
},
|
||||
}}
|
||||
upstream := &recordingHTTPUpstream{}
|
||||
svc := NewCNProviderQuotaService(repo, nil, upstream, nil)
|
||||
_, err := svc.QueryUsage(context.Background(), 1)
|
||||
require.ErrorContains(t, err, "CN_QUOTA_NOT_CODING_PLAN")
|
||||
require.Zero(t, upstream.calls)
|
||||
}
|
||||
|
||||
type kimiQuotaTestRepo struct {
|
||||
fakeCNProbeAccountRepo
|
||||
updates map[string]any
|
||||
}
|
||||
|
||||
func (r *kimiQuotaTestRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
|
||||
r.updates = updates
|
||||
return nil
|
||||
}
|
||||
|
||||
type kimiQuotaTestUpstream struct {
|
||||
recordingHTTPUpstream
|
||||
request *http.Request
|
||||
}
|
||||
|
||||
func (u *kimiQuotaTestUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
|
||||
u.calls++
|
||||
u.request = req
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(`{"limits":[{"detail":{"limit":100,"remaining":75}}],"usage":{"limit":1000,"remaining":800}}`)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func TestCNProviderQuotaService_KimiOfficialHosts(t *testing.T) {
|
||||
for _, host := range []string{"api.kimi.com", "api.kimi.ai"} {
|
||||
for _, path := range []string{"/coding", "/coding/", "/coding/v1", "/coding/v1/"} {
|
||||
t.Run(host+path, func(t *testing.T) {
|
||||
repo := &kimiQuotaTestRepo{fakeCNProbeAccountRepo: fakeCNProbeAccountRepo{account: &Account{
|
||||
ID: 1, Platform: PlatformKimi, Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"account_mode": AccountModeCoding,
|
||||
"api_key": "test-only-placeholder",
|
||||
"base_url": "https://" + host + path,
|
||||
},
|
||||
}}}
|
||||
upstream := &kimiQuotaTestUpstream{}
|
||||
svc := NewCNProviderQuotaService(repo, nil, upstream, cnProbeAllowlistConfig(host))
|
||||
result, err := svc.QueryUsage(context.Background(), 1)
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.Success)
|
||||
require.True(t, result.CredentialValid)
|
||||
require.True(t, result.Persisted)
|
||||
require.Equal(t, PlatformKimi, result.Provider)
|
||||
require.Len(t, result.Tiers, 2)
|
||||
require.Equal(t, float64(25), result.Tiers[0].UsedPercent)
|
||||
require.Equal(t, float64(20), result.Tiers[1].UsedPercent)
|
||||
require.Equal(t, float64(25), repo.updates["kimi_5h_used_percent"])
|
||||
require.Equal(t, 1, upstream.calls)
|
||||
require.Equal(t, http.MethodGet, upstream.request.Method)
|
||||
require.Equal(t, "https://"+host+"/coding/v1/usages", upstream.request.URL.String())
|
||||
require.Equal(t, "Bearer test-only-placeholder", upstream.request.Header.Get("Authorization"))
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// deepseek payg 账号自定义 base_url → 余额端点落在中转主机上,必须先过策略。
|
||||
func TestCNProviderBalanceService_RejectsURLBlockedByPolicy(t *testing.T) {
|
||||
repo := &fakeCNProbeAccountRepo{account: &Account{
|
||||
|
||||
@@ -720,7 +720,18 @@ func (s *ContentModerationService) UpdateConfig(ctx context.Context, input Updat
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal content moderation config: %w", err)
|
||||
}
|
||||
if err := s.settingRepo.Set(ctx, SettingKeyContentModerationConfig, string(raw)); err != nil {
|
||||
var saveErr error
|
||||
if writer, ok := s.settingRepo.(interface {
|
||||
SetContentModerationConfig(context.Context, string, []int64) error
|
||||
}); ok {
|
||||
saveErr = writer.SetContentModerationConfig(ctx, string(raw), cfg.GroupIDs)
|
||||
if errors.Is(saveErr, ErrGroupNotFound) {
|
||||
return nil, infraerrors.BadRequest("INVALID_CONTENT_MODERATION_GROUP", "审计分组不存在")
|
||||
}
|
||||
} else {
|
||||
saveErr = s.settingRepo.Set(ctx, SettingKeyContentModerationConfig, string(raw))
|
||||
}
|
||||
if err := saveErr; err != nil {
|
||||
return nil, fmt.Errorf("save content moderation config: %w", err)
|
||||
}
|
||||
s.replaceRuntimeConfig(cfg, raw)
|
||||
|
||||
@@ -184,7 +184,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
|
||||
logger.LegacyPrintf("service.gateway", "[Anthropic Passthrough] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d RequestID=%s Body=%s",
|
||||
account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(respBody), 1000))
|
||||
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account)
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account, input.RequestModel)
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
ProxyID: opsUpstreamProxyID(account),
|
||||
ProxyName: opsUpstreamProxyName(account),
|
||||
@@ -209,7 +209,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
return s.handleRetryExhaustedError(ctx, resp, c, account)
|
||||
return s.handleRetryExhaustedError(ctx, resp, c, account, input.RequestModel)
|
||||
}
|
||||
|
||||
if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
|
||||
@@ -130,7 +130,9 @@ func (s *GatewayService) forwardBedrock(
|
||||
|
||||
// 错误/failover 处理
|
||||
if resp.StatusCode >= 400 {
|
||||
return s.handleBedrockUpstreamErrors(ctx, resp, c, account)
|
||||
// 账号状态按模型隔离时用调度器的模型键(账号映射后的模型名),不是发往上游的 Bedrock 模型 ID;
|
||||
// 两者不一致会让模型级冷却写到调度器不会读取的键上。
|
||||
return s.handleBedrockUpstreamErrors(ctx, resp, c, account, account.GetMappedModel(reqModel))
|
||||
}
|
||||
|
||||
// Bedrock 分支绕过通用 Forward 成功路径,这里保持上游接受回调语义一致。
|
||||
@@ -263,12 +265,14 @@ func (s *GatewayService) executeBedrockUpstream(
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// handleBedrockUpstreamErrors 处理 Bedrock 上游 4xx/5xx 错误(failover + 错误响应)
|
||||
// handleBedrockUpstreamErrors 处理 Bedrock 上游 4xx/5xx 错误(failover + 错误响应);
|
||||
// modelKey 是调度器按模型隔离账号状态时使用的模型键(account.GetMappedModel 的结果)。
|
||||
func (s *GatewayService) handleBedrockUpstreamErrors(
|
||||
ctx context.Context,
|
||||
resp *http.Response,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
modelKey string,
|
||||
) (*ForwardResult, error) {
|
||||
// retry exhausted + failover
|
||||
if s.shouldRetryUpstreamError(account, resp.StatusCode) {
|
||||
@@ -280,7 +284,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
|
||||
logger.LegacyPrintf("service.gateway", "[Bedrock] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d Body=%s",
|
||||
account.ID, account.Name, resp.StatusCode, truncateString(string(respBody), 1000))
|
||||
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account)
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account, modelKey)
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
ProxyID: opsUpstreamProxyID(account),
|
||||
ProxyName: opsUpstreamProxyName(account),
|
||||
@@ -297,7 +301,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
return s.handleRetryExhaustedError(ctx, resp, c, account)
|
||||
return s.handleRetryExhaustedError(ctx, resp, c, account, modelKey)
|
||||
}
|
||||
|
||||
// non-retryable failover
|
||||
@@ -306,7 +310,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
|
||||
_ = resp.Body.Close()
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
|
||||
s.handleFailoverSideEffects(ctx, resp, account)
|
||||
s.handleFailoverSideEffects(ctx, resp, account, modelKey)
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
ProxyID: opsUpstreamProxyID(account),
|
||||
ProxyName: opsUpstreamProxyName(account),
|
||||
@@ -325,7 +329,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
|
||||
}
|
||||
|
||||
// other errors
|
||||
return s.handleErrorResponse(ctx, resp, c, account)
|
||||
return s.handleErrorResponse(ctx, resp, c, account, modelKey)
|
||||
}
|
||||
|
||||
// buildUpstreamRequestBedrock 构建 Bedrock 上游请求
|
||||
|
||||
@@ -263,7 +263,11 @@ func normalizeClaudeOAuthRequestBody(body []byte, modelID string, opts claudeOAu
|
||||
// temperature:真实 Claude Code CLI 总是发送 temperature(默认 1,客户端可覆盖)。
|
||||
// 之前的实现直接 delete 会导致 payload 缺字段,与真实 CLI 字节级不一致。
|
||||
// 策略:客户端传了什么就透传;没传则补默认 1。
|
||||
if !gjson.GetBytes(out, "temperature").Exists() && !claude.IsOpus55(modelID) {
|
||||
// 例外:采样锁定的模型(sonnet-5-5/haiku-5-5)上 temperature+top_p 共存即 400,
|
||||
// 客户端只传 top_p 时不得再补 temperature,否则网关自己制造 400(issue #7919)。
|
||||
samplingLocked := claude.IsSonnet55(modelID) || claude.IsHaiku55(modelID)
|
||||
if !gjson.GetBytes(out, "temperature").Exists() && !claude.IsOpus55(modelID) &&
|
||||
(!samplingLocked || !gjson.GetBytes(out, "top_p").Exists()) {
|
||||
if next, ok := setJSONValueBytes(out, "temperature", 1); ok {
|
||||
out = next
|
||||
modified = true
|
||||
|
||||
@@ -371,6 +371,25 @@ func TestNormalizeClaudeOAuthRequestBody_HaikuShortModelStillNormalizesToDatedID
|
||||
require.Equal(t, "claude-haiku-4-5-20251001", gjson.GetBytes(out, "model").String())
|
||||
}
|
||||
|
||||
// issue #7919:采样锁定的模型上 temperature+top_p 共存即 400,
|
||||
// 只传 top_p 时不得再补 temperature,否则网关自己制造 400。
|
||||
func TestNormalizeClaudeOAuthRequestBody_SkipsTemperatureTopUpWithTopPOnLockedModels(t *testing.T) {
|
||||
for _, model := range []string{"claude-haiku-5-5", "claude-sonnet-5-5"} {
|
||||
body := []byte(`{"model":"` + model + `","top_p":0.99,"messages":[]}`)
|
||||
out, _ := normalizeClaudeOAuthRequestBody(body, model, claudeOAuthNormalizeOptions{})
|
||||
require.False(t, gjson.GetBytes(out, "temperature").Exists(),
|
||||
"model %s with top_p must not gain temperature", model)
|
||||
require.Equal(t, 0.99, gjson.GetBytes(out, "top_p").Float())
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeClaudeOAuthRequestBody_StillTopsUpTemperatureWithoutTopP(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-haiku-5-5","messages":[]}`)
|
||||
out, _ := normalizeClaudeOAuthRequestBody(body, "claude-haiku-5-5", claudeOAuthNormalizeOptions{})
|
||||
require.Equal(t, int64(1), gjson.GetBytes(out, "temperature").Int(),
|
||||
"无 top_p 时 temperature:1 补齐保持不变")
|
||||
}
|
||||
|
||||
func TestApplyClaudeCodeOAuthMimicryToBody_HaikuRewritesSystem(t *testing.T) {
|
||||
account := &Account{ID: 405, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
|
||||
body := []byte(`{"model":"claude-haiku-4-5","system":"Pi project instructions","messages":[{"role":"user","content":"hello"}]}`)
|
||||
|
||||
@@ -42,13 +42,15 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
|
||||
|
||||
if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() {
|
||||
passthroughBody := parsed.Body.Bytes()
|
||||
passthroughModel := parsed.Model
|
||||
if reqModel := parsed.Model; reqModel != "" {
|
||||
if mappedModel := account.GetMappedModel(reqModel); mappedModel != reqModel {
|
||||
passthroughBody = s.replaceModelInBody(passthroughBody, mappedModel)
|
||||
passthroughModel = mappedModel
|
||||
logger.LegacyPrintf("service.gateway", "CountTokens passthrough model mapping: %s -> %s (account: %s)", reqModel, mappedModel, account.Name)
|
||||
}
|
||||
}
|
||||
return s.forwardCountTokensAnthropicAPIKeyPassthrough(ctx, c, account, passthroughBody)
|
||||
return s.forwardCountTokensAnthropicAPIKeyPassthrough(ctx, c, account, passthroughBody, passthroughModel)
|
||||
}
|
||||
|
||||
// Bedrock 不支持 count_tokens 端点
|
||||
@@ -220,8 +222,8 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
|
||||
|
||||
// 处理错误响应
|
||||
if resp.StatusCode >= 400 {
|
||||
// 标记账号状态(429/529等)
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
// 标记账号状态(429/529等);传入实际发往上游的模型,临时不可调度规则才能按模型隔离
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, reqModel)
|
||||
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
@@ -267,7 +269,9 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx context.Context, c *gin.Context, account *Account, body []byte) error {
|
||||
// forwardCountTokensAnthropicAPIKeyPassthrough 透传 count_tokens 请求;upstreamModel 是 body 中实际发往上游的模型,
|
||||
// 仅用于上游错误处理按模型隔离。
|
||||
func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx context.Context, c *gin.Context, account *Account, body []byte, upstreamModel string) error {
|
||||
token, tokenType, err := s.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token")
|
||||
@@ -322,7 +326,7 @@ func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx contex
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
if s.rateLimitService != nil {
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel)
|
||||
}
|
||||
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
|
||||
@@ -689,7 +689,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d RequestID=%s Body=%s",
|
||||
account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(respBody), 1000))
|
||||
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account)
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account, reqModel)
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
ProxyID: opsUpstreamProxyID(account),
|
||||
ProxyName: opsUpstreamProxyName(account),
|
||||
@@ -713,7 +713,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
return s.handleRetryExhaustedError(ctx, resp, c, account)
|
||||
return s.handleRetryExhaustedError(ctx, resp, c, account, reqModel)
|
||||
}
|
||||
|
||||
// 处理可切换账号的错误
|
||||
|
||||
@@ -131,7 +131,7 @@ func TestCalculateRecordUsageCost_MediaReasoningPricing(t *testing.T) {
|
||||
}
|
||||
var err error
|
||||
cost, err = svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey,
|
||||
[]string{model}, 0.5, 0.5, 0.5, 0.5, UsageTokens{}, "", nil, time.Time{})
|
||||
[]string{model}, 0.5, 0.5, 0.5, 0.5, UsageTokens{}, "", nil, false, time.Time{})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.NotNil(t, cost)
|
||||
|
||||
@@ -570,7 +570,7 @@ func StripEmptyTextBlocks(body []byte) []byte {
|
||||
// isClaude55SignedThinkingModel identifies models whose default thinking mode
|
||||
// requires signed history to survive protocol conversion and request filtering.
|
||||
func isClaude55SignedThinkingModel(model string) bool {
|
||||
return claude.IsOpus55(model) || claude.IsSonnet55(model)
|
||||
return claude.IsOpus55(model) || claude.IsSonnet55(model) || claude.IsHaiku55(model)
|
||||
}
|
||||
|
||||
// validateClaude55Request rejects settings that the upstream cannot honor.
|
||||
@@ -580,14 +580,31 @@ func validateClaude55Request(body []byte, model string) error {
|
||||
return nil
|
||||
}
|
||||
isSonnet55 := claude.IsSonnet55(model)
|
||||
switch gjson.GetBytes(body, "thinking.type").String() {
|
||||
isHaiku55 := claude.IsHaiku55(model)
|
||||
thinkingType := gjson.GetBytes(body, "thinking.type").String()
|
||||
switch thinkingType {
|
||||
case "disabled", "enabled":
|
||||
if isHaiku55 {
|
||||
// haiku-5-5:enabled 一律 400(用 adaptive);disabled 仅 high 及以下
|
||||
// 可用,xhigh/max 上 400。见官方 migration guide。
|
||||
if thinkingType == "enabled" {
|
||||
return fmt.Errorf("claude-haiku-5-5 requires adaptive thinking; omit thinking or use thinking.type=adaptive and output_config.effort")
|
||||
}
|
||||
switch gjson.GetBytes(body, "output_config.effort").String() {
|
||||
case "xhigh", "max":
|
||||
return fmt.Errorf("claude-haiku-5-5 thinking.type=disabled supports only low, medium or high effort")
|
||||
}
|
||||
break
|
||||
}
|
||||
if isSonnet55 {
|
||||
return fmt.Errorf("claude-sonnet-5-5 requires adaptive thinking or thinking.type=between_tools; omit thinking or use one of those modes")
|
||||
}
|
||||
return fmt.Errorf("claude-opus-5-5 requires adaptive thinking; omit thinking or use thinking.type=adaptive and output_config.effort")
|
||||
case "between_tools":
|
||||
if !isSonnet55 {
|
||||
if isHaiku55 {
|
||||
return fmt.Errorf("claude-haiku-5-5 requires adaptive thinking; thinking.type=between_tools is unsupported (use thinking.type=disabled at high effort or below)")
|
||||
}
|
||||
return fmt.Errorf("claude-opus-5-5 requires adaptive thinking; thinking.type=between_tools is unsupported")
|
||||
}
|
||||
effort := gjson.GetBytes(body, "output_config.effort").String()
|
||||
@@ -604,22 +621,34 @@ func validateClaude55Request(body []byte, model string) error {
|
||||
if isSonnet55 {
|
||||
modelName = "claude-sonnet-5-5"
|
||||
}
|
||||
if gjson.GetBytes(body, "tool_choice").String() == "required" {
|
||||
return fmt.Errorf("%s does not support forced tool_choice; use auto or none", modelName)
|
||||
if isHaiku55 {
|
||||
modelName = "claude-haiku-5-5"
|
||||
}
|
||||
switch gjson.GetBytes(body, "tool_choice.type").String() {
|
||||
case "any", "tool", "function", "custom", "namespace":
|
||||
return fmt.Errorf("%s does not support forced tool_choice; use auto or none", modelName)
|
||||
}
|
||||
if isSonnet55 {
|
||||
if temperature := gjson.GetBytes(body, "temperature"); temperature.Exists() && (temperature.Type != gjson.Number || temperature.Float() != 1) {
|
||||
return fmt.Errorf("claude-sonnet-5-5 does not support non-default temperature")
|
||||
// haiku-5-5 接受强制 tool_choice(响应以 tool call 开头、无 thinking 块),
|
||||
// 与 opus/sonnet-5-5 不同,不拦截。见官方 migration guide。
|
||||
if !isHaiku55 {
|
||||
if gjson.GetBytes(body, "tool_choice").String() == "required" {
|
||||
return fmt.Errorf("%s does not support forced tool_choice; use auto or none", modelName)
|
||||
}
|
||||
if topP := gjson.GetBytes(body, "top_p"); topP.Exists() && (topP.Type != gjson.Number || topP.Float() < 0.99 || topP.Float() > 1) {
|
||||
return fmt.Errorf("claude-sonnet-5-5 does not support non-default top_p")
|
||||
switch gjson.GetBytes(body, "tool_choice.type").String() {
|
||||
case "any", "tool", "function", "custom", "namespace":
|
||||
return fmt.Errorf("%s does not support forced tool_choice; use auto or none", modelName)
|
||||
}
|
||||
}
|
||||
if isSonnet55 || isHaiku55 {
|
||||
if temperature := gjson.GetBytes(body, "temperature"); temperature.Exists() && (temperature.Type != gjson.Number || temperature.Float() != 1) {
|
||||
return fmt.Errorf("%s does not support non-default temperature", modelName)
|
||||
}
|
||||
// Haiku 5.5 只接受显式 top_p=0.99;保留 Sonnet 5.5 既有的区间策略。
|
||||
if topP := gjson.GetBytes(body, "top_p"); topP.Exists() && (topP.Type != gjson.Number ||
|
||||
(isHaiku55 && topP.Float() != 0.99) || (isSonnet55 && (topP.Float() < 0.99 || topP.Float() > 1))) {
|
||||
return fmt.Errorf("%s does not support non-default top_p", modelName)
|
||||
}
|
||||
if gjson.GetBytes(body, "top_k").Exists() {
|
||||
return fmt.Errorf("claude-sonnet-5-5 does not support top_k")
|
||||
return fmt.Errorf("%s does not support top_k", modelName)
|
||||
}
|
||||
if isHaiku55 && gjson.GetBytes(body, "temperature").Exists() && gjson.GetBytes(body, "top_p").Exists() {
|
||||
return fmt.Errorf("%s does not allow temperature and top_p together; omit both", modelName)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// issue #7919:claude-haiku-5-5 请求校验(官方 migration guide 口径)。
|
||||
func TestValidateClaude55Request_Haiku55(t *testing.T) {
|
||||
const model = "claude-haiku-5-5"
|
||||
const prefix = `{"model":"claude-haiku-5-5","max_tokens":16000,"messages":[{"role":"user","content":"hi"}]`
|
||||
cases := []struct {
|
||||
name string
|
||||
// extra 为追加进最小合法请求的 JSON 片段(含前导逗号)
|
||||
extra string
|
||||
wantErr string // 为空表示期望通过
|
||||
}{
|
||||
{name: "adaptive thinking passes", extra: `,"thinking":{"type":"adaptive"},"output_config":{"effort":"medium"}`},
|
||||
{name: "omitted thinking passes", extra: ``},
|
||||
{name: "enabled thinking rejected", extra: `,"thinking":{"type":"enabled","budget_tokens":8000}`, wantErr: "requires adaptive thinking"},
|
||||
{name: "拒绝 between_tools 并返回 Haiku 错误说明", extra: `,"thinking":{"type":"between_tools"}`, wantErr: "claude-haiku-5-5 requires adaptive thinking; thinking.type=between_tools is unsupported"},
|
||||
{name: "disabled thinking allowed at medium effort", extra: `,"thinking":{"type":"disabled"},"output_config":{"effort":"medium"}`},
|
||||
{name: "disabled thinking allowed when effort omitted", extra: `,"thinking":{"type":"disabled"}`},
|
||||
{name: "disabled thinking rejected at xhigh", extra: `,"thinking":{"type":"disabled"},"output_config":{"effort":"xhigh"}`, wantErr: "supports only low, medium or high effort"},
|
||||
{name: "disabled thinking rejected at max", extra: `,"thinking":{"type":"disabled"},"output_config":{"effort":"max"}`, wantErr: "supports only low, medium or high effort"},
|
||||
{name: "temperature 1 passes", extra: `,"temperature":1`},
|
||||
{name: "non-default temperature rejected", extra: `,"temperature":0.5`, wantErr: "non-default temperature"},
|
||||
{name: "top_p 0.99 passes", extra: `,"top_p":0.99`},
|
||||
{name: "拒绝 top_p=0.995", extra: `,"top_p":0.995`, wantErr: "non-default top_p"},
|
||||
{name: "拒绝 top_p=1", extra: `,"top_p":1`, wantErr: "non-default top_p"},
|
||||
{name: "拒绝字符串 top_p", extra: `,"top_p":"0.99"`, wantErr: "non-default top_p"},
|
||||
{name: "top_p 0.5 rejected", extra: `,"top_p":0.5`, wantErr: "non-default top_p"},
|
||||
{name: "top_k rejected", extra: `,"top_k":5`, wantErr: "does not support top_k"},
|
||||
{name: "temperature and top_p together rejected", extra: `,"temperature":1,"top_p":0.99`, wantErr: "does not allow temperature and top_p together"},
|
||||
{name: "forced tool_choice any allowed", extra: `,"tool_choice":{"type":"any"}`},
|
||||
{name: "forced tool_choice required allowed", extra: `,"tool_choice":"required"`},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
body := []byte(prefix + tc.extra + "}")
|
||||
err := validateClaude55Request(body, model)
|
||||
if tc.wantErr == "" {
|
||||
require.NoError(t, err)
|
||||
return
|
||||
}
|
||||
require.Error(t, err)
|
||||
require.True(t, strings.Contains(err.Error(), tc.wantErr),
|
||||
"error %q should contain %q", err.Error(), tc.wantErr)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -649,6 +649,9 @@ type ForwardResult struct {
|
||||
ImageSizeBreakdown map[string]int
|
||||
SearchCount int
|
||||
AudioUsage *AudioUsage
|
||||
// CacheStorageTokenHours 显式上下文缓存的存储量(缓存 token 数 × 存储小时数),按计费模型的
|
||||
// 存储单价折算为费用后叠加在 token 费用之上。
|
||||
CacheStorageTokenHours float64
|
||||
}
|
||||
|
||||
// GatewayFailureStage identifies which request stage failed. The zero value is
|
||||
@@ -696,6 +699,7 @@ type UpstreamFailoverError struct {
|
||||
SameAccountRetryMax int // 可选的错误级同账号重试上限,低于 handler 默认预算时优先采用
|
||||
RequestScopedTransient bool // 故障因素与账号无关(如上游按客户端身份/模型容量降载):可同账号重试,但不得据此对账号做临时封禁
|
||||
SafeToFailoverAfterWrite bool // 仅写出 SSE 注释等非语义字节时,仍可在同一客户端流中切换账号
|
||||
BoundUpstreamPassthrough bool // 绑定显式缓存的请求在上游网关账号上收到的 429 / 5xx:反映持有缓存的上游账号状态,原样回写不走错误映射
|
||||
Stage GatewayFailureStage
|
||||
Scope GatewayFailureScope
|
||||
Reason GatewayFailureReason
|
||||
|
||||
@@ -419,11 +419,13 @@ func defaultAPIKeyBetaHeader(body []byte) string {
|
||||
|
||||
// The legacy streaming beta is incompatible with the stable computer and
|
||||
// browser toolsets. The per-tool eager_input_streaming field remains available.
|
||||
// Applies to Sonnet 5.5 and Haiku 5.5, which share the computer_toolset_20260801
|
||||
// / browser_toolset_20260801 toolsets (Haiku 5.5 migration guide).
|
||||
func filterSonnet55ToolsetBeta(header string, body []byte, modelID string) string {
|
||||
if at := strings.IndexByte(modelID, '@'); at >= 0 {
|
||||
modelID = modelID[:at] // Vertex's dated model ID
|
||||
}
|
||||
if !claude.IsSonnet55(modelID) {
|
||||
if !claude.IsSonnet55(modelID) && !claude.IsHaiku55(modelID) {
|
||||
return header
|
||||
}
|
||||
for _, tool := range gjson.GetBytes(body, "tools").Array() {
|
||||
|
||||
@@ -527,13 +527,19 @@ func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Res
|
||||
return nil, fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
|
||||
}
|
||||
|
||||
func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account) {
|
||||
// handleRetryExhaustedSideEffects 在重试耗尽后标记账号状态;requestedModel 是实际发往上游的模型,
|
||||
// 临时不可调度规则据此按模型隔离。
|
||||
func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) {
|
||||
body, _ := s.readUpstreamErrorBody(resp)
|
||||
statusCode := resp.StatusCode
|
||||
|
||||
// OAuth/Setup Token 账号的 403:标记账号异常
|
||||
if account.IsOAuth() && statusCode == 403 {
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, resp.Header, body)
|
||||
if len(requestedModel) > 0 {
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, resp.Header, body, requestedModel[0])
|
||||
} else {
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, resp.Header, body)
|
||||
}
|
||||
logger.LegacyPrintf("service.gateway", "Account %d: marked as error after %d retries for status %d", account.ID, maxRetryAttempts, statusCode)
|
||||
} else {
|
||||
// API Key 未配置错误码:不标记账号状态
|
||||
@@ -553,14 +559,14 @@ func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *ht
|
||||
// handleRetryExhaustedError 处理重试耗尽后的错误
|
||||
// OAuth 403:标记账号异常
|
||||
// API Key 未配置错误码:仅返回错误,不标记账号
|
||||
func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account) (*ForwardResult, error) {
|
||||
func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, requestedModel ...string) (*ForwardResult, error) {
|
||||
MarkResponseCommitted(c)
|
||||
// Capture upstream error body before side-effects consume the stream.
|
||||
respBody, _ := s.readUpstreamErrorBody(resp)
|
||||
_ = resp.Body.Close()
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account)
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account, requestedModel...)
|
||||
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
|
||||
@@ -241,7 +241,8 @@ func isForcedUsageBillingRequestID(requestID string) bool {
|
||||
return strings.HasPrefix(id, "web_search:") ||
|
||||
strings.HasPrefix(id, "grok-video:") ||
|
||||
strings.HasPrefix(id, "grok_audio:") ||
|
||||
strings.HasPrefix(id, "grok_realtime:")
|
||||
strings.HasPrefix(id, "grok_realtime:") ||
|
||||
strings.HasPrefix(id, GeminiCachedContentUsageRequestIDPrefix)
|
||||
}
|
||||
|
||||
// StableGrokAudioBillingRequestID is the durable usage_logs / dedup key for one
|
||||
@@ -787,42 +788,9 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
cacheTTLOverridden = (result.Usage.CacheCreation5mTokens + result.Usage.CacheCreation1hTokens) > 0
|
||||
}
|
||||
|
||||
// 获取费率倍数(优先级:用户专属 > 分组默认 > 系统默认)
|
||||
multiplier := 1.0
|
||||
if s.cfg != nil {
|
||||
multiplier = s.cfg.Default.RateMultiplier
|
||||
}
|
||||
if apiKey.GroupID != nil && apiKey.Group != nil {
|
||||
groupDefault := apiKey.Group.RateMultiplier
|
||||
multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault)
|
||||
}
|
||||
// token 倍率叠加高峰因子(token 计费含图片 token,图片按次倍率不受影响)。高峰因子按请求时刻现算,
|
||||
// 不并入上面的 getUserGroupRateMultiplier,以免污染 user:group 倍率缓存。
|
||||
pricingAt := input.PricingAt
|
||||
if pricingAt.IsZero() {
|
||||
pricingAt = timezone.Now()
|
||||
}
|
||||
multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, pricingAt)
|
||||
|
||||
// 确定计费模型
|
||||
concreteBillingModel := forwardResultBillingModel(result.Model, result.UpstreamModel)
|
||||
billingModel := concreteBillingModel
|
||||
if input.BillingModelSource == BillingModelSourceChannelMapped && input.ChannelMappedModel != "" {
|
||||
billingModel = input.ChannelMappedModel
|
||||
}
|
||||
if input.BillingModelSource == BillingModelSourceRequested && input.OriginalModel != "" {
|
||||
billingModel = input.OriginalModel
|
||||
}
|
||||
// composite 分组的公开别名(如 all/claude)会经 OriginalModel/ChannelMappedModel
|
||||
// 进入上面的来源覆盖:任意别名查无价会静默落 $0,含家族词的别名则被价格表的
|
||||
// 家族模糊匹配错计(如 Opus 流量按 Sonnet 兜底价)。除非管理员为别名显式配置了
|
||||
// 渠道定价(OpenRouter 式自定价),composite 请求一律按实际转发的具体模型计费。
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == PlatformComposite {
|
||||
billingModel = s.compositeBillableModel(ctx, apiKey, billingModel, concreteBillingModel)
|
||||
}
|
||||
// 通用兜底(与 OpenAI 路径的 usageBillingModelCandidates 语义对齐):
|
||||
// 选定模型查不到任何价格时回退到实际转发的具体模型。已定价流量不受影响。
|
||||
billingModel = s.billableModelWithFallback(ctx, apiKey, billingModel, result.UpstreamModel, result.Model)
|
||||
multiplier, imageMultiplier, pricingAt := s.recordUsageMultipliers(ctx, apiKey, user, input.PricingAt)
|
||||
billingModel := s.recordUsageBillingModel(ctx, apiKey, result, input.ChannelUsageFields)
|
||||
cacheStorageCost := s.cacheStorageCost(result, billingModel)
|
||||
|
||||
// 确定 RequestedModel(渠道映射前的原始模型)
|
||||
requestedModel := result.Model
|
||||
@@ -840,7 +808,7 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
input.BillingModelSource,
|
||||
result.UpstreamResponseModel,
|
||||
result.UpstreamResponseModelConflict,
|
||||
result.ImageCount > 0 || result.AudioUsage != nil || result.SearchCount > 0,
|
||||
result.ImageCount > 0 || result.AudioUsage != nil || result.SearchCount > 0 || result.CacheStorageTokenHours > 0,
|
||||
); responseModel != "" && !strings.EqualFold(responseModel, strings.TrimSpace(billingModel)) {
|
||||
if identified, responseChannelPriced := s.hasIdentifiedResponseModelPricing(ctx, responseModel, apiKey); identified {
|
||||
responseCost := s.calculateRecordUsageCost(ctx, result, apiKey, responseModel, multiplier, imageMultiplier, pricingAt)
|
||||
@@ -866,8 +834,9 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
usageLog := s.buildRecordUsageLog(ctx, input, result, apiKey, user, account, subscription,
|
||||
requestedModel, multiplier, imageMultiplier, accountRateMultiplier, billingType, cacheTTLOverridden, cost)
|
||||
|
||||
// 计算账号统计定价费用(使用最终上游模型匹配自定义规则)
|
||||
if apiKey.GroupID != nil {
|
||||
// 计算账号统计定价费用(使用最终上游模型匹配自定义规则)。纯存储费不是一次模型请求,
|
||||
// 账号统计成本回落 total_cost(官方存储价)。
|
||||
if apiKey.GroupID != nil && !isCacheStorageOnlyResult(result) {
|
||||
applyAccountStatsCost(ctx, usageLog, s.channelService, s.billingService,
|
||||
account.ID, *apiKey.GroupID, result.UpstreamModel, result.Model,
|
||||
// Anthropic's input_tokens excludes cache_read and cache_creation (billed separately);
|
||||
@@ -879,9 +848,10 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
CacheReadTokens: result.Usage.CacheReadInputTokens,
|
||||
ImageOutputTokens: result.Usage.ImageOutputTokens,
|
||||
},
|
||||
cost.TotalCost, pricingAt,
|
||||
cost.TotalCost-cacheStorageCost, pricingAt,
|
||||
accountStatsLongContextPricingEnabled(nil),
|
||||
)
|
||||
addCacheStorageToAccountStatsCost(usageLog, cacheStorageCost)
|
||||
}
|
||||
|
||||
simpleModeKeyRateLimitOnly := simpleModeKeyRateLimitBillingEnabled(s.cfg, apiKey)
|
||||
@@ -927,6 +897,112 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage
|
||||
return nil
|
||||
}
|
||||
|
||||
// recordUsageMultipliers 返回 token / 图片售价倍率(用户专属 > 分组默认 > 系统默认,token 倍率叠加高峰因子)
|
||||
// 与计价时刻;pricingAt 为零值时取当前时刻。
|
||||
func (s *GatewayService) recordUsageMultipliers(ctx context.Context, apiKey *APIKey, user *User, pricingAt time.Time) (float64, float64, time.Time) {
|
||||
multiplier := 1.0
|
||||
if s.cfg != nil {
|
||||
multiplier = s.cfg.Default.RateMultiplier
|
||||
}
|
||||
if apiKey.GroupID != nil && apiKey.Group != nil {
|
||||
multiplier = apiKey.Group.RateMultiplier
|
||||
if user != nil {
|
||||
multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, multiplier)
|
||||
}
|
||||
}
|
||||
// 高峰因子按请求时刻现算,不并入 getUserGroupRateMultiplier,以免污染 user:group 倍率缓存。
|
||||
if pricingAt.IsZero() {
|
||||
pricingAt = timezone.Now()
|
||||
}
|
||||
multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, pricingAt)
|
||||
return multiplier, imageMultiplier, pricingAt
|
||||
}
|
||||
|
||||
// recordUsageBillingModel 按计费模型来源、composite 规则与无价兜底确定计费模型。
|
||||
func (s *GatewayService) recordUsageBillingModel(ctx context.Context, apiKey *APIKey, result *ForwardResult, fields ChannelUsageFields) string {
|
||||
concreteBillingModel := forwardResultBillingModel(result.Model, result.UpstreamModel)
|
||||
billingModel := concreteBillingModel
|
||||
if fields.BillingModelSource == BillingModelSourceChannelMapped && fields.ChannelMappedModel != "" {
|
||||
billingModel = fields.ChannelMappedModel
|
||||
}
|
||||
if fields.BillingModelSource == BillingModelSourceRequested && fields.OriginalModel != "" {
|
||||
billingModel = fields.OriginalModel
|
||||
}
|
||||
// composite 分组的公开别名(如 all/claude)会经 OriginalModel/ChannelMappedModel
|
||||
// 进入上面的来源覆盖:任意别名查无价会静默落 $0,含家族词的别名则被价格表的
|
||||
// 家族模糊匹配错计(如 Opus 流量按 Sonnet 兜底价)。除非管理员为别名显式配置了
|
||||
// 渠道定价(OpenRouter 式自定价),composite 请求一律按实际转发的具体模型计费。
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == PlatformComposite {
|
||||
billingModel = s.compositeBillableModel(ctx, apiKey, billingModel, concreteBillingModel)
|
||||
}
|
||||
// 通用兜底(与 OpenAI 路径的 usageBillingModelCandidates 语义对齐):
|
||||
// 选定模型查不到任何价格时回退到实际转发的具体模型。已定价流量不受影响。
|
||||
return s.billableModelWithFallback(ctx, apiKey, billingModel, result.UpstreamModel, result.Model)
|
||||
}
|
||||
|
||||
// cacheStoragePrice 返回计费模型的显式缓存存储单价;计费模型未配置时依次回退到实际转发模型与请求模型。
|
||||
func (s *GatewayService) cacheStoragePrice(result *ForwardResult, billingModel string) float64 {
|
||||
if s.billingService == nil {
|
||||
return 0
|
||||
}
|
||||
for _, model := range []string{billingModel, result.UpstreamModel, result.Model} {
|
||||
if strings.TrimSpace(model) == "" {
|
||||
continue
|
||||
}
|
||||
if price := s.billingService.CacheStoragePricePerTokenHour(model); price > 0 {
|
||||
return price
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// cacheStorageCost 返回 result 的显式缓存存储费(USD,倍率前)。
|
||||
func (s *GatewayService) cacheStorageCost(result *ForwardResult, billingModel string) float64 {
|
||||
if result.CacheStorageTokenHours <= 0 {
|
||||
return 0
|
||||
}
|
||||
return result.CacheStorageTokenHours * s.cacheStoragePrice(result, billingModel)
|
||||
}
|
||||
|
||||
// isCacheStorageOnlyResult 判断 result 是否只有显式缓存存储费(延长有效期):不含 token 用量,不算一次模型请求。
|
||||
func isCacheStorageOnlyResult(result *ForwardResult) bool {
|
||||
u := result.Usage
|
||||
return result.CacheStorageTokenHours > 0 && u.InputTokens == 0 && u.OutputTokens == 0 &&
|
||||
u.CacheCreationInputTokens == 0 && u.CacheReadInputTokens == 0 && u.ImageOutputTokens == 0 &&
|
||||
result.ImageCount == 0 && result.AudioUsage == nil && result.SearchCount == 0
|
||||
}
|
||||
|
||||
// CacheStoragePriced 判断 RecordUsage 对 input.Result 采用的计费模型(含回退候选)是否配置了显式缓存存储单价。
|
||||
func (s *GatewayService) CacheStoragePriced(ctx context.Context, input *RecordUsageInput) bool {
|
||||
billingModel := s.recordUsageBillingModel(ctx, input.APIKey, input.Result, input.ChannelUsageFields)
|
||||
return s.cacheStoragePrice(input.Result, billingModel) > 0
|
||||
}
|
||||
|
||||
// QuoteUsageCost 按与 RecordUsage 相同的倍率、计费模型与定价计算 input.Result 的应付金额(倍率后 USD),不记录也不扣费。
|
||||
func (s *GatewayService) QuoteUsageCost(ctx context.Context, input *RecordUsageInput) float64 {
|
||||
user := input.User
|
||||
if user == nil {
|
||||
user = input.APIKey.User
|
||||
}
|
||||
multiplier, imageMultiplier, pricingAt := s.recordUsageMultipliers(ctx, input.APIKey, user, input.PricingAt)
|
||||
billingModel := s.recordUsageBillingModel(ctx, input.APIKey, input.Result, input.ChannelUsageFields)
|
||||
cost := s.calculateRecordUsageCost(ctx, input.Result, input.APIKey, billingModel, multiplier, imageMultiplier, pricingAt)
|
||||
if cost == nil {
|
||||
return 0
|
||||
}
|
||||
return cost.ActualCost
|
||||
}
|
||||
|
||||
// addCacheStorageToAccountStatsCost 把显式缓存存储费按官方价叠加到账号统计成本。
|
||||
// 账号统计定价只覆盖 token 部分;统计成本为空时成本统计回落 total_cost,其中已含存储费。
|
||||
func addCacheStorageToAccountStatsCost(usageLog *UsageLog, storageCost float64) {
|
||||
if usageLog == nil || usageLog.AccountStatsCost == nil || storageCost <= 0 {
|
||||
return
|
||||
}
|
||||
cost := *usageLog.AccountStatsCost + storageCost
|
||||
usageLog.AccountStatsCost = &cost
|
||||
}
|
||||
|
||||
// calculateRecordUsageCost 根据请求类型计算费用。
|
||||
func (s *GatewayService) calculateRecordUsageCost(
|
||||
ctx context.Context,
|
||||
@@ -964,8 +1040,13 @@ func (s *GatewayService) calculateRecordUsageCost(
|
||||
return s.billingService.CalculateAudioCost(result.AudioUsage.Mode, result.AudioUsage.DurationOrUnits, cfg, multiplier)
|
||||
}
|
||||
|
||||
// Token 计费;SearchCount 为叠加 surcharge(不替代 token)。
|
||||
tokenCost := s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, pricingAt)
|
||||
// Token 计费;SearchCount 为叠加 surcharge(不替代 token)。纯存储费不是一次模型请求,不计 token / 按次费用。
|
||||
var tokenCost *CostBreakdown
|
||||
if isCacheStorageOnlyResult(result) {
|
||||
tokenCost = &CostBreakdown{}
|
||||
} else {
|
||||
tokenCost = s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, pricingAt)
|
||||
}
|
||||
if result.SearchCount > 0 {
|
||||
price := groupSearchPricePer1kFromAPIKey(apiKey)
|
||||
if price != nil && *price == 0 {
|
||||
@@ -980,6 +1061,14 @@ func (s *GatewayService) calculateRecordUsageCost(
|
||||
tokenCost.ActualCost += searchCost.ActualCost
|
||||
}
|
||||
}
|
||||
// 显式缓存存储费同为叠加 surcharge,与 token 费用使用同一计费模型与倍率。
|
||||
if storageCost := s.cacheStorageCost(result, billingModel); storageCost > 0 {
|
||||
if tokenCost == nil {
|
||||
tokenCost = &CostBreakdown{}
|
||||
}
|
||||
tokenCost.TotalCost += storageCost
|
||||
tokenCost.ActualCost += storageCost * multiplier
|
||||
}
|
||||
return tokenCost
|
||||
}
|
||||
|
||||
|
||||
@@ -31,9 +31,27 @@ func TestIsForcedUsageBillingRequestID(t *testing.T) {
|
||||
require.True(t, isForcedUsageBillingRequestID("grok-video:task-1"))
|
||||
require.True(t, isForcedUsageBillingRequestID("grok_audio:up-1"))
|
||||
require.True(t, isForcedUsageBillingRequestID("grok_realtime:sess-1"))
|
||||
require.True(t, isForcedUsageBillingRequestID(GeminiCachedContentCreateUsageRequestID("abc")))
|
||||
require.False(t, isForcedUsageBillingRequestID("resp_abc"))
|
||||
}
|
||||
|
||||
func TestGeminiCachedContentUsageRequestIDs(t *testing.T) {
|
||||
t.Parallel()
|
||||
publicID, err := NewGeminiCachedContentPublicID()
|
||||
require.NoError(t, err)
|
||||
create := GeminiCachedContentCreateUsageRequestID(publicID)
|
||||
patch := GeminiCachedContentPatchUsageRequestID(publicID)
|
||||
require.Equal(t, "gcache:create:"+publicID, create)
|
||||
require.True(t, strings.HasPrefix(patch, "gcache:patch:"+publicID+":"))
|
||||
require.NotEqual(t, patch, GeminiCachedContentPatchUsageRequestID(publicID), "每次延长各自唯一")
|
||||
require.LessOrEqual(t, len(create), 64, "usage_logs.request_id 列长 64")
|
||||
require.LessOrEqual(t, len(patch), 64, "usage_logs.request_id 列长 64")
|
||||
|
||||
ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-shared-id")
|
||||
require.Equal(t, create, resolveUsageBillingRequestID(ctx, create), "缓存费用行不被客户端请求 id 覆盖")
|
||||
require.Equal(t, patch, resolveUsageBillingRequestID(ctx, patch))
|
||||
}
|
||||
|
||||
func TestStableGrokAudioBillingRequestID(t *testing.T) {
|
||||
t.Parallel()
|
||||
require.Equal(t, "grok_audio:up-1", StableGrokAudioBillingRequestID("up-1"))
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Gemini 显式上下文缓存存储的内置官方单价(USD / token / 小时,Standard 档)。
|
||||
// 存储价与上下文长度无关(≤200K 与 >200K 同价);Priority 档的更高存储价不在此列。
|
||||
// 来源:https://ai.google.dev/gemini-api/docs/pricing
|
||||
// 与 https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing
|
||||
const (
|
||||
geminiProCacheStoragePrice = 4.5e-6 // $4.50 / 1M tokens / 小时
|
||||
geminiFlashCacheStoragePrice = 1e-6 // $1.00 / 1M tokens / 小时,Flash-Lite 同价
|
||||
|
||||
// 3.6/3.7/3.8 Flash 的优惠价 $0.50 / 1M tokens / 小时,有效期至 2026-12-31,
|
||||
// 2027-01-01 起恢复为 geminiFlashCacheStoragePrice。AI Studio 价目页直接标注该优惠价;
|
||||
// Vertex 价目页存储价标 $1.00,优惠以「按该模型净消费返还 50% 抵扣金」发放,
|
||||
// 官方未明确存储费是否在返还范围内,两类账号统一按优惠价计。
|
||||
geminiFlashPromoCacheStoragePrice = 0.5e-6
|
||||
)
|
||||
|
||||
// geminiCacheStorageModelPrices 是已发布型号的存储单价。模型名等于前缀,或以「前缀-」开头
|
||||
// (-preview、-customtools、-001 等变体)即命中;更具体的前缀排在前面。
|
||||
var geminiCacheStorageModelPrices = []struct {
|
||||
prefix string
|
||||
price float64
|
||||
}{
|
||||
{"gemini-3.8-flash-cyber", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3.8-flash", geminiFlashPromoCacheStoragePrice},
|
||||
{"gemini-3.7-flash", geminiFlashPromoCacheStoragePrice},
|
||||
{"gemini-3.6-flash", geminiFlashPromoCacheStoragePrice},
|
||||
{"gemini-3.5-flash-lite", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3.5-flash", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3.1-flash-lite", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3.1-pro", geminiProCacheStoragePrice},
|
||||
{"gemini-3-flash", geminiFlashCacheStoragePrice},
|
||||
{"gemini-3-pro", geminiProCacheStoragePrice},
|
||||
{"gemini-2.5-flash-lite", geminiFlashCacheStoragePrice},
|
||||
{"gemini-2.5-flash", geminiFlashCacheStoragePrice},
|
||||
{"gemini-2.5-pro", geminiProCacheStoragePrice},
|
||||
{"gemini-2.0-flash-lite", geminiFlashCacheStoragePrice},
|
||||
{"gemini-2.0-flash", geminiFlashCacheStoragePrice},
|
||||
}
|
||||
|
||||
// builtinGeminiCacheStoragePrice 返回内置的显式缓存存储单价:先按已发布型号匹配,未命中时按系列兜底,
|
||||
// Pro 系列取 geminiProCacheStoragePrice,Flash 与 Flash-Lite 系列取 geminiFlashCacheStoragePrice(标准价,
|
||||
// 不套用优惠价)。其他 Gemini 模型与非 Gemini 模型返回 0。
|
||||
func builtinGeminiCacheStoragePrice(model string) float64 {
|
||||
name := lastSegment(strings.ToLower(strings.TrimSpace(model)))
|
||||
if !strings.HasPrefix(name, "gemini-") {
|
||||
return 0
|
||||
}
|
||||
parts := strings.Split(name, "-")
|
||||
lite := slices.Contains(parts, "lite")
|
||||
for _, entry := range geminiCacheStorageModelPrices {
|
||||
// Flash-Lite 变体只匹配 Flash-Lite 条目:gemini-3.8-flash-lite 不能落到 3.8 Flash 的优惠价。
|
||||
if lite != strings.HasSuffix(entry.prefix, "-lite") {
|
||||
continue
|
||||
}
|
||||
if name == entry.prefix || strings.HasPrefix(name, entry.prefix+"-") {
|
||||
return entry.price
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case slices.Contains(parts, "pro"):
|
||||
return geminiProCacheStoragePrice
|
||||
case slices.Contains(parts, "flash"):
|
||||
return geminiFlashCacheStoragePrice
|
||||
}
|
||||
return 0
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user