合并 main 并解决 PR #7828 的测试文件冲突

完整保留临时不可调度规则作用范围测试与 Gemini 显式缓存测试;保留双方生产代码改动。
This commit is contained in:
shaw
2026-10-13 14:59:02 +08:00
453 changed files with 31919 additions and 2085 deletions
+6 -6
View File
@@ -3,16 +3,16 @@ exceptions:
- package: xlsx
advisory: "GHSA-4r6h-8v6p-xvw6"
severity: high
reason: "Admin export only; switched to dynamic import to reduce exposure (CVE-2023-30533)"
mitigation: "Load only on export; restrict export permissions and data scope"
expires_on: "2026-10-06"
reason: "Write-only: admin UsageView lazily imports xlsx solely for XLSX.write export; the advisory requires parsing a crafted file and no XLSX.read path exists (CVE-2023-30533)"
mitigation: "Never parse untrusted spreadsheets with xlsx; no patched version on npm, migrate to SheetJS CDN build or another writer before expiry"
expires_on: "2027-01-06"
owner: "security@your-domain"
- package: xlsx
advisory: "GHSA-5pgg-2g8v-p4x9"
severity: high
reason: "Admin export only; switched to dynamic import to reduce exposure (CVE-2024-22363)"
mitigation: "Load only on export; restrict export permissions and data scope"
expires_on: "2026-10-06"
reason: "Write-only: admin UsageView lazily imports xlsx solely for XLSX.write export; the advisory requires parsing a crafted file and no XLSX.read path exists (CVE-2024-22363)"
mitigation: "Never parse untrusted spreadsheets with xlsx; no patched version on npm, migrate to SheetJS CDN build or another writer before expiry"
expires_on: "2027-01-06"
owner: "security@your-domain"
- package: lodash
advisory: "GHSA-r5fr-rjxr-66jc"
+3 -3
View File
@@ -35,7 +35,7 @@ jobs:
cache-dependency-path: backend/go.sum
- name: Verify Go version
run: |
go version | grep -q 'go1.27.0'
go version | grep -q 'go1.27.2'
- name: Unit tests
working-directory: backend
run: make test-unit
@@ -75,11 +75,11 @@ jobs:
cache-dependency-path: backend/go.sum
- name: Verify Go version
run: |
go version | grep -q 'go1.27.0'
go version | grep -q 'go1.27.2'
- name: golangci-lint
uses: golangci/golangci-lint-action@v9
with:
version: v2.13
version: v2.14.0
args: --timeout=30m
working-directory: backend
+1 -1
View File
@@ -137,7 +137,7 @@ jobs:
check-latest: false
cache: false
- name: Verify Go version
run: go version | grep -q 'go1.27.0'
run: go version | grep -q 'go1.27.2'
- name: Restore target-specific Go cache
id: go-cache
uses: actions/cache/restore@v6
+1 -1
View File
@@ -23,7 +23,7 @@ jobs:
cache-dependency-path: backend/go.sum
- name: Verify Go version
run: |
go version | grep -q 'go1.27.0'
go version | grep -q 'go1.27.2'
- name: Run govulncheck
working-directory: backend
run: |
+1 -1
View File
@@ -8,7 +8,7 @@
# =============================================================================
ARG NODE_IMAGE=node:24-alpine
ARG GOLANG_IMAGE=golang:1.27.0-alpine
ARG GOLANG_IMAGE=golang:1.27.2-alpine
ARG ALPINE_IMAGE=alpine:3.21
ARG POSTGRES_IMAGE=postgres:18-alpine
ARG GOPROXY=https://goproxy.cn,direct
+10 -1
View File
@@ -5,9 +5,16 @@ 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 \
src/constants/__tests__/platforms.spec.ts \
src/components/account/__tests__/credentialsBuilder.platformCatalog.spec.ts \
src/components/account/__tests__/CreateAccountModal.spec.ts \
src/components/account/__tests__/EditAccountModal.spec.ts \
src/components/account/__tests__/credentialsBuilder.spec.ts \
src/components/account/__tests__/OpenCodeGoProtocolRulesEditor.spec.ts \
src/components/keys/__tests__/BulkEditKeysModal.spec.ts \
src/components/admin/user/__tests__/UserPlatformQuotaModal.spec.ts \
src/views/user/__tests__/KeysView.spec.ts \
@@ -18,10 +25,12 @@ FRONTEND_CRITICAL_VITEST := \
src/views/user/__tests__/PaymentResultView.spec.ts \
src/views/user/__tests__/ChannelStatusView.mode.spec.ts \
src/components/user/profile/__tests__/ProfileInfoCard.spec.ts \
src/components/user/profile/__tests__/ProfileIdentityBindingsSection.spec.ts \
src/views/admin/__tests__/SettingsView.spec.ts \
src/features/channel-monitor-v2/__tests__/designSystem.structure.spec.ts \
src/features/channel-monitor-v2/__tests__/monitorFormat.spec.ts \
src/features/channel-monitor-v2/__tests__/monitorZoom.spec.ts
src/features/channel-monitor-v2/__tests__/monitorZoom.spec.ts \
src/components/admin/channel/__tests__/PricingEntryCard.modelDefaultPrice.spec.ts
# 一键编译前后端
build: build-backend build-frontend
+6 -4
View File
@@ -354,8 +354,10 @@ JWT_SECRET=your_jwt_secret_here
TOTP_ENCRYPTION_KEY=your_totp_key_here
# Optional: Admin account
ADMIN_EMAIL=admin@example.com
ADMIN_PASSWORD=your_admin_password
# Leave empty to auto-generate a random email (login username) and password, shown in logs on first startup.
# Avoid guessable values such as admin@example.com: they are brute-force targets.
ADMIN_EMAIL=
ADMIN_PASSWORD=
# Optional: Custom port
SERVER_PORT=8080
@@ -404,9 +406,9 @@ docker compose -f docker-compose.local.yml logs -f sub2api
Open `http://YOUR_SERVER_IP:8080` in your browser.
If admin password was auto-generated, find it in logs:
If the admin email (login username) or password was auto-generated, find them in logs:
```bash
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
docker compose -f docker-compose.local.yml logs sub2api | grep "Generated admin"
```
#### Upgrade
+6 -4
View File
@@ -380,8 +380,10 @@ JWT_SECRET=your_jwt_secret_here
TOTP_ENCRYPTION_KEY=your_totp_key_here
# 可选:管理员账号
ADMIN_EMAIL=admin@example.com
ADMIN_PASSWORD=your_admin_password
# 留空则首次启动时自动生成随机邮箱(登录用户名)和密码,并输出到日志。
# 请勿使用 admin@example.com 这类可被猜中的值,它们是暴力破解的目标。
ADMIN_EMAIL=
ADMIN_PASSWORD=
# 可选:自定义端口
SERVER_PORT=8080
@@ -442,9 +444,9 @@ docker compose -f docker-compose.local.yml logs -f sub2api
在浏览器中打开 `http://你的服务器IP:8080`
如果管理员密码是自动生成的,在日志中查找:
如果管理员邮箱(登录用户名)或密码是自动生成的,在日志中查找:
```bash
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
docker compose -f docker-compose.local.yml logs sub2api | grep "Generated admin"
```
#### 升级
+6 -4
View File
@@ -355,8 +355,10 @@ JWT_SECRET=your_jwt_secret_here
TOTP_ENCRYPTION_KEY=your_totp_key_here
# オプション: 管理者アカウント
ADMIN_EMAIL=admin@example.com
ADMIN_PASSWORD=your_admin_password
# 空欄の場合、初回起動時にランダムなメールアドレス(ログインユーザー名)とパスワードが自動生成され、ログに出力されます。
# admin@example.com のような推測されやすい値は総当たり攻撃の標的になるため避けてください。
ADMIN_EMAIL=
ADMIN_PASSWORD=
# オプション: カスタムポート
SERVER_PORT=8080
@@ -405,9 +407,9 @@ docker compose -f docker-compose.local.yml logs -f sub2api
ブラウザで `http://YOUR_SERVER_IP:8080` を開いてください。
管理者パスワードが自動生成された場合は、ログで確認できます:
管理者メールアドレス(ログインユーザー名)またはパスワードが自動生成された場合は、ログで確認できます:
```bash
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
docker compose -f docker-compose.local.yml logs sub2api | grep "Generated admin"
```
#### アップグレード
+9
View File
@@ -123,6 +123,15 @@ linters:
exclusions:
rules:
# x/net v0.60 deprecated the HTTP/2 compatibility adapters. Keep their
# existing per-protocol idle timeout and PING configuration in this CI fix.
- path: 'internal/(repository/http_upstream(_http2_keepalive_test)?|server/http)\.go$'
linters: [ staticcheck ]
text: 'SA1019: (\(golang\.org/x/net/http2\.(Server|Transport)\)\.|golang\.org/x/net/http2\.(ConfigureServer|ConfigureTransports|Server|Transport)\b)'
# Upstream libraries can still return the legacy typed GOAWAY error.
- path: 'internal/service/openai_codex_models_service(_test)?\.go$'
linters: [ staticcheck ]
text: 'SA1019: golang\.org/x/net/http2\.GoAwayError\b'
# G703/G704 污点分析在测试文件中只会命中 httptest mock 与测试者自设的
# 环境变量路径,不构成攻击面;且该分析跨环境结果不稳定(本地/CI 报告的
# 位置集合不同),逐点 nolint 无法收敛,故按路径豁免。生产代码不豁免,
+1 -1
View File
@@ -1,4 +1,4 @@
FROM golang:1.27.0-alpine
FROM golang:1.27.2-alpine
WORKDIR /app
+1 -1
View File
@@ -1 +1 @@
0.2.13
0.2.15
+7
View File
@@ -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()
+11 -2
View File
@@ -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()
+1
View File
@@ -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
View File
@@ -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(", ")
+5
View File
@@ -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.
+57
View File
@@ -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) {
+18
View File
@@ -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)
}
+7 -6
View File
@@ -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
View File
@@ -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
+65 -47
View File
@@ -906,7 +906,21 @@ func init() {
// compositemodelroute.DefaultTargetPlatform holds the default value on creation for the target_platform field.
compositemodelroute.DefaultTargetPlatform = compositemodelrouteDescTargetPlatform.Default.(string)
// compositemodelroute.TargetPlatformValidator is a validator for the "target_platform" field. It is called by the builders before save.
compositemodelroute.TargetPlatformValidator = compositemodelrouteDescTargetPlatform.Validators[0].(func(string) error)
compositemodelroute.TargetPlatformValidator = func() func(string) error {
validators := compositemodelrouteDescTargetPlatform.Validators
fns := [...]func(string) error{
validators[0].(func(string) error),
validators[1].(func(string) error),
}
return func(target_platform string) error {
for _, fn := range fns {
if err := fn(target_platform); err != nil {
return err
}
}
return nil
}
}()
// compositemodelrouteDescUpstreamModel is the schema descriptor for upstream_model field.
compositemodelrouteDescUpstreamModel := compositemodelrouteFields[4].Descriptor()
// compositemodelroute.DefaultUpstreamModel holds the default value on creation for the upstream_model field.
@@ -1023,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()
@@ -1,6 +1,8 @@
package schema
import (
"fmt"
"github.com/Wei-Shaw/sub2api/ent/schema/mixins"
"github.com/Wei-Shaw/sub2api/internal/domain"
@@ -45,6 +47,13 @@ func (CompositeModelRoute) Fields() []ent.Field {
field.String("target_platform").
MaxLen(50).
Default(domain.PlatformOpenAI).
Validate(func(s string) error {
// 目标平台须为平台清单中的具体平台;数据库不再维护 CHECK 约束。
if !domain.IsConcretePlatform(s) {
return fmt.Errorf("target_platform %q is not a concrete platform", s)
}
return nil
}).
Comment("Concrete provider platform."),
field.String("upstream_model").
MaxLen(200).
+5
View File
@@ -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),
+5 -7
View File
@@ -11,6 +11,7 @@ import (
"entgo.io/ent/schema/field"
"entgo.io/ent/schema/index"
"github.com/Wei-Shaw/sub2api/ent/schema/mixins"
"github.com/Wei-Shaw/sub2api/internal/domain"
)
// UserPlatformQuota holds the schema definition for per-user per-platform quota.
@@ -38,15 +39,12 @@ func (UserPlatformQuota) Fields() []ent.Field {
MaxLen(32).
NotEmpty().
Validate(func(s string) error {
// 注意:平台列表的单一权威源为 service.AllowedQuotaPlatforms;
// 此处为 ent 构建期约束,需与 service.AllowedQuotaPlatforms 保持同步。
switch s {
case "anthropic", "openai", "gemini", "antigravity", "grok",
"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe":
return nil
default:
// 平台白名单为平台清单(domain/platforms.go)中的全部具体平台,
// 与 service.AllowedQuotaPlatforms 同源;数据库不再维护 CHECK 约束。
if !domain.IsConcretePlatform(s) {
return fmt.Errorf("platform %q is not allowed", s)
}
return nil
}),
// 日 / 周 / 月 USD 上限:
+10 -10
View File
@@ -1,6 +1,6 @@
module github.com/Wei-Shaw/sub2api
go 1.27.0
go 1.27.2
require (
entgo.io/ent v0.14.5
@@ -19,6 +19,7 @@ require (
github.com/coder/websocket v1.8.14
github.com/dgraph-io/ristretto v0.2.0
github.com/gin-gonic/gin v1.9.1
github.com/go-playground/validator/v10 v10.14.0
github.com/go-webauthn/webauthn v0.17.4
github.com/golang-jwt/jwt/v5 v5.3.1
github.com/google/uuid v1.6.0
@@ -50,12 +51,12 @@ require (
github.com/wechatpay-apiv3/wechatpay-go v0.2.21
github.com/zeromicro/go-zero v1.9.4
go.uber.org/zap v1.24.0
golang.org/x/crypto v0.55.0
golang.org/x/crypto v0.57.0
golang.org/x/image v0.41.0
golang.org/x/mod v0.38.0
golang.org/x/net v0.58.0
golang.org/x/sync v0.22.0
golang.org/x/term v0.45.0
golang.org/x/mod v0.41.0
golang.org/x/net v0.60.0
golang.org/x/sync v0.23.0
golang.org/x/term v0.46.0
google.golang.org/grpc v1.83.2
google.golang.org/protobuf v1.36.11
gopkg.in/natefinch/lumberjack.v2 v2.2.1
@@ -120,7 +121,6 @@ require (
github.com/go-openapi/inflect v0.19.0 // indirect
github.com/go-playground/locales v0.14.1 // indirect
github.com/go-playground/universal-translator v0.18.1 // indirect
github.com/go-playground/validator/v10 v10.14.0 // indirect
github.com/go-viper/mapstructure/v2 v2.5.0 // indirect
github.com/go-webauthn/x v0.2.6 // indirect
github.com/goccy/go-json v0.10.2 // indirect
@@ -201,10 +201,10 @@ require (
go.uber.org/multierr v1.9.0 // indirect
golang.org/x/arch v0.3.0 // indirect
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect
golang.org/x/sys v0.47.0 // indirect
golang.org/x/text v0.41.0 // indirect
golang.org/x/sys v0.48.0 // indirect
golang.org/x/text v0.42.0 // indirect
golang.org/x/time v0.12.0 // indirect
golang.org/x/tools v0.48.0 // indirect
golang.org/x/tools v0.50.0 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect
gopkg.in/ini.v1 v1.67.0 // indirect
modernc.org/libc v1.67.6 // indirect
+18
View File
@@ -564,6 +564,8 @@ golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v
golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
golang.org/x/crypto v0.55.0 h1:+KWHjbgOaAQ66dh/YlkZKHlz9ZUlq61AFirAR9ntP8M=
golang.org/x/crypto v0.55.0/go.mod h1:uq0V9dE/fzQuJtbnL+2EhWOE63vo164FY8xqEnV9xis=
golang.org/x/crypto v0.57.0 h1:3ZVCjf8Ggz7zneR/EHRVx68Ctf+2pmIMP2UFhh9cC6M=
golang.org/x/crypto v0.57.0/go.mod h1:Fdz0i5U6CoizGwLda9DttjSk6qlZo25zYNtR+ycvuZA=
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY=
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70=
@@ -580,6 +582,8 @@ golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
golang.org/x/mod v0.38.0 h1:MECBjubtXD7yj4HrhIUcywNaGeNVUdfVnxmPajOk4yk=
golang.org/x/mod v0.38.0/go.mod h1:V6Xz0pq8TQ3dGqVQ1FVHuelZpAL0uNhSkk9ogYP3c40=
golang.org/x/mod v0.41.0 h1:qJmnOUb4YB+FsEuM3HcWucdZASCPGhsX6uljO6pog0c=
golang.org/x/mod v0.41.0/go.mod h1:Ek9pY8RKWXwsWvd3rQiHYtMqkjSUV+s1Rj7j4H5Ur6o=
golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -603,6 +607,8 @@ golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
golang.org/x/net v0.58.0 h1:ynWG7rqYi4ccpTEuPZ2QGWHktVEM9DMCj9yzDE0Q7To=
golang.org/x/net v0.58.0/go.mod h1:YwCddHnFlT7eLQqVprV19OnhLGtc5xOKgE0RyqgfWAU=
golang.org/x/net v0.60.0 h1:79p50tfZlm0J9YfoDsSi639qSXNGVwEzOPLCxM2FsYU=
golang.org/x/net v0.60.0/go.mod h1:2DA/G1UfVbCpQPeWTmMPGY7Cs2PkBkwu743bVX5PIVg=
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
@@ -616,6 +622,8 @@ golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk=
golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0=
golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
@@ -649,6 +657,8 @@ golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo=
golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og=
golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
@@ -663,6 +673,8 @@ golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0=
golang.org/x/term v0.45.0 h1:NwWyBmoJCbfTHpxrWoZ9C6/VxOf7ic219I8xZZFdrf0=
golang.org/x/term v0.45.0/go.mod h1:9aqxs0blBcrm/n0L9QW0aRVD+ktan8ssZromtqJC43w=
golang.org/x/term v0.46.0 h1:3+OXuTbaKDgwk8jTi3aSLHRlmWqHEUDUtxnbFigO4YE=
golang.org/x/term v0.46.0/go.mod h1:+K02xbkittuwc0Am4abfA3Fc+XRGXkvBXNO88NCXPoc=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
@@ -675,6 +687,8 @@ golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8=
golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M=
golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI=
golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E=
golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE=
golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
@@ -691,6 +705,10 @@ golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
golang.org/x/tools v0.48.0 h1:3+hClM1aLL5mjMKm5ovokw9epgRXPuu2tILgismM6RE=
golang.org/x/tools v0.48.0/go.mod h1:08xX0orndb/F7jJxGDicx061tyd5pcMto75YMAXr6lk=
golang.org/x/tools v0.49.0 h1:3NI7VXzL9+1WZD52Dx2ttoPwD5DWrFGpl9mFZDlmisI=
golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo=
golang.org/x/tools v0.50.0 h1:c2ifzfcuY7L90lZ2aKd8S4K2NpASF08SZx9ZuJkHmSU=
golang.org/x/tools v0.50.0/go.mod h1:7ulVMw3831Mwi5EZD6RomGyffr4VFjuNYXf2BbCEAV0=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+3
View File
@@ -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 处理。
@@ -2527,6 +2529,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) // 最大空闲连接总数(高并发场景可调大)
+8 -2
View File
@@ -30,9 +30,15 @@ const (
PlatformMiniMax = "minimax" // MiniMax (M 系列)
PlatformTypeSafe = "typesafe" // TypeSafe AI System One (Jev)
// PlatformOpenCodeGo 是 OpenCode 平台(账号类型 Zen 按量 / Go 订阅)。
// 值保持 opencode_go 以兼容已落库的分组、配额与 Composite 路由 CHECK。
// 值保持 opencode_go 以兼容已落库的账号、分组、配额与 Composite 路由。
PlatformOpenCodeGo = "opencode_go"
PlatformComposite = "composite"
// PlatformCommandCode 是 Command Code Provider API(多模型聚合,积分计费,
// 订阅套餐另有 5 小时 / 每周滚动窗口)。
PlatformCommandCode = "command_code"
// PlatformCline 是 Cline API(多模型聚合,OpenAI 兼容 Chat Completions;按量积分计费,
// 另有 ClinePass 订阅)。
PlatformCline = "cline"
PlatformComposite = "composite"
)
// Account mode constants 区分国产供应商的「按量付费(余额)」与「Coding Plan」两种接入方式。
+113
View File
@@ -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
}
+107
View File
@@ -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)
}
}
})
}
}
+137
View File
@@ -0,0 +1,137 @@
package domain
// 平台清单:平台的单一数据来源。"所有平台"类列表(分组平台、用户额度、组合
// 路由目标、调度快照、定价匹配等)与"某网关族"类列表都从这里派生;新增平台
// 登记一条即可在这些位置生效,无需逐处修改或新增数据库约束迁移。
//
// 依赖具体实现的能力(OAuth、渠道监控、余额 / Coding Plan 探测、用量窗口、
// Codex 模型目录等)仍由各自的代码显式列出,不在此处声明。
// PlatformGateway 标识平台请求由哪个网关服务处理。
type PlatformGateway string
const (
// PlatformGatewayAnthropic 由 GatewayService 处理(Claude Messages 协议,以及 TypeSafe 的 /v1/systemone)。
PlatformGatewayAnthropic PlatformGateway = "anthropic"
// PlatformGatewayOpenAI 由 OpenAIGatewayService 处理(OpenAI、Grok 与多协议 API Key 供应商)。
PlatformGatewayOpenAI PlatformGateway = "openai"
// PlatformGatewayGemini 由 Gemini 兼容网关处理。
PlatformGatewayGemini PlatformGateway = "gemini"
// PlatformGatewayAntigravity 由 Antigravity 网关处理。
PlatformGatewayAntigravity PlatformGateway = "antigravity"
)
// PlatformSpec 描述一个具体(非 composite)平台。
type PlatformSpec struct {
ID string `json:"id"`
DisplayName string `json:"display_name"`
Gateway PlatformGateway `json:"gateway"`
// CNProvider 标识国产单一厂商(Kimi / 智谱 / DeepSeek / MiniMax),共享国产
// 供应商的余额 / Coding Plan 等能力。
CNProvider bool `json:"cn_provider"`
// LiteLLMProvider 是定价目录中的 provider 名;空表示不支持从定价目录同步模型。
LiteLLMProvider string `json:"litellm_provider,omitempty"`
}
// platformList 的顺序即平台的展示顺序。
var platformList = []PlatformSpec{
{ID: PlatformAnthropic, DisplayName: "Anthropic", Gateway: PlatformGatewayAnthropic, LiteLLMProvider: "anthropic"},
{ID: PlatformOpenAI, DisplayName: "OpenAI", Gateway: PlatformGatewayOpenAI, LiteLLMProvider: "openai"},
{ID: PlatformGemini, DisplayName: "Gemini", Gateway: PlatformGatewayGemini, LiteLLMProvider: "gemini"},
{ID: PlatformAntigravity, DisplayName: "Antigravity", Gateway: PlatformGatewayAntigravity, LiteLLMProvider: "anthropic"},
{ID: PlatformGrok, DisplayName: "Grok", Gateway: PlatformGatewayOpenAI, LiteLLMProvider: "xai"},
{ID: PlatformKimi, DisplayName: "Kimi", Gateway: PlatformGatewayOpenAI, CNProvider: true, LiteLLMProvider: "moonshot"},
{ID: PlatformZhipu, DisplayName: "Zhipu GLM", Gateway: PlatformGatewayOpenAI, CNProvider: true, LiteLLMProvider: "zhipu"},
{ID: PlatformDeepseek, DisplayName: "DeepSeek", Gateway: PlatformGatewayOpenAI, CNProvider: true, LiteLLMProvider: "deepseek"},
{ID: PlatformMiniMax, DisplayName: "MiniMax", Gateway: PlatformGatewayOpenAI, CNProvider: true, LiteLLMProvider: "minimax"},
{ID: PlatformOpenCodeGo, DisplayName: "OpenCode", Gateway: PlatformGatewayOpenAI, LiteLLMProvider: "opencode-go"},
{ID: PlatformTypeSafe, DisplayName: "TypeSafe / Jev", Gateway: PlatformGatewayAnthropic, LiteLLMProvider: "typesafe"},
{ID: PlatformCommandCode, DisplayName: "Command Code", Gateway: PlatformGatewayOpenAI},
{ID: PlatformCline, DisplayName: "Cline", Gateway: PlatformGatewayOpenAI},
}
var platformIndex = func() map[string]int {
index := make(map[string]int, len(platformList))
for i, spec := range platformList {
index[spec.ID] = i
}
return index
}()
// Platforms 返回全部具体平台(副本),按展示顺序。
func Platforms() []PlatformSpec {
out := make([]PlatformSpec, len(platformList))
copy(out, platformList)
return out
}
// LookupPlatform 返回具体平台的登记信息;composite 与未知值返回 false。
func LookupPlatform(id string) (PlatformSpec, bool) {
i, ok := platformIndex[id]
if !ok {
return PlatformSpec{}, false
}
return platformList[i], true
}
// IsConcretePlatform 报告 id 是否为已登记的具体平台(不含 composite)。
func IsConcretePlatform(id string) bool {
_, ok := platformIndex[id]
return ok
}
// IsGroupPlatform 报告 id 是否可作为分组平台(具体平台或 composite)。
func IsGroupPlatform(id string) bool {
return id == PlatformComposite || IsConcretePlatform(id)
}
// ConcretePlatformIDs 返回全部具体平台标识,按展示顺序。
func ConcretePlatformIDs() []string {
return PlatformIDsWhere(func(PlatformSpec) bool { return true })
}
// PlatformIDsWhere 返回满足条件的具体平台标识,按展示顺序。
func PlatformIDsWhere(match func(PlatformSpec) bool) []string {
out := make([]string, 0, len(platformList))
for _, spec := range platformList {
if match(spec) {
out = append(out, spec.ID)
}
}
return out
}
// compositePrecedenceHead 是组合分组查找顺序的固定前缀(历史顺序:Gemini 先于 OpenAI)。
var compositePrecedenceHead = []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI}
// CompositePrecedencePlatformIDs 返回按平台依次回退时的顺序:组合分组在请求目标
// 解析前逐平台查找定价 / 模型映射(先命中者生效)、拼接默认模型列表,以及调度快照
// 遍历。沿用历史顺序 anthropic、gemini、openai,其余按展示顺序,新登记的平台追加在末尾。
func CompositePrecedencePlatformIDs() []string {
out := make([]string, 0, len(platformList))
seen := make(map[string]struct{}, len(platformList))
for _, id := range compositePrecedenceHead {
if IsConcretePlatform(id) {
out = append(out, id)
seen[id] = struct{}{}
}
}
for _, spec := range platformList {
if _, ok := seen[spec.ID]; !ok {
out = append(out, spec.ID)
}
}
return out
}
// UsesOpenAIGateway 报告平台是否由 OpenAI 网关处理。
func UsesOpenAIGateway(id string) bool {
spec, ok := LookupPlatform(id)
return ok && spec.Gateway == PlatformGatewayOpenAI
}
// IsCNProviderPlatform 报告平台是否为国产单一厂商。
func IsCNProviderPlatform(id string) bool {
spec, ok := LookupPlatform(id)
return ok && spec.CNProvider
}
+74
View File
@@ -0,0 +1,74 @@
package domain
import (
"testing"
"github.com/stretchr/testify/require"
)
// 重构前各处手写的平台列表,作为平台清单派生结果的等价基准;重构后新登记的平台
// (Command Code、Cline)按同类平台(OpenCode)的位置补入。
var (
legacyDisplayOrder = []string{
PlatformAnthropic, PlatformOpenAI, PlatformGemini, PlatformAntigravity, PlatformGrok,
PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe,
PlatformCommandCode,
PlatformCline,
}
legacyCompositePrecedence = []string{
PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok,
PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe,
PlatformCommandCode,
PlatformCline,
}
legacyOpenAIGateway = []string{
PlatformOpenAI, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo,
PlatformCommandCode,
PlatformCline,
}
legacyCNProviders = []string{PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax}
legacyLiteLLMByPlatf = map[string]string{
PlatformAnthropic: "anthropic", PlatformOpenAI: "openai", PlatformGemini: "gemini",
PlatformAntigravity: "anthropic", PlatformGrok: "xai", PlatformKimi: "moonshot",
PlatformZhipu: "zhipu", PlatformDeepseek: "deepseek", PlatformMiniMax: "minimax",
PlatformOpenCodeGo: "opencode-go", PlatformTypeSafe: "typesafe", PlatformCommandCode: "", PlatformCline: "",
}
)
func TestPlatformListMatchesLegacyLists(t *testing.T) {
require.Equal(t, legacyDisplayOrder, ConcretePlatformIDs())
require.Equal(t, legacyCompositePrecedence, CompositePrecedencePlatformIDs())
require.Equal(t, legacyOpenAIGateway, PlatformIDsWhere(func(spec PlatformSpec) bool {
return spec.Gateway == PlatformGatewayOpenAI
}))
require.Equal(t, legacyCNProviders, PlatformIDsWhere(func(spec PlatformSpec) bool { return spec.CNProvider }))
for platform, provider := range legacyLiteLLMByPlatf {
spec, ok := LookupPlatform(platform)
require.True(t, ok, platform)
require.Equal(t, provider, spec.LiteLLMProvider, platform)
}
}
func TestPlatformListPredicates(t *testing.T) {
for _, platform := range legacyDisplayOrder {
require.True(t, IsConcretePlatform(platform), platform)
require.True(t, IsGroupPlatform(platform), platform)
spec, ok := LookupPlatform(platform)
require.True(t, ok, platform)
require.NotEmpty(t, spec.DisplayName, platform)
require.NotEmpty(t, spec.Gateway, platform)
}
require.False(t, IsConcretePlatform(PlatformComposite))
require.True(t, IsGroupPlatform(PlatformComposite))
for _, invalid := range []string{"", "moonshot", "Kimi", "openai ", "glm", "bogus"} {
require.False(t, IsConcretePlatform(invalid), invalid)
require.False(t, IsGroupPlatform(invalid), invalid)
require.False(t, UsesOpenAIGateway(invalid), invalid)
}
}
func TestPlatformsReturnsCopy(t *testing.T) {
platforms := Platforms()
platforms[0].ID = "mutated"
require.Equal(t, PlatformAnthropic, Platforms()[0].ID)
}
@@ -1226,7 +1226,7 @@ func (h *AccountHandler) Update(c *gin.Context) {
// 网关会按"现状即证据"默认走 Responses。
func (h *AccountHandler) scheduleOpenAIResponsesProbe(account *service.Account) {
if account == nil || account.Type != service.AccountTypeAPIKey ||
(account.Platform != service.PlatformOpenAI && !service.IsCNProvider(account.Platform)) {
(account.Platform != service.PlatformOpenAI && !account.RoutesProtocolByInbound()) {
return
}
if h.accountTestService == nil {
@@ -2894,10 +2894,9 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
return
}
// Handle Antigravity accounts: return Claude + Gemini models
// Explicit account mappings expose their request-side names to connectivity tests.
if account.Platform == service.PlatformAntigravity {
// 直接复用 antigravity.DefaultModels(),与 /v1/models 端点保持同步
response.Success(c, antigravity.DefaultModels())
response.Success(c, antigravityAccountTestModels(account.Credentials["model_mapping"]))
return
}
@@ -2998,6 +2997,48 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
response.Success(c, models)
}
// antigravityAccountTestModels uses the stored mapping rather than GetModelMapping,
// which supplies defaults and compatibility aliases that the administrator did not configure.
func antigravityAccountTestModels(rawMapping any) []antigravity.ClaudeModel {
var mappedIDs []string
switch mapping := rawMapping.(type) {
case map[string]any:
for id := range mapping {
if strings.TrimSpace(id) != "" {
mappedIDs = append(mappedIDs, id)
}
}
case map[string]string:
for id := range mapping {
if strings.TrimSpace(id) != "" {
mappedIDs = append(mappedIDs, id)
}
}
}
if len(mappedIDs) == 0 {
return antigravity.DefaultModels()
}
sort.Strings(mappedIDs)
defaultByID := make(map[string]antigravity.ClaudeModel)
for _, model := range antigravity.DefaultModels() {
defaultByID[model.ID] = model
}
models := make([]antigravity.ClaudeModel, 0, len(mappedIDs))
for _, id := range mappedIDs {
if model, ok := defaultByID[id]; ok {
models = append(models, model)
continue
}
models = append(models, antigravity.ClaudeModel{
ID: id,
Type: "model",
DisplayName: id,
})
}
return models
}
// SyncUpstreamModels handles syncing live supported models from an account's upstream.
// POST /api/v1/admin/accounts/:id/models/sync-upstream
func (h *AccountHandler) SyncUpstreamModels(c *gin.Context) {
@@ -10,6 +10,7 @@ import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
@@ -330,6 +331,49 @@ func TestAccountHandlerGetAvailableModels_GeminiGoogleOneUsesConservativeCatalog
require.NotContains(t, ids, "gemini-2.5-flash-image")
}
func TestAccountHandlerGetAvailableModels_Antigravity(t *testing.T) {
defaults := antigravity.DefaultModels()
known := defaults[0]
custom := antigravity.ClaudeModel{ID: "claude-custom-alias", Type: "model", DisplayName: "claude-custom-alias"}
cases := []struct {
name string
credentials map[string]any
want []antigravity.ClaudeModel
}{
{name: "nil credentials", want: defaults},
{name: "missing mapping", credentials: map[string]any{}, want: defaults},
{name: "empty JSON mapping", credentials: map[string]any{"model_mapping": map[string]any{}}, want: defaults},
{name: "empty string mapping", credentials: map[string]any{"model_mapping": map[string]string{}}, want: defaults},
{name: "blank key", credentials: map[string]any{"model_mapping": map[string]any{" ": "upstream-model"}}, want: defaults},
{name: "JSON mapping keys and metadata", credentials: map[string]any{"model_mapping": map[string]any{
known.ID: known.ID, custom.ID: "upstream-model",
}}, want: []antigravity.ClaudeModel{custom, known}},
{name: "string mapping keys and metadata", credentials: map[string]any{"model_mapping": map[string]string{
known.ID: known.ID, custom.ID: "upstream-model",
}}, want: []antigravity.ClaudeModel{custom, known}},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
for _, accountType := range []string{service.AccountTypeOAuth, service.AccountTypeAPIKey} {
t.Run(accountType, func(t *testing.T) {
svc := &availableModelsAdminService{
stubAdminService: newStubAdminService(),
account: service.Account{ID: 51, Platform: service.PlatformAntigravity, Type: accountType, Credentials: tc.credentials},
}
rec := httptest.NewRecorder()
setupAvailableModelsRouter(svc).ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/51/models", nil))
require.Equal(t, http.StatusOK, rec.Code)
var resp struct {
Data []antigravity.ClaudeModel `json:"data"`
}
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, tc.want, resp.Data)
})
}
})
}
}
func TestAccountHandlerSyncUpstreamModels_ConfigErrorReturnsBadRequest(t *testing.T) {
svc := &availableModelsAdminService{
stubAdminService: newStubAdminService(),
@@ -5,6 +5,7 @@ import (
"strconv"
"strings"
"github.com/Wei-Shaw/sub2api/internal/domain"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
@@ -631,20 +632,14 @@ func (h *ChannelHandler) GetModelDefaultPricing(c *gin.Context) {
})
}
// platformToLiteLLMProvider maps a channel platform name to the corresponding
// LiteLLM provider string used as the key in the pricing catalog.
var platformToLiteLLMProvider = map[string]string{
service.PlatformAnthropic: "anthropic",
service.PlatformOpenAI: "openai",
service.PlatformGemini: "gemini",
service.PlatformAntigravity: "anthropic",
service.PlatformGrok: "xai",
service.PlatformKimi: "moonshot",
service.PlatformZhipu: "zhipu",
service.PlatformDeepseek: "deepseek",
service.PlatformMiniMax: "minimax",
service.PlatformOpenCodeGo: "opencode-go",
service.PlatformTypeSafe: "typesafe",
// platformLiteLLMProvider maps a channel platform name to the corresponding
// LiteLLM provider string used as the key in the pricing catalog (platform list).
func platformLiteLLMProvider(platform string) (string, bool) {
spec, ok := domain.LookupPlatform(platform)
if !ok || spec.LiteLLMProvider == "" {
return "", false
}
return spec.LiteLLMProvider, true
}
// SyncPricingModels 返回 LiteLLM 定价目录中指定平台的最新模型列表
@@ -657,7 +652,7 @@ func (h *ChannelHandler) SyncPricingModels(c *gin.Context) {
return
}
provider, ok := platformToLiteLLMProvider[platform]
provider, ok := platformLiteLLMProvider(platform)
if !ok {
response.ErrorFrom(c, infraerrors.BadRequest("UNSUPPORTED_PLATFORM",
fmt.Sprintf("unsupported platform: %s", platform)).
@@ -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,14 +180,15 @@ 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"`
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go typesafe composite"`
I18n service.GroupI18n `json:"i18n"` // 名称与描述的界面语言译文
Platform string `json:"platform" binding:"omitempty,group_platform"`
RateMultiplier float64 `json:"rate_multiplier"`
IsExclusive bool `json:"is_exclusive"`
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
@@ -258,7 +262,8 @@ type CreateGroupRequest struct {
type UpdateGroupRequest struct {
Name string `json:"name"`
Description *string `json:"description"`
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go typesafe composite"`
I18n *service.GroupI18n `json:"i18n"` // 缺省不修改,传入则整体替换
Platform string `json:"platform" binding:"omitempty,group_platform"`
RateMultiplier *float64 `json:"rate_multiplier"`
IsExclusive *bool `json:"is_exclusive"`
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
@@ -332,7 +337,7 @@ type UpdateGroupRequest struct {
type CompositeRouteRequest struct {
PublicModel string `json:"public_model" binding:"required"`
MatchType string `json:"match_type" binding:"omitempty,oneof=exact prefix"`
TargetPlatform string `json:"target_platform" binding:"required,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go typesafe"`
TargetPlatform string `json:"target_platform" binding:"required,concrete_platform"`
UpstreamModel string `json:"upstream_model"`
Endpoint string `json:"endpoint" binding:"omitempty,oneof=any messages count_tokens responses chat_completions embeddings images gemini"`
Priority int `json:"priority"`
@@ -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,
@@ -29,6 +29,7 @@ func TestGroupPlatformBinding_AllowedPlatforms(t *testing.T) {
"anthropic", "openai", "gemini", "antigravity", "grok",
"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "composite",
"typesafe",
"command_code", "cline",
}
for _, platform := range allowed {
t.Run("create_"+platform, func(t *testing.T) {
@@ -73,7 +74,7 @@ func TestGroupPlatformBinding_RejectsInvalidPlatforms(t *testing.T) {
}
func TestCompositeRouteTargetPlatform_AllowsConcreteProviders(t *testing.T) {
for _, platform := range []string{"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe"} {
for _, platform := range []string{"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe", "command_code", "cline"} {
var req CompositeRouteRequest
body := fmt.Sprintf(`{"public_model":"m","target_platform":%q}`, platform)
require.NoError(t, bindGroupPlatformJSON(t, &req, body))
@@ -0,0 +1,30 @@
package admin
import (
"fmt"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/gin-gonic/gin/binding"
"github.com/go-playground/validator/v10"
)
// 平台相关的请求绑定校验标签,由平台清单(domain/platforms.go)驱动,替代在
// binding 标签里手写 oneof 白名单(新增平台时易漏加,导致分组无法创建):
// - group_platform:已登记的具体平台或 composite;
// - concrete_platform:已登记的具体平台。
func init() {
engine, ok := binding.Validator.Engine().(*validator.Validate)
if !ok {
panic("platform validation requires a go-playground validator engine")
}
if err := engine.RegisterValidation("group_platform", func(fl validator.FieldLevel) bool {
return domain.IsGroupPlatform(fl.Field().String())
}); err != nil {
panic(fmt.Errorf("register group_platform validation: %w", err))
}
if err := engine.RegisterValidation("concrete_platform", func(fl validator.FieldLevel) bool {
return domain.IsConcretePlatform(fl.Field().String())
}); err != nil {
panic(fmt.Errorf("register concrete_platform validation: %w", err))
}
}
+5 -5
View File
@@ -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)
}
+1
View File
@@ -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,
+4
View File
@@ -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"`
+8 -1
View File
@@ -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},
+37 -22
View File
@@ -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
@@ -1166,7 +1168,7 @@ func (h *GatewayHandler) Models(c *gin.Context) {
}
if platform == service.PlatformComposite {
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID, true)
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID, "", true)
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
source := availableModels
if len(source) == 0 {
@@ -1265,19 +1267,19 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
platform = group.Platform
}
if platform == service.PlatformComposite {
availableModels := h.compositeAvailableModels(ctx, groupID, false)
availableModels := h.compositeAvailableModels(ctx, groupID, service.CompositeRouteEndpointResponses, false)
fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite)
models := availableModels
if len(models) == 0 {
models = fallbackModels
}
if group.ModelAllowlistEnabled() {
source := availableModels
if len(source) == 0 {
source = fallbackModels
}
return group.ModelAllowlist.FilterForListing(source)
models = group.ModelAllowlist.FilterForListing(models)
}
if len(availableModels) > 0 {
return availableModels
if filtered, err := h.gatewayService.FilterCompositeCodexModels(ctx, group.ID, models); err == nil {
return filtered
}
return fallbackModels
return models
}
availableModels := h.gatewayService.GetAvailableModels(ctx, groupID, platform)
@@ -1294,14 +1296,14 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
// compositeAvailableModels lists the models the composite group can serve.
// includeSystemOne adds TypeSafe models, which only work through /v1/systemone;
// LLM client catalogs (Codex) must exclude them.
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64, includeSystemOne bool) []string {
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64, endpoint string, includeSystemOne bool) []string {
if h == nil || h.gatewayService == nil {
return nil
}
seen := make(map[string]struct{})
models := make([]string, 0)
schedulablePlatforms := h.gatewayService.GetSchedulablePlatforms(ctx, groupID)
for _, platform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax, service.PlatformOpenCodeGo, service.PlatformTypeSafe} {
for _, platform := range domain.CompositePrecedencePlatformIDs() {
if platform == service.PlatformTypeSafe && !includeSystemOne {
continue
}
@@ -1325,6 +1327,16 @@ func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *
models = append(models, model)
}
}
// A route can expose a public ID that no account model mapping contains.
// On lookup failure, retain the existing account-derived catalog only.
if routeModels, err := h.gatewayService.GetCompositeRouteModels(ctx, groupID, endpoint, includeSystemOne); err == nil {
for _, model := range routeModels {
if _, ok := seen[model]; !ok {
seen[model] = struct{}{}
models = append(models, model)
}
}
}
return models
}
@@ -1493,10 +1505,13 @@ func defaultModelIDsForPlatform(platform string) []string {
case service.PlatformComposite:
ids := make([]string, 0)
seen := make(map[string]struct{})
// TypeSafe is deliberately absent: jev-latest only works through
// /v1/systemone, so the static fallback never advertises it to LLM
// clients. compositeAvailableModels lists it when the group can serve it.
for _, concretePlatform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax, service.PlatformOpenCodeGo} {
for _, concretePlatform := range domain.CompositePrecedencePlatformIDs() {
// TypeSafe is deliberately skipped: jev-latest only works through
// /v1/systemone, so the static fallback never advertises it to LLM
// clients. compositeAvailableModels lists it when the group can serve it.
if concretePlatform == service.PlatformTypeSafe {
continue
}
for _, id := range defaultModelIDsForPlatform(concretePlatform) {
if _, ok := seen[id]; ok {
continue
+262 -2
View File
@@ -3,6 +3,7 @@ package handler
import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"strings"
@@ -22,6 +23,16 @@ type gatewayModelsAccountRepoStub struct {
byGroup map[int64][]service.Account
}
type gatewayModelsRouteRepoStub struct {
service.CompositeModelRouteRepository
routes []service.CompositeModelRoute
err error
}
func (s *gatewayModelsRouteRepoStub) ListByGroup(_ context.Context, _ int64, _ bool) ([]service.CompositeModelRoute, error) {
return s.routes, s.err
}
type gatewayModelsResponseForTest struct {
Object string `json:"object"`
Data []gatewayModelItemForTest `json:"data"`
@@ -81,16 +92,80 @@ func (s *gatewayModelsAccountRepoStub) ListModelAvailabilityCandidates(ctx conte
return s.ListSchedulableByGroupID(ctx, *groupID)
}
func newGatewayModelsHandlerForTest(repo service.AccountRepository) *GatewayHandler {
func newGatewayModelsHandlerForTest(repo service.AccountRepository, routeRepo ...service.CompositeModelRouteRepository) *GatewayHandler {
var resolver *service.CompositeRouteResolver
if len(routeRepo) > 0 {
resolver = service.NewCompositeRouteResolver(routeRepo[0])
}
return &GatewayHandler{
gatewayService: service.NewGatewayService(
repo,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, resolver, nil, nil,
),
}
}
func TestGatewayCompositeRouteModelsAppearInBothCatalogs(t *testing.T) {
gin.SetMode(gin.TestMode)
const groupID int64 = 7784
accounts := &gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{
groupID: {{ID: 1, Platform: service.PlatformOpenAI, Credentials: map[string]any{
"model_mapping": map[string]any{"mapped-model": "upstream"},
}}},
}}
routes := &gatewayModelsRouteRepoStub{routes: []service.CompositeModelRoute{
{PublicModel: "route-only", MatchType: service.CompositeRouteMatchExact, Enabled: true},
{PublicModel: "route-only", MatchType: service.CompositeRouteMatchExact, Enabled: true},
{PublicModel: "mapped-model", MatchType: service.CompositeRouteMatchExact, Enabled: true},
{PublicModel: "image-only", MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointImages, Enabled: true},
{PublicModel: "prefix-", MatchType: service.CompositeRouteMatchPrefix, Enabled: true},
{PublicModel: "disabled", MatchType: service.CompositeRouteMatchExact, Enabled: false},
}}
for _, tc := range []struct {
name string
repo service.CompositeModelRouteRepository
allowlist []string
want []string
}{
{"routes", routes, nil, []string{"mapped-model", "route-only", "image-only"}},
{"allowlist", routes, []string{"route-only", "disabled", "unknown"}, []string{"route-only"}},
{"lookup failure", &gatewayModelsRouteRepoStub{err: errors.New("unavailable")}, nil, []string{"mapped-model"}},
{"allowlisted lookup failure", &gatewayModelsRouteRepoStub{err: errors.New("unavailable")}, []string{"mapped-model", "unknown"}, []string{"mapped-model"}},
{"no route repo", nil, nil, []string{"mapped-model"}},
} {
t.Run(tc.name, func(t *testing.T) {
h := newGatewayModelsHandlerForTest(accounts, tc.repo)
group := &service.Group{ID: groupID, Platform: service.PlatformComposite}
if tc.allowlist != nil {
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: tc.allowlist}
}
for _, endpoint := range []string{"models", "codex"} {
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group})
if endpoint == "models" {
h.Models(c)
var got gatewayModelsResponseForTest
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
require.Equal(t, tc.want, modelIDsForTest(got.Data))
} else {
h.CodexModels(c)
var got codexModelsResponseForTest
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
want := tc.want
if tc.name == "routes" {
want = []string{"mapped-model", "route-only"}
}
require.ElementsMatch(t, want, codexModelSlugsForTest(got.Models))
}
require.Equal(t, http.StatusOK, rec.Code)
}
})
}
}
func TestDefaultModelIDsForCompositeIncludesAntigravityDefaults(t *testing.T) {
antigravityIDs := defaultModelIDsForPlatform(service.PlatformAntigravity)
require.NotEmpty(t, antigravityIDs)
@@ -1570,3 +1645,188 @@ func TestDefaultModelIDsForPlatform_CompositeFallbackExcludesTypeSafe(t *testing
require.NotContains(t, defaultModelIDsForPlatform(service.PlatformComposite), "jev-latest")
require.NotContains(t, defaultCodexModelIDsForPlatform(service.PlatformComposite), "jev-latest")
}
func TestGatewayModels_CompositeExactRoutesRespectTypeSafeIsolation(t *testing.T) {
gin.SetMode(gin.TestMode)
const groupID int64 = 7784
routes := &gatewayModelsRouteRepoStub{routes: []service.CompositeModelRoute{
{PublicModel: "systemone-any-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointAny, Enabled: true},
{PublicModel: "systemone-responses-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
{PublicModel: "systemone-legacy-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Enabled: true},
{PublicModel: "jev-latest", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Enabled: true},
{PublicModel: "responses-alias", TargetPlatform: service.PlatformOpenAI, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
{PublicModel: "image-alias", TargetPlatform: service.PlatformOpenAI, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointImages, Enabled: true},
}}
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{}, routes)
wantModels := make([]string, 0, len(routes.routes))
for _, route := range routes.routes {
wantModels = append(wantModels, route.PublicModel)
}
for _, allowlisted := range []bool{false, true} {
name := "unrestricted"
if allowlisted {
name = "allowlisted"
}
t.Run(name, func(t *testing.T) {
group := &service.Group{ID: groupID, Platform: service.PlatformComposite}
if allowlisted {
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{
"systemone-any-alias", "systemone-responses-alias", "systemone-legacy-alias", "jev-latest", "responses-alias", "image-alias",
}}
}
for _, codex := range []bool{false, true} {
rec := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(rec)
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
ctx.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group})
if codex {
h.CodexModels(ctx)
var manifest codexModelsResponseForTest
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &manifest))
require.Equal(t, []string{"responses-alias"}, codexModelSlugsForTest(manifest.Models))
} else {
h.Models(ctx)
var models gatewayModelsResponseForTest
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &models))
require.ElementsMatch(t, wantModels, modelIDsForTest(models.Data))
}
require.Equal(t, http.StatusOK, rec.Code)
}
})
}
}
func TestGatewayModels_CompositeCodexFiltersEffectiveResponsesRoutes(t *testing.T) {
gin.SetMode(gin.TestMode)
const groupID int64 = 7793
defaultModel := defaultModelIDsForPlatform(service.PlatformOpenAI)[0]
for _, scenario := range []struct {
name string
model string
accounts []service.Account
routes []service.CompositeModelRoute
}{
{
name: "colliding exact routes",
model: "shared-alias",
routes: []service.CompositeModelRoute{
{PublicModel: "shared-alias", TargetPlatform: service.PlatformOpenAI, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointAny, Priority: 1, Enabled: true},
{PublicModel: "shared-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Priority: 100, Enabled: true},
},
},
{
name: "account alias exact override",
model: "account-alias",
accounts: []service.Account{{ID: 1, Platform: service.PlatformOpenAI, Credentials: map[string]any{
"model_mapping": map[string]any{"account-alias": "gpt-5"},
}}},
routes: []service.CompositeModelRoute{
{PublicModel: "account-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
},
},
{
name: "account alias prefix override",
model: "account-alias",
accounts: []service.Account{{ID: 1, Platform: service.PlatformOpenAI, Credentials: map[string]any{
"model_mapping": map[string]any{"account-alias": "gpt-5"},
}}},
routes: []service.CompositeModelRoute{
{PublicModel: "account-", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchPrefix, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
},
},
{
name: "account default override",
model: defaultModel,
accounts: []service.Account{{ID: 1, Platform: service.PlatformOpenAI}},
routes: []service.CompositeModelRoute{
{PublicModel: defaultModel, TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
},
},
{
name: "fallback default override",
model: defaultModel,
routes: []service.CompositeModelRoute{
{PublicModel: defaultModel, TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchPrefix, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
},
},
} {
t.Run(scenario.name, func(t *testing.T) {
for _, allowlisted := range []bool{false, true} {
name := "unrestricted"
if allowlisted {
name = "allowlisted"
}
t.Run(name, func(t *testing.T) {
routes := append([]service.CompositeModelRoute(nil), scenario.routes...)
if scenario.name != "fallback default override" {
routes = append(routes, service.CompositeModelRoute{
PublicModel: "responses-control", TargetPlatform: service.PlatformOpenAI, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true,
})
}
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{
byGroup: map[int64][]service.Account{groupID: scenario.accounts},
}, &gatewayModelsRouteRepoStub{routes: routes})
group := &service.Group{ID: groupID, Platform: service.PlatformComposite}
if allowlisted {
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{scenario.model, "responses-control"}}
}
for _, codex := range []bool{false, true} {
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
ctx.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group})
if codex {
h.CodexModels(ctx)
var manifest codexModelsResponseForTest
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &manifest))
slugs := codexModelSlugsForTest(manifest.Models)
require.NotContains(t, slugs, scenario.model)
if scenario.name != "fallback default override" {
require.Contains(t, slugs, "responses-control")
} else if allowlisted {
require.Empty(t, slugs)
}
} else {
h.Models(ctx)
var models gatewayModelsResponseForTest
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &models))
require.Contains(t, modelIDsForTest(models.Data), scenario.model)
}
require.Equal(t, http.StatusOK, recorder.Code)
}
})
}
})
}
}
func TestGatewayModels_CompositeCodexRouteLookupErrorKeepsFallback(t *testing.T) {
gin.SetMode(gin.TestMode)
const groupID int64 = 7793
model := defaultCodexModelIDsForPlatform(service.PlatformComposite)[0]
for _, allowlisted := range []bool{false, true} {
name := "unrestricted"
if allowlisted {
name = "allowlisted"
}
t.Run(name, func(t *testing.T) {
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{}, &gatewayModelsRouteRepoStub{err: errors.New("unavailable")})
group := &service.Group{ID: groupID, Platform: service.PlatformComposite}
if allowlisted {
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{model}}
}
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
ctx.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group})
h.CodexModels(ctx)
require.Equal(t, http.StatusOK, recorder.Code)
var manifest codexModelsResponseForTest
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &manifest))
require.Contains(t, codexModelSlugsForTest(manifest.Models), model)
if allowlisted {
require.Len(t, manifest.Models, 1)
}
})
}
}
@@ -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"))
}
}
@@ -1001,3 +1001,32 @@ func TestCodexModelsPinnedAccountsETagMatchReturns304(t *testing.T) {
require.Equal(t, http.StatusNotModified, second.Code, second.Body.String())
require.Empty(t, second.Body.Bytes())
}
func TestCodexModelsDiscoveryProjectsWildcardMappingWithoutAllowlist(t *testing.T) {
for _, mode := range []string{"ordinary", "pinned", "fallback"} {
t.Run(mode, func(t *testing.T) {
account := newPinnedCodexAccount(1, service.StatusActive, true, false)
account.Credentials["model_mapping"] = map[string]any{"gpt-*": "gpt-5.4"}
upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
1: `{"models":[{"slug":"gpt-5.4"},{"slug":"blocked"},{"slug":"codex-auto-review"},{"slug":"gpt-image-2"}]}`,
}}
handler := newPinnedCodexTestHandler([]service.Account{account}, upstream, 3)
group := &service.Group{ID: 84, Platform: service.PlatformOpenAI,
ModelAllowlist: service.GroupModelAllowlist{Models: []string{"blocked"}},
}
if mode != "ordinary" {
group.CodexModelsManifestConfig = service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{1}}
}
if mode == "fallback" {
group.CodexModelsManifestConfig.AccountIDs = []int64{999}
group.CodexModelsManifestConfig.FallbackToScheduler = true
}
first := performCodexModelsRequestForGroup(t, handler, group, "")
require.Equal(t, http.StatusOK, first.Code, first.Body.String())
require.Equal(t, []string{"gpt-5.4"}, codexHandlerManifestSlugs(t, first))
require.NotEmpty(t, first.Header().Get("ETag"))
second := performCodexModelsRequestForGroup(t, handler, group, first.Header().Get("ETag"))
require.Equal(t, http.StatusNotModified, second.Code)
})
}
}
@@ -17,6 +17,7 @@ import (
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
@@ -330,9 +331,9 @@ func allowOpenAICompatibleMessagesDispatch(c *gin.Context, apiKey *service.APIKe
func openAICompatibleTextTargetAllowed(c *gin.Context, apiKey *service.APIKey, model string) bool {
return compositeTargetPlatformAllowed(c, apiKey, model,
service.PlatformOpenAI, service.PlatformGrok,
service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek,
service.PlatformMiniMax, service.PlatformOpenCodeGo)
domain.PlatformIDsWhere(func(spec domain.PlatformSpec) bool {
return spec.Gateway == domain.PlatformGatewayOpenAI
})...)
}
// isResponsesWebSocketCompositePlatform 限定 composite 分组在 Responses WebSocket
@@ -2125,6 +2126,101 @@ func (p *openAIWSTurnPricing) currentOr(fallback time.Time) time.Time {
return fallback
}
// openAIWSTurnAPIKeyLookup 是后续 turn 重取 API Key 认证快照的入口
// (APIKeyService.GetByKey:经 L1/L2 认证缓存,未命中才回源)。
type openAIWSTurnAPIKeyLookup interface {
GetByKey(ctx context.Context, key string) (*service.APIKey, error)
}
// openAIWSTurnBillingAPIKeys 按 turn 号保存每个 turn 计费用的 API Key 快照。
//
// API Key 认证快照只在建连时取一次。若连接内所有 turn 共用它,分组定价(倍率、
// 高峰、图片定价、利润门参数)会停留在建连时刻:管理员调价后,已打开的长连接
// 在客户端重连前一直按旧价计费、按旧售价过利润门,而 HTTP 请求每次都经认证缓存
// 取快照,缓存失效后即生效。
//
// 因此后续 turn 在 BeforeTurn(与 HTTP 请求进入认证中间件同位)经同一认证缓存
// 重取快照,只采用其中的分组,作为该 turn 利润门准入与用量计费共同的分组;用户、
// Key 限额、订阅仍沿用建连快照。只在同一把 Key、同一分组且平台与订阅类型未变时
// 采用:Key 换组或分组改平台/订阅类型不是调价,该连接按建连分组调度,继续按建连
// 快照计费;重取失败同样保留建连快照,不断连。
//
// 首轮沿用刚经认证中间件取得的建连快照;没有经过 BeforeTurn 的 turn 回退建连
// 快照。只保留当前与上一个 turn:下一 turn 的 BeforeTurn 先于上一 turn 的
// AfterTurn 执行时,上一 turn 仍取到自己的快照。
type openAIWSTurnBillingAPIKeys struct {
mu sync.Mutex
keys map[int]*service.APIKey
}
// begin 在 BeforeTurn 重装利润门前调用:后续 turn 重取计费分组并按 turn 记下,
// 返回换入该分组的上下文供本 turn 的利润门使用。
func (k *openAIWSTurnBillingAPIKeys) begin(ctx context.Context, apiKeyService *service.APIKeyService, turn int, conn *service.APIKey) context.Context {
turnKey := conn
if turn > 1 && apiKeyService != nil {
turnKey = refreshOpenAIWSTurnBillingAPIKey(ctx, apiKeyService, conn)
}
k.set(turn, turnKey)
return withOpenAIWSTurnBillingGroup(ctx, turnKey)
}
func (k *openAIWSTurnBillingAPIKeys) set(turn int, key *service.APIKey) {
k.mu.Lock()
defer k.mu.Unlock()
if k.keys == nil {
k.keys = make(map[int]*service.APIKey, 2)
}
for t := range k.keys {
if t < turn-1 {
delete(k.keys, t)
}
}
k.keys[turn] = key
}
func (k *openAIWSTurnBillingAPIKeys) forTurn(turn int, conn *service.APIKey) *service.APIKey {
k.mu.Lock()
defer k.mu.Unlock()
if key := k.keys[turn]; key != nil {
return key
}
return conn
}
// refreshOpenAIWSTurnBillingAPIKey 返回本 turn 计费用的 API Key:分组取当前认证
// 快照,其余字段与建连快照共享。不满足采用条件时原样返回建连快照。
func refreshOpenAIWSTurnBillingAPIKey(ctx context.Context, lookup openAIWSTurnAPIKeyLookup, conn *service.APIKey) *service.APIKey {
if lookup == nil || conn == nil || conn.Key == "" || conn.GroupID == nil || conn.Group == nil {
return conn
}
latest, err := lookup.GetByKey(ctx, conn.Key)
if err != nil || latest == nil || latest.ID != conn.ID || latest.GroupID == nil || *latest.GroupID != *conn.GroupID {
return conn
}
group := latest.Group
if group == nil || group.ID != conn.Group.ID ||
group.Platform != conn.Group.Platform ||
group.SubscriptionType != conn.Group.SubscriptionType {
return conn
}
turnKey := *conn
turnKey.Group = group
return &turnKey
}
// withOpenAIWSTurnBillingGroup 把本 turn 的计费分组换进认证分组上下文,使利润门
// 的售价与本 turn 计费同源。上下文里没有同 ID 的认证分组时不改动。
func withOpenAIWSTurnBillingGroup(ctx context.Context, turnKey *service.APIKey) context.Context {
if turnKey == nil || turnKey.Group == nil {
return ctx
}
current, ok := ctx.Value(ctxkey.Group).(*service.Group)
if !ok || current == nil || current == turnKey.Group || current.ID != turnKey.Group.ID {
return ctx
}
return context.WithValue(ctx, ctxkey.Group, turnKey.Group)
}
// recordOpenAIProfitVeto 记录 OpenAI 侧选号循环的一次利润门终检否决:把账号
// 加入本请求排除集并递增否决计数。返回 false 表示否决次数已达
// maxProfitVetoAttempts,调用方必须停止重选并按「无可用账号」终止。
@@ -2856,6 +2952,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
// turn 级定价:首轮回退到 TurnStarted 的所属 turn 时刻;后续 turn 由
// BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号。
var turnPricing openAIWSTurnPricing
// 后续 turn 的计费分组同样在 BeforeTurn 经认证缓存重取,使分组调价对
// 已打开的连接生效。
var turnBillingAPIKeys openAIWSTurnBillingAPIKeys
// Passthrough ingress does not invoke BeforeTurn for the first frame.
if err := checkSimpleModeTurnBilling(); err != nil {
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "billing check failed")
@@ -2943,7 +3042,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
// 长连接跨峰谷/倍率刷新防护:每个 turn 按当前时刻重装门并复核
// 当前账号,越线即要求客户端重连重选(连接绑定单一上游账号,
// 无法中途换号)。本 turn 的准入与计费共用同一 pricingAt。
turnCtx, turnAt := h.gatewayService.WithOpenAITurnPricingContext(ctx, apiKey.GroupID)
turnBillingCtx := turnBillingAPIKeys.begin(ctx, h.apiKeyService, turn, apiKey)
turnCtx, turnAt := h.gatewayService.WithOpenAITurnPricingContext(turnBillingCtx, apiKey.GroupID)
if _, vetoed, reason := h.gatewayService.ProfitControlVetoLatest(turnCtx, account); vetoed {
reqLog.Info("openai.websocket_turn_profit_vetoed",
zap.Int("turn", turn),
@@ -3059,11 +3159,12 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
turnRecordPricingAt := turnPricing.currentOr(turnStart)
turnBillingAPIKey := turnBillingAPIKeys.forTurn(turn, apiKey)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) {
if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{
Result: result,
APIKey: apiKey,
APIKey: turnBillingAPIKey,
User: apiKey.User,
Account: account,
Subscription: subscription,
@@ -15,6 +15,7 @@ import (
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
@@ -1948,6 +1949,12 @@ type openAIResponsesWSUsageLogCase struct {
afterFirstUpstreamRequest func(channelSvc *service.ChannelService) error
// group 覆盖 apiKey.Group(分组级模型白名单测试用);nil 保持原有无分组行为。
group *service.Group
// apiKeyService 非 nil 时模拟 API Key 认证中间件:连接认证快照经它按
// apiKeyCredential 取得,并把其分组放入请求 ctx;handler 也使用它。
apiKeyService *service.APIKeyService
apiKeyCredential string
// accountRateMultiplier 覆盖账号倍率(利润门测试用)。
accountRateMultiplier *float64
// firstFrameCloseExpected:首帧即被拒(连接被 1008 关闭),不期待任何响应帧。
firstFrameCloseExpected bool
// secondTurnCloseExpected:第二个 turn 被拒(连接被 1008 关闭)。
@@ -2962,6 +2969,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
if tc.accountPlatform != "" {
account.Platform = tc.accountPlatform
}
account.RateMultiplier = tc.accountRateMultiplier
if strings.TrimSpace(tc.ingressMode) != "" {
account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
}
@@ -3050,6 +3058,12 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
GroupID: &groupID,
User: &service.User{ID: 1701, Status: service.StatusActive},
}
if tc.apiKeyService != nil {
h.apiKeyService = tc.apiKeyService
authKey, err := tc.apiKeyService.GetByKey(context.Background(), tc.apiKeyCredential)
require.NoError(t, err)
apiKey = authKey
}
if tc.simpleModeRejectAtRead > 0 {
apiKey.RateLimit5h = 1
}
@@ -3060,6 +3074,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
router.Use(func(c *gin.Context) {
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
if tc.apiKeyService != nil && apiKey.Group != nil {
c.Request = c.Request.WithContext(context.WithValue(c.Request.Context(), ctxkey.Group, apiKey.Group))
}
c.Next()
})
router.GET("/openai/v1/responses", h.ResponsesWebSocket)
@@ -0,0 +1,293 @@
package handler
import (
"context"
"errors"
"sync"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
coderws "github.com/coder/websocket"
"github.com/stretchr/testify/require"
)
// 分组调价对已打开的 WebSocket 连接生效:同一条连接的两个 turn 之间由管理端
// 改分组倍率(并按分组失效认证缓存),第二个 turn 必须按新倍率计费、按新售价
// 过利润门;HTTP 请求每次经认证缓存取快照,本来就是这个行为。
const (
wsGroupPricingGroupID = int64(4201) // 与 runOpenAIResponsesWebSocketUsageLogCase 的调度分组一致
wsGroupPricingKey = "sk-ws-group-pricing"
wsGroupPricingFirst = `{"type":"response.create","model":"gpt-5.6-sol","stream":false}`
wsGroupPricingSecond = `{"type":"response.create","model":"gpt-5.6-sol","stream":false}`
)
// wsGroupPricingAPIKeyRepoStub 只实现认证查询与按分组列 Key(认证缓存失效用)。
type wsGroupPricingAPIKeyRepoStub struct {
service.APIKeyRepository
mu sync.Mutex
apiKey service.APIKey
group service.Group
err error
lookups int
}
func newWSGroupPricingAPIKeyRepoStub(rate float64) *wsGroupPricingAPIKeyRepoStub {
groupID := wsGroupPricingGroupID
return &wsGroupPricingAPIKeyRepoStub{
apiKey: service.APIKey{
ID: 1801,
UserID: 1701,
Key: wsGroupPricingKey,
Status: service.StatusActive,
GroupID: &groupID,
User: &service.User{ID: 1701, Status: service.StatusActive, Concurrency: 1},
},
group: service.Group{
ID: groupID,
Name: "ws-group-pricing",
Platform: service.PlatformOpenAI,
Status: service.StatusActive,
RateMultiplier: rate,
Hydrated: true,
},
}
}
func (s *wsGroupPricingAPIKeyRepoStub) GetByKeyForAuth(ctx context.Context, key string) (*service.APIKey, error) {
s.mu.Lock()
defer s.mu.Unlock()
s.lookups++
if s.err != nil {
return nil, s.err
}
if key != s.apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
apiKey := s.apiKey
user := *s.apiKey.User
apiKey.User = &user
group := s.group
apiKey.Group = &group
return &apiKey, nil
}
func (s *wsGroupPricingAPIKeyRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
return []string{s.apiKey.Key}, nil
}
// adminUpdate 模拟管理端改 Key/分组后按分组失效认证缓存(与 admin 分组更新同一入口)。
func (s *wsGroupPricingAPIKeyRepoStub) adminUpdate(apiKeyService *service.APIKeyService, fn func(apiKey *service.APIKey, group *service.Group)) error {
s.mu.Lock()
fn(&s.apiKey, &s.group)
s.mu.Unlock()
apiKeyService.InvalidateAuthCacheByGroupID(context.Background(), wsGroupPricingGroupID)
return nil
}
func (s *wsGroupPricingAPIKeyRepoStub) lookupCount() int {
s.mu.Lock()
defer s.mu.Unlock()
return s.lookups
}
func newWSGroupPricingAPIKeyService(repo *wsGroupPricingAPIKeyRepoStub) *service.APIKeyService {
return service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, &config.Config{})
}
func TestOpenAIResponsesWebSocket_GroupRateChangeAppliesToNextTurn(t *testing.T) {
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeCtxPool} {
t.Run(mode, func(t *testing.T) {
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
apiKeyService := newWSGroupPricingAPIKeyService(repo)
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: wsGroupPricingFirst,
secondPayload: wsGroupPricingSecond,
ingressMode: mode,
apiKeyService: apiKeyService,
apiKeyCredential: wsGroupPricingKey,
afterFirstUpstreamRequest: func(*service.ChannelService) error {
return repo.adminUpdate(apiKeyService, func(_ *service.APIKey, group *service.Group) { group.RateMultiplier = 0.3 })
},
})
require.Len(t, got.logs, 2)
require.InDelta(t, 3.0, got.logs[0].RateMultiplier, 1e-12, "turn 1 bills at the rate in force at connect")
require.InDelta(t, 0.3, got.logs[1].RateMultiplier, 1e-12, "turn 2 must bill at the group rate in force when it starts")
require.Greater(t, got.logs[0].ActualCost, 0.0)
require.InDelta(t, got.logs[0].ActualCost/10, got.logs[1].ActualCost, 1e-15)
})
}
}
func TestOpenAIResponsesWebSocket_GroupRateChangeReachesProfitGateOnNextTurn(t *testing.T) {
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
repo.group.ProfitControlEnabled = true
repo.group.ProfitMinMargin = 0.1
apiKeyService := newWSGroupPricingAPIKeyService(repo)
accountRate := 1.0
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: wsGroupPricingFirst,
secondPayload: wsGroupPricingSecond,
apiKeyService: apiKeyService,
apiKeyCredential: wsGroupPricingKey,
accountRateMultiplier: &accountRate,
// 3.0 → 0.3:账号倍率 1.0 在新售价下已越线,第二个 turn 必须在转发上游前被拒。
afterFirstUpstreamRequest: func(*service.ChannelService) error {
return repo.adminUpdate(apiKeyService, func(_ *service.APIKey, group *service.Group) { group.RateMultiplier = 0.3 })
},
secondTurnCloseExpected: true,
closeStatus: coderws.StatusTryAgainLater,
closeReason: "no longer eligible for this connection",
})
}
func TestOpenAIResponsesWebSocket_KeyMovedToAnotherGroupKeepsConnectionGroup(t *testing.T) {
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
apiKeyService := newWSGroupPricingAPIKeyService(repo)
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: wsGroupPricingFirst,
secondPayload: wsGroupPricingSecond,
apiKeyService: apiKeyService,
apiKeyCredential: wsGroupPricingKey,
afterFirstUpstreamRequest: func(*service.ChannelService) error {
return repo.adminUpdate(apiKeyService, func(apiKey *service.APIKey, group *service.Group) {
other := wsGroupPricingGroupID + 1
apiKey.GroupID = &other
group.ID = other
group.RateMultiplier = 0.3
})
},
})
require.Len(t, got.logs, 2)
require.InDelta(t, 3.0, got.logs[1].RateMultiplier, 1e-12,
"a key moved to another group keeps the connection's group: the connection was scheduled from it")
}
func TestOpenAIResponsesWebSocket_FailedKeyRefreshKeepsConnectionGroup(t *testing.T) {
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
apiKeyService := newWSGroupPricingAPIKeyService(repo)
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: wsGroupPricingFirst,
secondPayload: wsGroupPricingSecond,
apiKeyService: apiKeyService,
apiKeyCredential: wsGroupPricingKey,
afterFirstUpstreamRequest: func(*service.ChannelService) error {
return repo.adminUpdate(apiKeyService, func(_ *service.APIKey, group *service.Group) {
group.RateMultiplier = 0.3
repo.err = errors.New("auth store unavailable")
})
},
})
require.Len(t, got.logs, 2)
require.InDelta(t, 3.0, got.logs[1].RateMultiplier, 1e-12, "a failed refresh keeps the connection snapshot")
require.GreaterOrEqual(t, repo.lookupCount(), 2, "the second turn must have attempted a refresh")
}
type wsTurnAPIKeyLookupFunc func(ctx context.Context, key string) (*service.APIKey, error)
func (f wsTurnAPIKeyLookupFunc) GetByKey(ctx context.Context, key string) (*service.APIKey, error) {
return f(ctx, key)
}
func TestRefreshOpenAIWSTurnBillingAPIKey(t *testing.T) {
groupID := int64(77)
conn := &service.APIKey{
ID: 5,
Key: "sk-conn",
GroupID: &groupID,
User: &service.User{ID: 9},
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, SubscriptionType: service.SubscriptionTypeStandard, RateMultiplier: 3},
}
latestWith := func(mutate func(k *service.APIKey)) *service.APIKey {
gid := groupID
k := &service.APIKey{
ID: 5,
Key: "sk-conn",
GroupID: &gid,
User: &service.User{ID: 9, Balance: 123},
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, SubscriptionType: service.SubscriptionTypeStandard, RateMultiplier: 0.3},
}
if mutate != nil {
mutate(k)
}
return k
}
t.Run("same group adopts the latest group and nothing else", func(t *testing.T) {
got := refreshOpenAIWSTurnBillingAPIKey(context.Background(), wsTurnAPIKeyLookupFunc(func(ctx context.Context, key string) (*service.APIKey, error) {
require.Equal(t, "sk-conn", key)
return latestWith(nil), nil
}), conn)
require.NotSame(t, conn, got)
require.InDelta(t, 0.3, got.Group.RateMultiplier, 1e-12)
require.Same(t, conn.User, got.User, "only the group snapshot is refreshed")
require.Equal(t, conn.ID, got.ID)
require.InDelta(t, 3.0, conn.Group.RateMultiplier, 1e-12, "the connection snapshot is never mutated")
})
keep := map[string]wsTurnAPIKeyLookupFunc{
"lookup error": func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, errors.New("boom")
},
"not found": func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, service.ErrAPIKeyNotFound
},
"nil key": func(ctx context.Context, key string) (*service.APIKey, error) { return nil, nil },
"different key id": func(ctx context.Context, key string) (*service.APIKey, error) {
return latestWith(func(k *service.APIKey) { k.ID = 6 }), nil
},
"moved group": func(ctx context.Context, key string) (*service.APIKey, error) {
return latestWith(func(k *service.APIKey) { other := int64(78); k.GroupID = &other; k.Group.ID = other }), nil
},
"no group": func(ctx context.Context, key string) (*service.APIKey, error) {
return latestWith(func(k *service.APIKey) { k.Group = nil }), nil
},
"platform changed": func(ctx context.Context, key string) (*service.APIKey, error) {
return latestWith(func(k *service.APIKey) { k.Group.Platform = service.PlatformAnthropic }), nil
},
"billing type changed": func(ctx context.Context, key string) (*service.APIKey, error) {
return latestWith(func(k *service.APIKey) { k.Group.SubscriptionType = service.SubscriptionTypeSubscription }), nil
},
}
for name, lookup := range keep {
t.Run(name+" keeps the connection snapshot", func(t *testing.T) {
require.Same(t, conn, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, conn))
})
}
t.Run("no credential or no group skips the lookup", func(t *testing.T) {
called := false
lookup := wsTurnAPIKeyLookupFunc(func(ctx context.Context, key string) (*service.APIKey, error) {
called = true
return latestWith(nil), nil
})
noKey := *conn
noKey.Key = ""
require.Same(t, &noKey, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, &noKey))
noGroup := *conn
noGroup.Group = nil
require.Same(t, &noGroup, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, &noGroup))
require.Nil(t, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, nil))
require.Same(t, conn, refreshOpenAIWSTurnBillingAPIKey(context.Background(), nil, conn))
require.False(t, called)
})
}
func TestOpenAIWSTurnBillingAPIKeysKeepPreviousTurnUntilItIsRecorded(t *testing.T) {
conn := &service.APIKey{ID: 1}
turn2 := &service.APIKey{ID: 2}
turn3 := &service.APIKey{ID: 3}
var keys openAIWSTurnBillingAPIKeys
require.Same(t, conn, keys.forTurn(1, conn), "a turn without BeforeTurn bills with the connection snapshot")
keys.set(2, turn2)
keys.set(3, turn3) // turn 3 starts before turn 2's usage is submitted
require.Same(t, turn2, keys.forTurn(2, conn))
require.Same(t, turn3, keys.forTurn(3, conn))
keys.set(4, conn)
require.Same(t, conn, keys.forTurn(2, conn), "only the current and previous turn are kept")
require.Len(t, keys.keys, 2)
}
@@ -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)
}
+2 -2
View File
@@ -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,
+3 -3
View File
@@ -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
}
+2 -2
View File
@@ -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
+2
View File
@@ -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
}
@@ -264,10 +264,11 @@ func loginTestUser(t *testing.T) string {
t.Helper()
// 先尝试用管理员账户登录
adminEmail := getEnv("ADMIN_EMAIL", "admin@sub2api.local")
// 管理员邮箱不再有固定默认值(未设置时安装阶段随机生成),需同时提供邮箱与密码。
adminEmail := getEnv("ADMIN_EMAIL", "")
adminPassword := getEnv("ADMIN_PASSWORD", "")
if adminPassword == "" {
if adminEmail == "" || adminPassword == "" {
// 尝试用测试用户
adminEmail = testUserEmail
adminPassword = testUserPassword
@@ -49,21 +49,9 @@ const (
PlatformTypeSafe = domain.PlatformTypeSafe
)
// AllPlatforms 返回所有支持的平台列表
// AllPlatforms 返回所有支持的平台列表(平台清单,按展示顺序)。
func AllPlatforms() []string {
return []string{
PlatformAnthropic,
PlatformOpenAI,
PlatformGemini,
PlatformAntigravity,
PlatformGrok,
PlatformKimi,
PlatformZhipu,
PlatformDeepseek,
PlatformMiniMax,
PlatformOpenCodeGo,
PlatformTypeSafe,
}
return domain.ConcretePlatformIDs()
}
// Validate 验证规则配置的有效性
@@ -19,5 +19,7 @@ func TestAllPlatformsIncludesEveryConcretePlatform(t *testing.T) {
"minimax",
"opencode_go",
"typesafe",
"command_code",
"cline",
}, AllPlatforms())
}
@@ -376,6 +376,9 @@ func (e *EasyPay) VerifyNotification(_ context.Context, rawBody string, _ map[st
// url.ParseQuery already decodes values — no additional decode needed.
params := make(map[string]string)
for k := range values {
if !easyPayNotifyAllowedParams[k] {
return nil, fmt.Errorf("unexpected notify param: %s", k)
}
params[k] = values.Get(k)
}
sign := params["sign"]
@@ -601,3 +604,26 @@ func easyPaySign(params map[string]string, pkey string) string {
func easyPayVerifySign(params map[string]string, pkey string, sign string) bool {
return hmac.Equal([]byte(easyPaySign(params, pkey)), []byte(sign))
}
// easyPayNotifyAllowedParams is the exact parameter set a genuine EasyPay
// (彩虹易支付-compatible) async notification carries. Order-creation-only
// fields (notify_url, return_url, cid, device, clientip, ...) must never
// appear in a callback: because the sign base string concatenates values
// unescaped, a signed order URL whose return_url embeds e.g.
// "trade_status=TRADE_SUCCESS" could otherwise be replayed as a forged
// payment-success notification (issue #7881). Rejecting unknown keys closes
// the whole smuggling class; genuinely paid orders rejected by an exotic
// upstream variant are still recovered by the upstream QueryOrder reconcile
// path.
var easyPayNotifyAllowedParams = map[string]bool{
"pid": true,
"trade_no": true,
"out_trade_no": true,
"type": true,
"name": true,
"money": true,
"trade_status": true,
"param": true,
"sign": true,
"sign_type": true,
}
@@ -0,0 +1,172 @@
package provider
// Regression tests for the EasyPay forged-callback vulnerability
// (Wei-Shaw/sub2api issue #7875).
//
// Attack recap: the order-creation signature (visible to the payer in the
// submit.php popup URL) signs return_url among other fields, and the sign
// base string concatenates values unescaped. A client-supplied return_url
// ending in "&trade_status=TRADE_SUCCESS" therefore produced a signature
// that was byte-identical to a payment-success notification's signature
// once the callback decoded return_url partially and promoted the smuggled
// pair to a top-level parameter. Merchant-key rotation could not help
// because the attacker never needed the key.
//
// Fixes covered here:
// 1. service.CanonicalizeReturnURL drops client-supplied query parameters
// (see payment_resume_service_test.go).
// 2. VerifyNotification rejects any parameter outside the genuine async
// notify set, so smuggled order-creation fields (return_url et al.)
// fail closed even if a signed value smuggles a fake pair.
import (
"context"
"net/url"
"testing"
"github.com/Wei-Shaw/sub2api/internal/payment"
)
// easyPayPoCProvider returns a provider with a fixed test credential set.
func easyPayPoCProvider() *EasyPay {
return &EasyPay{config: map[string]string{
"pid": "1000",
"pkey": "MERCHANT_SECRET_KEY",
"apiBase": "https://pay.example.com",
"notifyUrl": "https://site.example.com/api/v1/payment/webhook/easypay",
"returnUrl": "https://site.example.com/payment/result",
}}
}
// TestEasyPayNotifyRejectsForgedSignReuseCallback is the exact PoC payload
// that was accepted before the fix: the order's own submit.php signature
// replayed with trade_status smuggled out of the return_url value.
func TestEasyPayNotifyRejectsForgedSignReuseCallback(t *testing.T) {
t.Parallel()
e := easyPayPoCProvider()
// Order creation (popup mode): return_url carries the smuggled pair as
// its last sorted inner query key, exactly as buildPaymentReturnURL
// would have produced before CanonicalizeReturnURL stripped user query.
returnURL := "https://site.example.com/payment/result?order_id=99&out_trade_no=ORDER123&status=success&trade_status=TRADE_SUCCESS"
createParams := map[string]string{
"pid": "1000",
"type": "alipay",
"out_trade_no": "ORDER123",
"notify_url": e.config["notifyUrl"],
"return_url": returnURL,
"name": "balance recharge",
"money": "650.00",
}
sign := easyPaySign(createParams, e.config["pkey"])
// Forged callback: return_url encoded only up to status=success, then a
// raw &trade_status=TRADE_SUCCESS promotes it to a top-level parameter.
prefix := "https://site.example.com/payment/result?order_id=99&out_trade_no=ORDER123&status=success"
cb := url.Values{}
cb.Set("pid", "1000")
cb.Set("type", "alipay")
cb.Set("out_trade_no", "ORDER123")
cb.Set("notify_url", e.config["notifyUrl"])
cb.Set("name", "balance recharge")
cb.Set("money", "650.00")
cb.Set("return_url", prefix)
rawCallback := cb.Encode() + "&trade_status=TRADE_SUCCESS" + "&sign=" + sign + "&sign_type=MD5"
if _, err := e.VerifyNotification(context.Background(), rawCallback, nil); err == nil {
t.Fatal("forged sign-reuse callback must be rejected")
}
}
// TestEasyPayNotifyRejectsOrderURLReplay covers the milder variant where the
// attacker replays the complete signed pay URL unchanged: return_url itself
// is not a legitimate notify parameter and must also be rejected.
func TestEasyPayNotifyRejectsOrderURLReplay(t *testing.T) {
t.Parallel()
e := easyPayPoCProvider()
createParams := map[string]string{
"pid": "1000",
"type": "alipay",
"out_trade_no": "ORDER123",
"notify_url": e.config["notifyUrl"],
"return_url": "https://site.example.com/payment/result",
"name": "balance recharge",
"money": "650.00",
}
createParams["sign"] = easyPaySign(createParams, e.config["pkey"])
createParams["sign_type"] = signTypeMD5
q := url.Values{}
for k, v := range createParams {
q.Set(k, v)
}
if _, err := e.VerifyNotification(context.Background(), q.Encode(), nil); err == nil {
t.Fatal("replayed order URL must be rejected")
}
}
// TestEasyPayNotifyAcceptsGenuineCallback locks in the legitimate notify
// contract: the canonical parameter set with a valid signature succeeds.
func TestEasyPayNotifyAcceptsGenuineCallback(t *testing.T) {
t.Parallel()
e := easyPayPoCProvider()
params := map[string]string{
"pid": "1000",
"trade_no": "2026100622001400000001",
"out_trade_no": "ORDER123",
"type": "alipay",
"name": "balance recharge",
"money": "650.00",
"trade_status": tradeStatusSuccess,
}
sign := easyPaySign(params, e.config["pkey"])
params["sign"] = sign
params["sign_type"] = signTypeMD5
q := url.Values{}
for k, v := range params {
q.Set(k, v)
}
n, err := e.VerifyNotification(context.Background(), q.Encode(), nil)
if err != nil {
t.Fatalf("genuine callback rejected: %v", err)
}
if n.Status != payment.ProviderStatusSuccess {
t.Fatalf("status = %v, want success", n.Status)
}
if n.OrderID != "ORDER123" || n.TradeNo != "2026100622001400000001" || n.Amount != 650.00 {
t.Fatalf("unexpected notification: %+v", n)
}
}
// TestEasyPayNotifyRejectsUnknownParam ensures any parameter outside the
// canonical notify set fails closed, including empty-valued ones that the
// signer itself would skip.
func TestEasyPayNotifyRejectsUnknownParam(t *testing.T) {
t.Parallel()
e := easyPayPoCProvider()
params := map[string]string{
"pid": "1000",
"trade_no": "T1",
"out_trade_no": "ORDER123",
"type": "alipay",
"name": "balance recharge",
"money": "650.00",
"trade_status": tradeStatusSuccess,
}
sign := easyPaySign(params, e.config["pkey"])
q := url.Values{}
for k, v := range params {
q.Set(k, v)
}
q.Set("sign", sign)
q.Set("sign_type", signTypeMD5)
q.Set("device", "") // empty value: invisible to the signer, still rejected
if _, err := e.VerifyNotification(context.Background(), q.Encode(), nil); err == nil {
t.Fatal("callback with unknown param must be rejected")
}
}
@@ -122,6 +122,20 @@ type ClaudeContentItem struct {
Input any `json:"input,omitempty"`
}
// MarshalJSON 保证 thinking 块始终带 thinking/signature 字段(即使为空),
// 与 Anthropic 官方一致;Grok Build 等严格客户端缺少 signature 会解析失败。
func (c ClaudeContentItem) MarshalJSON() ([]byte, error) {
type claudeContentItem ClaudeContentItem
if c.Type != "thinking" {
return json.Marshal(claudeContentItem(c))
}
return json.Marshal(struct {
Thinking string `json:"thinking"`
Signature string `json:"signature"`
claudeContentItem
}{Thinking: c.Thinking, Signature: c.Signature, claudeContentItem: claudeContentItem(c)})
}
// ClaudeUsage Claude 用量统计
type ClaudeUsage struct {
InputTokens int `json:"input_tokens"`
@@ -1,6 +1,9 @@
package antigravity
import "testing"
import (
"encoding/json"
"testing"
)
func TestDefaultModels_ContainsNewAndLegacyImageModels(t *testing.T) {
t.Parallel()
@@ -44,3 +47,23 @@ func TestDefaultModels_ContainsNewAndLegacyImageModels(t *testing.T) {
}
}
}
func TestClaudeContentItem_ThinkingAlwaysHasSignature(t *testing.T) {
t.Parallel()
data, err := json.Marshal(ClaudeContentItem{Type: "thinking", Thinking: "plan"})
if err != nil {
t.Fatal(err)
}
if got, want := string(data), `{"thinking":"plan","signature":"","type":"thinking"}`; got != want {
t.Fatalf("thinking block = %s, want %s", got, want)
}
data, err = json.Marshal(ClaudeContentItem{Type: "text", Text: "hi"})
if err != nil {
t.Fatal(err)
}
if got, want := string(data), `{"type":"text","text":"hi"}`; got != want {
t.Fatalf("text block = %s, want %s", got, want)
}
}
@@ -296,8 +296,9 @@ func (p *StreamingProcessor) processThinking(text, signature string) []byte {
// 开始或继续 thinking 块
if p.blockType != BlockTypeThinking {
_, _ = result.Write(p.startBlock(BlockTypeThinking, map[string]any{
"type": "thinking",
"thinking": "",
"type": "thinking",
"thinking": "",
"signature": "",
}))
}
@@ -475,8 +476,9 @@ func (p *StreamingProcessor) emitEmptyThinkingWithSignature(signature string) []
var result bytes.Buffer
_, _ = result.Write(p.startBlock(BlockTypeThinking, map[string]any{
"type": "thinking",
"thinking": "",
"type": "thinking",
"thinking": "",
"signature": "",
}))
_, _ = result.Write(p.emitDelta("thinking_delta", map[string]any{
"thinking": "",
@@ -844,7 +844,7 @@ func TestStreamingReasoning(t *testing.T) {
sse, err := ResponsesAnthropicEventToSSE(events[0])
require.NoError(t, err)
assert.Contains(t, sse, `"content_block":{"thinking":"","type":"thinking"}`)
assert.Contains(t, sse, `"content_block":{"thinking":"","signature":"","type":"thinking"}`)
// reasoning text delta
events = ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
@@ -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
}
@@ -88,10 +88,11 @@ func ChatCompletionsToResponses(req *ChatCompletionsRequest) (*ResponsesRequest,
out.Tools = convertChatToolsToResponses(req.Tools, req.Functions)
}
// tool_choice: already compatible format — pass through directly.
// tool_choice: strings and Responses-shaped objects pass through; a named
// Chat choice nests the name under "function" and must be flattened.
// Legacy function_call needs mapping.
if len(req.ToolChoice) > 0 {
out.ToolChoice = req.ToolChoice
out.ToolChoice = convertChatToolChoiceToResponses(req.ToolChoice)
} else if len(req.FunctionCall) > 0 {
tc, err := convertChatFunctionCallToToolChoice(req.FunctionCall)
if err != nil {
@@ -107,7 +108,18 @@ func ChatCompletionsToResponses(req *ChatCompletionsRequest) (*ResponsesRequest,
// array into a Responses API input items array.
func convertChatMessagesToResponsesInput(msgs []ChatMessage) ([]ResponsesInputItem, error) {
var out []ResponsesInputItem
legacyIDs := legacyFunctionCallIDs{msgs: msgs}
for _, m := range msgs {
switch {
case m.Role == "assistant" && m.FunctionCall != nil && len(m.ToolCalls) == 0:
m.ToolCalls = []ChatToolCall{{
ID: legacyIDs.assign(m.FunctionCall.Name),
Type: "function",
Function: *m.FunctionCall,
}}
case m.Role == "function" && m.ToolCallID == "":
m.ToolCallID = legacyIDs.claim(m.Name)
}
items, err := chatMessageToResponsesItems(m)
if err != nil {
return nil, err
@@ -117,11 +129,63 @@ func convertChatMessagesToResponsesInput(msgs []ChatMessage) ([]ResponsesInputIt
return out, nil
}
// legacyFunctionCallIDs pairs legacy function calls with their results. Legacy
// function calling carries no call IDs: the assistant turn has function_call
// and the role=function result refers to it only by name, while the Responses
// API pairs function_call and function_call_output by call_id. IDs follow
// message order, so a replayed history converts to the same input, and skip
// IDs already used by tool_calls in the same conversation.
type legacyFunctionCallIDs struct {
msgs []ChatMessage
used map[string]bool
next int
pending map[string][]string
}
// assign returns a new call ID for a legacy function call and queues it for
// the next result with the same function name.
func (ids *legacyFunctionCallIDs) assign(name string) string {
if ids.used == nil {
ids.used = make(map[string]bool)
ids.pending = make(map[string][]string)
for _, m := range ids.msgs {
for _, tc := range m.ToolCalls {
ids.used[tc.ID] = true
}
if m.ToolCallID != "" {
ids.used[m.ToolCallID] = true
}
}
}
var id string
for {
ids.next++
id = fmt.Sprintf("call_legacy_%d", ids.next)
if !ids.used[id] {
break
}
}
ids.used[id] = true
ids.pending[name] = append(ids.pending[name], id)
return id
}
// claim returns the oldest unanswered call ID for name, or "" when no legacy
// call with that name precedes the result.
func (ids *legacyFunctionCallIDs) claim(name string) string {
queue := ids.pending[name]
if len(queue) == 0 {
return ""
}
ids.pending[name] = queue[1:]
return queue[0]
}
// chatMessageToResponsesItems converts a single ChatMessage into one or more
// ResponsesInputItem values.
func chatMessageToResponsesItems(m ChatMessage) ([]ResponsesInputItem, error) {
switch m.Role {
case "system":
case "system", "developer":
return chatSystemToResponses(m)
case "user":
return chatUserToResponses(m)
@@ -136,7 +200,9 @@ func chatMessageToResponsesItems(m ChatMessage) ([]ResponsesInputItem, error) {
}
}
// chatSystemToResponses converts a system message.
// chatSystemToResponses converts a system or developer message, keeping its
// role: the Responses API accepts both, and developer must not be demoted to
// a user turn.
func chatSystemToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
parsed, err := parseChatMessageContent(m.Content)
if err != nil {
@@ -146,7 +212,7 @@ func chatSystemToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
if err != nil {
return nil, err
}
return []ResponsesInputItem{{Type: "message", Role: "system", Content: content}}, nil
return []ResponsesInputItem{{Type: "message", Role: m.Role, Content: content}}, nil
}
// chatUserToResponses converts a user message, handling both plain strings and
@@ -303,8 +369,9 @@ func chatToolToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
}
// chatFunctionToResponses converts a legacy function result message
// (role=function) into a function_call_output item. The Name field is used as
// call_id since legacy function calls do not carry a separate call_id.
// (role=function) into a function_call_output item. ToolCallID holds the ID
// assigned to the matching legacy call (see legacyFunctionCallIDs); without a
// matching call the Name field is used as call_id.
func chatFunctionToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
output, err := parseChatContent(m.Content)
if err != nil {
@@ -313,9 +380,13 @@ func chatFunctionToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
if output == "" {
output = "(empty)"
}
callID := m.ToolCallID
if callID == "" {
callID = m.Name
}
return []ResponsesInputItem{{
Type: "function_call_output",
CallID: m.Name,
CallID: callID,
Output: output,
}}, nil
}
@@ -490,6 +561,34 @@ func defaultStrictFalse(src *bool) *bool {
return src
}
// convertChatToolChoiceToResponses maps a Chat Completions tool_choice to the
// Responses API shape.
//
// {"type":"function","function":{"name":"X"}} → {"type":"function","name":"X"}
//
// Strings ("auto", "none", "required") and objects that already use the
// Responses shape are returned unchanged.
func convertChatToolChoiceToResponses(raw json.RawMessage) json.RawMessage {
var choice struct {
Type string `json:"type"`
Function *struct {
Name string `json:"name"`
} `json:"function"`
}
if err := json.Unmarshal(raw, &choice); err != nil || choice.Type != "function" || choice.Function == nil {
return raw
}
name := strings.TrimSpace(choice.Function.Name)
if name == "" {
return raw
}
flat, err := json.Marshal(map[string]string{"type": "function", "name": name})
if err != nil {
return raw
}
return flat
}
// convertChatFunctionCallToToolChoice maps the legacy function_call field to a
// Responses API tool_choice value.
//
@@ -0,0 +1,43 @@
package apicompat
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
// The Responses API accepts role "developer" in input messages. Converting it
// to "user" silently demotes the caller's instructions to user turns.
func TestChatCompletionsToResponses_DeveloperRolePreserved(t *testing.T) {
tests := []struct {
name string
content string
wantContent string
}{
{name: "string content", content: `"Answer in French."`, wantContent: `"Answer in French."`},
{name: "text parts", content: `[{"type":"text","text":"Answer in French."}]`, wantContent: `[{"type":"input_text","text":"Answer in French."}]`},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := &ChatCompletionsRequest{
Model: "gpt-4o",
Messages: []ChatMessage{
{Role: "developer", Content: json.RawMessage(tt.content)},
{Role: "user", Content: json.RawMessage(`"Hello"`)},
},
}
resp, err := ChatCompletionsToResponses(req)
require.NoError(t, err)
var items []ResponsesInputItem
require.NoError(t, json.Unmarshal(resp.Input, &items))
require.Len(t, items, 2)
require.Equal(t, "message", items[0].Type)
require.Equal(t, "developer", items[0].Role)
require.JSONEq(t, tt.wantContent, string(items[0].Content))
require.Equal(t, "user", items[1].Role)
})
}
}
@@ -0,0 +1,87 @@
package apicompat
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
func convertLegacyFunctionHistory(t *testing.T, msgs []ChatMessage) []ResponsesInputItem {
t.Helper()
resp, err := ChatCompletionsToResponses(&ChatCompletionsRequest{Model: "gpt-4o", Messages: msgs})
require.NoError(t, err)
var items []ResponsesInputItem
require.NoError(t, json.Unmarshal(resp.Input, &items))
return items
}
// Legacy function calling has no call IDs: the assistant turn carries
// function_call and the result is a role=function message keyed by name. The
// Responses API needs a function_call item and a function_call_output with the
// same call_id, otherwise the output is rejected as having no matching call.
func TestChatCompletionsToResponses_LegacyFunctionCallHistoryPaired(t *testing.T) {
msgs := []ChatMessage{
{Role: "user", Content: json.RawMessage(`"Weather in Paris and Rome?"`)},
{Role: "assistant", Content: json.RawMessage(`"Checking Paris."`), FunctionCall: &ChatFunctionCall{Name: "get_weather", Arguments: `{"city":"Paris"}`}},
{Role: "function", Name: "get_weather", Content: json.RawMessage(`"sunny"`)},
{Role: "assistant", FunctionCall: &ChatFunctionCall{Name: "get_weather", Arguments: `{"city":"Rome"}`}},
{Role: "function", Name: "get_weather", Content: json.RawMessage(`"rainy"`)},
}
items := convertLegacyFunctionHistory(t, msgs)
require.Len(t, items, 6)
require.Equal(t, "user", items[0].Role)
require.Equal(t, "assistant", items[1].Role)
firstCall, firstOutput := items[2], items[3]
require.Equal(t, "function_call", firstCall.Type)
require.Equal(t, "get_weather", firstCall.Name)
require.Equal(t, `{"city":"Paris"}`, firstCall.Arguments)
require.Equal(t, "function_call_output", firstOutput.Type)
require.Equal(t, "sunny", firstOutput.Output)
require.NotEmpty(t, firstCall.CallID)
require.Equal(t, firstCall.CallID, firstOutput.CallID)
secondCall, secondOutput := items[4], items[5]
require.Equal(t, "function_call", secondCall.Type)
require.Equal(t, `{"city":"Rome"}`, secondCall.Arguments)
require.Equal(t, "function_call_output", secondOutput.Type)
require.Equal(t, "rainy", secondOutput.Output)
require.Equal(t, secondCall.CallID, secondOutput.CallID)
require.NotEqual(t, firstCall.CallID, secondCall.CallID)
// Synthesized IDs must be stable across turns so the replayed history keeps
// the same prompt prefix.
require.Equal(t, items, convertLegacyFunctionHistory(t, msgs))
}
func TestChatCompletionsToResponses_LegacyFunctionCallIDAvoidsToolCallIDs(t *testing.T) {
items := convertLegacyFunctionHistory(t, []ChatMessage{
{Role: "user", Content: json.RawMessage(`"hi"`)},
{Role: "assistant", FunctionCall: &ChatFunctionCall{Name: "lookup", Arguments: `{}`}},
{Role: "function", Name: "lookup", Content: json.RawMessage(`"legacy result"`)},
{Role: "assistant", ToolCalls: []ChatToolCall{{ID: "call_legacy_1", Type: "function", Function: ChatFunctionCall{Name: "lookup", Arguments: `{}`}}}},
{Role: "tool", ToolCallID: "call_legacy_1", Content: json.RawMessage(`"tool result"`)},
})
require.Len(t, items, 5)
legacyCall, legacyOutput := items[1], items[2]
require.Equal(t, "function_call", legacyCall.Type)
require.Equal(t, legacyCall.CallID, legacyOutput.CallID)
require.NotEqual(t, "call_legacy_1", legacyCall.CallID)
require.Equal(t, "call_legacy_1", items[3].CallID)
require.Equal(t, "call_legacy_1", items[4].CallID)
}
func TestChatCompletionsToResponses_LegacyFunctionResultWithoutCallKeepsName(t *testing.T) {
items := convertLegacyFunctionHistory(t, []ChatMessage{
{Role: "user", Content: json.RawMessage(`"hi"`)},
{Role: "function", Name: "lookup", Content: json.RawMessage(`"orphan"`)},
})
require.Len(t, items, 2)
require.Equal(t, "function_call_output", items[1].Type)
require.Equal(t, "lookup", items[1].CallID)
}
@@ -0,0 +1,43 @@
package apicompat
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
// Chat Completions nests a named choice under "function"; the Responses API
// expects the name at the top level and rejects the nested shape.
func TestChatCompletionsToResponses_ToolChoiceShape(t *testing.T) {
tests := []struct {
name string
choice string
want string
}{
{name: "named chat function is flattened", choice: `{"type":"function","function":{"name":"get_weather"}}`, want: `{"type":"function","name":"get_weather"}`},
{name: "flat responses function is kept", choice: `{"type":"function","name":"get_weather"}`, want: `{"type":"function","name":"get_weather"}`},
{name: "auto string is kept", choice: `"auto"`, want: `"auto"`},
{name: "required string is kept", choice: `"required"`, want: `"required"`},
{name: "none string is kept", choice: `"none"`, want: `"none"`},
{name: "allowed_tools policy is kept", choice: `{"type":"allowed_tools","mode":"auto","tools":[{"type":"function","name":"get_weather"}]}`, want: `{"type":"allowed_tools","mode":"auto","tools":[{"type":"function","name":"get_weather"}]}`},
{name: "function without name is kept", choice: `{"type":"function","function":{}}`, want: `{"type":"function","function":{}}`},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := &ChatCompletionsRequest{
Model: "gpt-4o",
Messages: []ChatMessage{{Role: "user", Content: json.RawMessage(`"weather?"`)}},
Tools: []ChatTool{{
Type: "function",
Function: &ChatFunction{Name: "get_weather", Parameters: json.RawMessage(`{"type":"object"}`)},
}},
ToolChoice: json.RawMessage(tt.choice),
}
resp, err := ChatCompletionsToResponses(req)
require.NoError(t, err)
require.JSONEq(t, tt.want, string(resp.ToolChoice))
})
}
}
@@ -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:
@@ -15,7 +15,8 @@ import (
// ResponsesToChatCompletions converts a Responses API response into a Chat
// Completions response. Text output items are concatenated into
// choices[0].message.content; function_call items become tool_calls.
// choices[0].message.content, refusal parts into choices[0].message.refusal;
// function_call items become tool_calls.
func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatCompletionsResponse {
id := resp.ID
if id == "" {
@@ -31,6 +32,7 @@ func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatComp
}
var contentText string
var refusalText string
var reasoningText string
var toolCalls []ChatToolCall
@@ -38,8 +40,11 @@ func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatComp
switch item.Type {
case "message":
for _, part := range item.Content {
if part.Type == "output_text" && part.Text != "" {
switch {
case part.Type == "output_text" && part.Text != "":
contentText += part.Text
case part.Type == "refusal" && part.Refusal != "":
refusalText += part.Refusal
}
}
case "function_call":
@@ -62,7 +67,7 @@ func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatComp
}
}
msg := ChatMessage{Role: "assistant"}
msg := ChatMessage{Role: "assistant", Refusal: refusalText}
if len(toolCalls) > 0 {
msg.ToolCalls = toolCalls
}
@@ -149,6 +154,8 @@ func ResponsesEventToChatChunks(evt *ResponsesStreamEvent, state *ResponsesEvent
return resToChatHandleCreated(evt, state)
case "response.output_text.delta":
return resToChatHandleTextDelta(evt, state)
case "response.refusal.delta":
return resToChatHandleRefusalDelta(evt, state)
case "response.output_item.added":
return resToChatHandleOutputItemAdded(evt, state)
case "response.function_call_arguments.delta",
@@ -248,6 +255,14 @@ func resToChatHandleTextDelta(evt *ResponsesStreamEvent, state *ResponsesEventTo
return []ChatCompletionsChunk{makeChatDeltaChunk(state, ChatDelta{Content: &content})}
}
func resToChatHandleRefusalDelta(evt *ResponsesStreamEvent, state *ResponsesEventToChatState) []ChatCompletionsChunk {
if evt.Delta == "" {
return nil
}
refusal := evt.Delta
return []ChatCompletionsChunk{makeChatDeltaChunk(state, ChatDelta{Refusal: &refusal})}
}
func resToChatHandleOutputItemAdded(evt *ResponsesStreamEvent, state *ResponsesEventToChatState) []ChatCompletionsChunk {
// function_call 与 custom_tool_call(custom/freeform 工具)均按工具调用注册,
// 以便后续 *_input.delta / *_arguments.delta 能映射到正确的工具索引。
@@ -496,6 +511,7 @@ type bufferedFuncCall struct {
// (response.completed / response.done) carries an empty output array.
type BufferedResponseAccumulator struct {
text strings.Builder
refusal strings.Builder
reasoning strings.Builder
funcCalls []bufferedFuncCall
outputIndexToFuncIdx map[int]int
@@ -517,6 +533,10 @@ func (a *BufferedResponseAccumulator) ProcessEvent(event *ResponsesStreamEvent)
if event.Delta != "" {
_, _ = a.text.WriteString(event.Delta)
}
case "response.refusal.delta":
if event.Delta != "" {
_, _ = a.refusal.WriteString(event.Delta)
}
case "response.output_item.added":
if event.Item != nil && (event.Item.Type == "function_call" || event.Item.Type == "custom_tool_call") {
idx := len(a.funcCalls)
@@ -553,7 +573,7 @@ func (a *BufferedResponseAccumulator) ProcessEvent(event *ResponsesStreamEvent)
// HasContent reports whether any content has been accumulated.
func (a *BufferedResponseAccumulator) HasContent() bool {
return a.text.Len() > 0 || len(a.funcCalls) > 0 || a.reasoning.Len() > 0
return a.text.Len() > 0 || a.refusal.Len() > 0 || len(a.funcCalls) > 0 || a.reasoning.Len() > 0
}
// BuildOutput constructs a []ResponsesOutput from the accumulated delta
@@ -572,14 +592,18 @@ func (a *BufferedResponseAccumulator) BuildOutput() []ResponsesOutput {
})
}
if a.text.Len() > 0 {
if a.text.Len() > 0 || a.refusal.Len() > 0 {
var content []ResponsesContentPart
if a.text.Len() > 0 {
content = append(content, ResponsesContentPart{Type: "output_text", Text: a.text.String()})
}
if a.refusal.Len() > 0 {
content = append(content, ResponsesContentPart{Type: "refusal", Refusal: a.refusal.String()})
}
out = append(out, ResponsesOutput{
Type: "message",
Role: "assistant",
Content: []ResponsesContentPart{{
Type: "output_text",
Text: a.text.String(),
}},
Type: "message",
Role: "assistant",
Content: content,
})
}
@@ -0,0 +1,114 @@
package apicompat
import (
"encoding/json"
"strings"
"testing"
"github.com/stretchr/testify/require"
)
const refusalText = "I can't help with that."
func decodeResponsesEvent(t *testing.T, raw string) ResponsesStreamEvent {
t.Helper()
var evt ResponsesStreamEvent
require.NoError(t, json.Unmarshal([]byte(raw), &evt))
return evt
}
func chatMessageJSON(t *testing.T, resp *ChatCompletionsResponse) map[string]any {
t.Helper()
raw, err := json.Marshal(resp)
require.NoError(t, err)
var decoded struct {
Choices []struct {
Message map[string]any `json:"message"`
} `json:"choices"`
}
require.NoError(t, json.Unmarshal(raw, &decoded))
require.Len(t, decoded.Choices, 1)
return decoded.Choices[0].Message
}
// A refusal arrives as a "refusal" content part, not output_text. Dropping it
// leaves the Chat client with an empty assistant message.
func TestResponsesToChatCompletions_RefusalPreserved(t *testing.T) {
var resp ResponsesResponse
require.NoError(t, json.Unmarshal([]byte(`{
"id":"resp_refusal",
"status":"completed",
"output":[{"type":"message","role":"assistant","content":[{"type":"refusal","refusal":"`+refusalText+`"}]}]
}`), &resp))
chat := ResponsesToChatCompletions(&resp, "gpt-4o")
msg := chatMessageJSON(t, chat)
require.Equal(t, refusalText, msg["refusal"])
require.NotContains(t, msg, "content")
require.Equal(t, "stop", chat.Choices[0].FinishReason)
}
func TestResponsesEventToChatChunks_RefusalDelta(t *testing.T) {
state := NewResponsesEventToChatState()
events := []string{
`{"type":"response.created","response":{"id":"resp_refusal","model":"gpt-4o"}}`,
`{"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant"}}`,
`{"type":"response.refusal.delta","output_index":0,"content_index":0,"delta":"I can't "}`,
`{"type":"response.refusal.delta","output_index":0,"content_index":0,"delta":"help with that."}`,
`{"type":"response.refusal.done","output_index":0,"content_index":0,"refusal":"` + refusalText + `"}`,
`{"type":"response.completed","response":{"id":"resp_refusal","status":"completed"}}`,
}
var refusal strings.Builder
var finishReason string
for _, raw := range events {
evt := decodeResponsesEvent(t, raw)
for _, chunk := range ResponsesEventToChatChunks(&evt, state) {
encoded, err := json.Marshal(chunk)
require.NoError(t, err)
var decoded struct {
Choices []struct {
Delta struct {
Refusal *string `json:"refusal"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
}
require.NoError(t, json.Unmarshal(encoded, &decoded))
for _, choice := range decoded.Choices {
if choice.Delta.Refusal != nil {
_, _ = refusal.WriteString(*choice.Delta.Refusal)
}
if choice.FinishReason != nil {
finishReason = *choice.FinishReason
}
}
}
}
require.Equal(t, refusalText, refusal.String())
require.Equal(t, "stop", finishReason)
}
// Non-streaming Chat requests are served from a forced upstream stream; when
// the terminal event carries an empty output array the message is rebuilt from
// the accumulated deltas, which must include refusal deltas.
func TestBufferedResponseAccumulator_RefusalRebuiltFromDeltas(t *testing.T) {
acc := NewBufferedResponseAccumulator()
for _, raw := range []string{
`{"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant"}}`,
`{"type":"response.refusal.delta","output_index":0,"content_index":0,"delta":"I can't "}`,
`{"type":"response.refusal.delta","output_index":0,"content_index":0,"delta":"help with that."}`,
} {
evt := decodeResponsesEvent(t, raw)
acc.ProcessEvent(&evt)
}
require.True(t, acc.HasContent())
final := &ResponsesResponse{ID: "resp_refusal", Status: "completed"}
acc.SupplementResponseOutput(final)
chat := ResponsesToChatCompletions(final, "gpt-4o")
require.Equal(t, refusalText, chatMessageJSON(t, chat)["refusal"])
}
@@ -0,0 +1,24 @@
package apicompat
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
// Thinking blocks must always carry the `signature` key (Anthropic wire shape);
// strict clients such as Grok Build fail with "missing field `signature`".
func TestAnthropicContentBlock_ThinkingAlwaysHasSignature(t *testing.T) {
data, err := json.Marshal(AnthropicContentBlock{Type: "thinking", Thinking: "plan"})
require.NoError(t, err)
require.JSONEq(t, `{"type":"thinking","thinking":"plan","signature":""}`, string(data))
data, err = json.Marshal(AnthropicContentBlock{Type: "thinking", Signature: "sig"})
require.NoError(t, err)
require.JSONEq(t, `{"type":"thinking","thinking":"","signature":"sig"}`, string(data))
data, err = json.Marshal(AnthropicContentBlock{Type: "text", Text: "hi"})
require.NoError(t, err)
require.JSONEq(t, `{"type":"text","text":"hi"}`, string(data))
}
+10 -3
View File
@@ -95,10 +95,14 @@ func (b AnthropicContentBlock) MarshalJSON() ([]byte, error) {
anthropicContentBlock
}{Text: b.Text, anthropicContentBlock: anthropicContentBlock(b)})
case "thinking":
// Anthropic always sends `signature` on thinking blocks (empty on
// content_block_start); strict clients such as Grok Build reject the
// block with "missing field `signature`" when the key is absent.
return json.Marshal(struct {
Thinking string `json:"thinking"`
Thinking string `json:"thinking"`
Signature string `json:"signature"`
anthropicContentBlock
}{Thinking: b.Thinking, anthropicContentBlock: anthropicContentBlock(b)})
}{Thinking: b.Thinking, Signature: b.Signature, anthropicContentBlock: anthropicContentBlock(b)})
default:
return json.Marshal(base)
}
@@ -312,8 +316,9 @@ func (i *ResponsesInputItem) UnmarshalJSON(data []byte) error {
// ResponsesContentPart is a typed content part in a Responses message.
type ResponsesContentPart struct {
PromptCacheBreakpoint json.RawMessage `json:"prompt_cache_breakpoint,omitempty"`
Type string `json:"type"` // "input_text" | "output_text" | "input_image" | "input_file"
Type string `json:"type"` // "input_text" | "output_text" | "refusal" | "input_image" | "input_file"
Text string `json:"text,omitempty"`
Refusal string `json:"refusal,omitempty"` // type=refusal
ImageURL string `json:"image_url,omitempty"` // data URI for input_image
// input_file fields.
@@ -690,6 +695,7 @@ type ChatStreamOptions struct {
type ChatMessage struct {
Role string `json:"role"` // "system" | "user" | "assistant" | "tool" | "function"
Content json.RawMessage `json:"content,omitempty"`
Refusal string `json:"refusal,omitempty"`
ReasoningContent string `json:"reasoning_content,omitempty"`
Reasoning string `json:"reasoning,omitempty"`
Name string `json:"name,omitempty"`
@@ -830,6 +836,7 @@ type ChatChunkChoice struct {
type ChatDelta struct {
Role string `json:"role,omitempty"`
Content *string `json:"content,omitempty"` // pointer: omit when not present, null vs "" matters
Refusal *string `json:"refusal,omitempty"`
ReasoningContent *string `json:"reasoning_content,omitempty"`
Reasoning *string `json:"reasoning,omitempty"`
ToolCalls []ChatToolCall `json:"tool_calls,omitempty"`
+9
View File
@@ -20,6 +20,9 @@ const (
BetaFastMode = "fast-mode-2026-02-01"
// Legacy structured output compatibility; forwarded only when explicitly requested.
BetaStructuredOutputs = "structured-outputs-2025-11-13"
// Tool-change compatibility tokens; never enabled without an explicit client request.
BetaMidConversationToolChanges = "mid-conversation-tool-changes-2026-07-01"
BetaInlineTools = "inline-tools-2026-09-15"
// 新增(对齐官方 CLI 2.1.9x 以来的流量)
BetaPromptCachingScope = "prompt-caching-scope-2026-01-05"
@@ -220,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)
}
}
+54 -3
View File
@@ -1,6 +1,7 @@
package xai
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
@@ -14,6 +15,8 @@ import (
"strconv"
"strings"
"time"
"golang.org/x/net/html"
)
const (
@@ -147,7 +150,7 @@ func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) {
return nil, fmt.Errorf("open xAI device verification page: %w", SSOHTTPError{Status: status})
}
status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOVerifyURL, url.Values{"user_code": {device.UserCode}})
status, finalURL, body, err = f.do(ctx, http.MethodPost, SSOVerifyURL, url.Values{"user_code": {device.UserCode}})
if err != nil {
return nil, err
}
@@ -158,11 +161,19 @@ func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) {
return nil, errors.New("xAI device verification did not reach consent page")
}
status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOApproveURL, url.Values{
approval := url.Values{
"user_code": {device.UserCode},
"action": {"allow"},
"principal_type": {"User"},
"principal_id": {""},
}
if consentToken := ssoConsentToken(body); consentToken != "" {
approval.Set("consent_token", consentToken)
}
consentURL, _ := url.Parse(finalURL) // f.do only returns validated xAI URLs.
status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOApproveURL, approval, http.Header{
"Origin": {consentURL.Scheme + "://" + consentURL.Host},
"Referer": {finalURL},
})
if err != nil {
return nil, err
@@ -177,6 +188,37 @@ func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) {
return f.pollToken(ctx, device.DeviceCode, time.Duration(device.Interval)*time.Second, time.Duration(device.ExpiresIn)*time.Second)
}
// xAI binds device approval to a hidden token on the consent page. Keep it in
// this flow only; older consent pages without the field remain supported.
func ssoConsentToken(body []byte) string {
tokenizer := html.NewTokenizer(bytes.NewReader(body))
for {
switch tokenizer.Next() {
case html.ErrorToken:
return ""
case html.StartTagToken, html.SelfClosingTagToken:
token := tokenizer.Token()
if token.Data != "input" {
continue
}
var name, value, inputType string
for _, attr := range token.Attr {
switch attr.Key {
case "name":
name = attr.Val
case "value":
value = attr.Val
case "type":
inputType = attr.Val
}
}
if name == "consent_token" && strings.EqualFold(inputType, "hidden") {
return value
}
}
}
}
func (f *ssoDeviceFlow) pollToken(ctx context.Context, deviceCode string, interval, expiresIn time.Duration) (*TokenResponse, error) {
if interval < time.Second {
interval = time.Second
@@ -241,7 +283,7 @@ func (f *ssoDeviceFlow) pollToken(ctx context.Context, deviceCode string, interv
return nil, errors.New("xAI device flow token polling timed out")
}
func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form url.Values) (int, string, []byte, error) {
func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form url.Values, headers ...http.Header) (int, string, []byte, error) {
if !safeXAIAuthURL(endpoint) {
return 0, "", nil, errors.New("xAI OAuth URL is not trusted")
}
@@ -260,6 +302,15 @@ func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form ur
request.Header.Set("Accept", "application/json, text/html;q=0.9, */*;q=0.8")
request.Header.Set("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8")
request.Header.Set("User-Agent", f.userAgent)
if redirects == 0 {
for _, header := range headers {
for name, values := range header {
for _, value := range values {
request.Header.Add(name, value)
}
}
}
}
if cookie := f.cookieHeader(request.URL); cookie != "" {
request.Header.Set("Cookie", cookie)
}
+26 -1
View File
@@ -23,6 +23,10 @@ type ssoDeviceFakeClient struct {
func (c *ssoDeviceFakeClient) Do(req *http.Request) (*http.Response, error) {
c.cookieHeaders = append(c.cookieHeaders, req.Header.Get("Cookie"))
if req.URL.String() != SSOApproveURL {
require.Empty(c.t, req.Header.Get("Origin"))
require.Empty(c.t, req.Header.Get("Referer"))
}
switch req.URL.String() {
case SSOAccountsURL:
require.Equal(c.t, http.MethodGet, req.Method)
@@ -43,13 +47,16 @@ func (c *ssoDeviceFakeClient) Do(req *http.Request) (*http.Response, error) {
return ssoDeviceResponse(http.StatusFound, http.Header{"Location": {"/oauth2/device/consent"}}, ``), nil
case "https://auth.x.ai/oauth2/device/consent":
require.Equal(c.t, http.MethodGet, req.Method)
return ssoDeviceResponse(http.StatusOK, nil, `<html>consent</html>`), nil
return ssoDeviceResponse(http.StatusOK, nil, `<html><form><input type="hidden" name="consent_token" value="consent&amp;token"/></form></html>`), nil
case SSOApproveURL:
require.Equal(c.t, http.MethodPost, req.Method)
require.Equal(c.t, "https://auth.x.ai", req.Header.Get("Origin"))
require.Equal(c.t, "https://auth.x.ai/oauth2/device/consent", req.Header.Get("Referer"))
values := readSSODeviceForm(c.t, req)
require.Equal(c.t, "USER-1", values.Get("user_code"))
require.Equal(c.t, "allow", values.Get("action"))
require.Equal(c.t, "User", values.Get("principal_type"))
require.Equal(c.t, "consent&token", values.Get("consent_token"))
return ssoDeviceResponse(http.StatusSeeOther, http.Header{"Location": {"/oauth2/device/done"}}, ``), nil
case "https://auth.x.ai/oauth2/device/done":
require.Equal(c.t, http.MethodGet, req.Method)
@@ -87,6 +94,24 @@ func TestConvertSSOToBuildCompletesDeviceFlow(t *testing.T) {
require.Contains(t, client.cookieHeaders[0], "sso-rw=sso-token")
require.Contains(t, client.cookieHeaders[len(client.cookieHeaders)-1], "session=web-session")
require.Contains(t, client.cookieHeaders[len(client.cookieHeaders)-1], "csrf=csrf-token")
for _, cookie := range client.cookieHeaders {
require.NotContains(t, cookie, "consent&token")
}
}
func TestSSOConsentToken(t *testing.T) {
for _, tc := range []struct {
name, body, want string
}{
{"legacy page", `<html>consent</html>`, ""},
{"encoded hidden value", `<input value='a&amp;b&#43;c' name='consent_token' type='hidden'>`, "a&b+c"},
{"case insensitive HTML", `<INPUT TYPE="HIDDEN" NAME="consent_token" VALUE="token"/>`, "token"},
{"ignore script and visible fields", `<script>"<input type='hidden' name='consent_token' value='fake'>"</script><input name='consent_token' value='visible'><input type='hidden' name='user_code' value='user'><input type='hidden' name='consent_token' value='real'>`, "real"},
} {
t.Run(tc.name, func(t *testing.T) {
require.Equal(t, tc.want, ssoConsentToken([]byte(tc.body)))
})
}
}
func TestNormalizeSSOTokenAcceptsCookieHeader(t *testing.T) {
+3 -2
View File
@@ -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))
}

Some files were not shown because too many files have changed in this diff Show More