mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Compare commits
91
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
432a6a461f | ||
|
|
c2d5bbd93e | ||
|
|
ae501cc220 | ||
|
|
99efaa477f | ||
|
|
d9d4333795 | ||
|
|
458b92abd4 | ||
|
|
5106065716 | ||
|
|
46f5f5e8ba | ||
|
|
2529371bab | ||
|
|
43c4c882fd | ||
|
|
719013053a | ||
|
|
f3dc0835df | ||
|
|
5637f469fa | ||
|
|
34fe845f5a | ||
|
|
fa49647c6d | ||
|
|
9eb394e3e9 | ||
|
|
2d8a386811 | ||
|
|
bb97ab364a | ||
|
|
7225538358 | ||
|
|
0cecf3c367 | ||
|
|
1aa34d478d | ||
|
|
1b1039f4fd | ||
|
|
d6adebd22d | ||
|
|
53158587a6 | ||
|
|
ad05eda157 | ||
|
|
fdce5b8d01 | ||
|
|
00940743d2 | ||
|
|
2e3c2cf592 | ||
|
|
42bc7f6cff | ||
|
|
96f4c115c9 | ||
|
|
327c322184 | ||
|
|
48b7109588 | ||
|
|
4f0c079e27 | ||
|
|
b0cfff9826 | ||
|
|
017bbcb901 | ||
|
|
9ecb340822 | ||
|
|
a0f41f95a0 | ||
|
|
b5efbe3f4c | ||
|
|
e6d191a83f | ||
|
|
c91bb6124a | ||
|
|
9688571a83 | ||
|
|
514ad0aa8e | ||
|
|
dc2efc6b9a | ||
|
|
1872d44b58 | ||
|
|
196b47757a | ||
|
|
a60a29549f | ||
|
|
2f3fed2fdb | ||
|
|
fe9dadfc7d | ||
|
|
41dcfec489 | ||
|
|
052e08e764 | ||
|
|
1f955af353 | ||
|
|
ccb027b515 | ||
|
|
f50b99b66c | ||
|
|
883a53fea0 | ||
|
|
b35a611538 | ||
|
|
1dca86f287 | ||
|
|
dd6cb410e6 | ||
|
|
27a9421e6a | ||
|
|
860782e6fb | ||
|
|
35232c3949 | ||
|
|
421da3f0d4 | ||
|
|
48970f61b4 | ||
|
|
f8984b8a93 | ||
|
|
6a8f921650 | ||
|
|
2cc459e24f | ||
|
|
475c0839c6 | ||
|
|
b886fb8edb | ||
|
|
8490a81867 | ||
|
|
652a7d74f9 | ||
|
|
52c123e6f4 | ||
|
|
43f88c99aa | ||
|
|
6073704fd7 | ||
|
|
6a69dd051a | ||
|
|
2840cdeced | ||
|
|
b1ce3c6d59 | ||
|
|
d359196adc | ||
|
|
cf19f30d26 | ||
|
|
aa3a3ea787 | ||
|
|
bd48b2b2d2 | ||
|
|
9a62841fd1 | ||
|
|
58d605083e | ||
|
|
0da18e4a5e | ||
|
|
b31b435091 | ||
|
|
e2c35cf394 | ||
|
|
453b796df0 | ||
|
|
e878ab710f | ||
|
|
13fcfc0f3b | ||
|
|
0416c9ede4 | ||
|
|
ece1ecccca | ||
|
|
b13200d7a1 | ||
|
|
d1ba57977d |
@@ -0,0 +1,107 @@
|
||||
# Security Policy
|
||||
|
||||
[English](#english) | [中文](#中文)
|
||||
|
||||
## English
|
||||
|
||||
### Reporting a Vulnerability
|
||||
|
||||
**Please do not report security vulnerabilities through public GitHub issues, pull requests, or discussions.**
|
||||
|
||||
Report them privately through GitHub's private vulnerability reporting:
|
||||
|
||||
**[Report a vulnerability](https://github.com/Wei-Shaw/sub2api/security/advisories/new)**
|
||||
|
||||
(Repository → **Security** tab → **Report a vulnerability**.) Only the maintainers can see the report, and the advisory becomes the place to discuss, fix, and eventually publish the issue.
|
||||
|
||||
If you cannot use GitHub private vulnerability reporting, email **dev@sub2api.org** instead, with `[Security]` in the subject line.
|
||||
|
||||
A useful report includes:
|
||||
|
||||
- The affected component and version (release tag or commit)
|
||||
- Impact: what an attacker can do and which privileges they need
|
||||
- Minimal reproduction steps or a proof of concept
|
||||
- Any non-default configuration required to trigger the issue
|
||||
- A suggested fix, if you have one
|
||||
|
||||
### Supported Versions
|
||||
|
||||
Sub2API releases frequently, and fixes ship in a new release rather than being backported.
|
||||
|
||||
| Version | Supported |
|
||||
| ------- | --------- |
|
||||
| Latest release | ✅ |
|
||||
| Older releases | ❌ |
|
||||
|
||||
Please confirm that the issue reproduces on the latest release or on `main` before reporting.
|
||||
|
||||
### Scope
|
||||
|
||||
In scope: the Sub2API backend, frontend, official Docker images, and deployment files in this repository.
|
||||
|
||||
Usually out of scope:
|
||||
|
||||
- Issues that require already-compromised admin credentials or server access
|
||||
- Risks arising from explicitly weakening documented security settings (for example, disabling `security.url_allowlist`)
|
||||
- Vulnerabilities in upstream AI providers or third-party dependencies with no Sub2API-specific impact (please report those upstream)
|
||||
- Missing best-practice headers or scanner output without a demonstrated impact
|
||||
|
||||
### Disclosure Process
|
||||
|
||||
1. We acknowledge the report and assess its validity and severity.
|
||||
2. We develop a fix privately in the advisory and may ask you to verify it.
|
||||
3. We publish a release containing the fix, then publish the GitHub Security Advisory (with a CVE when applicable), crediting you unless you prefer to stay anonymous.
|
||||
|
||||
Please keep details private until the advisory is published. This is a volunteer-maintained project, so response times are best effort.
|
||||
|
||||
## 中文
|
||||
|
||||
### 报告漏洞
|
||||
|
||||
**请不要通过公开的 Issue、Pull Request 或 Discussion 报告安全漏洞。**
|
||||
|
||||
请使用 GitHub 私密漏洞报告功能提交:
|
||||
|
||||
**[提交漏洞报告](https://github.com/Wei-Shaw/sub2api/security/advisories/new)**
|
||||
|
||||
(仓库 → **Security** 标签页 → **Report a vulnerability**。)报告仅维护者可见,后续的讨论、修复与公告发布都在该安全公告中进行。
|
||||
|
||||
如果无法使用 GitHub 私密漏洞报告,也可以发送邮件至 **dev@sub2api.org**,邮件标题请注明 `[Security]`。
|
||||
|
||||
报告中建议包含:
|
||||
|
||||
- 受影响的组件与版本(Release tag 或 commit)
|
||||
- 影响:攻击者能做什么、需要什么权限
|
||||
- 最小复现步骤或 PoC
|
||||
- 触发问题所需的非默认配置
|
||||
- 修复建议(如有)
|
||||
|
||||
### 支持的版本
|
||||
|
||||
Sub2API 发版频繁,安全修复随新版本发布,不向旧版本回移。
|
||||
|
||||
| 版本 | 是否支持 |
|
||||
| ---- | -------- |
|
||||
| 最新 Release | ✅ |
|
||||
| 更早的版本 | ❌ |
|
||||
|
||||
报告前请先确认问题在最新 Release 或 `main` 上仍可复现。
|
||||
|
||||
### 范围
|
||||
|
||||
范围内:本仓库中的 Sub2API 后端、前端、官方 Docker 镜像与部署文件。
|
||||
|
||||
通常不在范围内:
|
||||
|
||||
- 需要事先已掌握管理员凭据或服务器权限才能利用的问题
|
||||
- 因主动关闭文档中的安全配置(如关闭 `security.url_allowlist`)而产生的风险
|
||||
- 上游 AI 服务商或第三方依赖自身的漏洞,且对 Sub2API 无特定影响(请向上游报告)
|
||||
- 仅缺少最佳实践响应头或扫描器输出,但未证明实际影响
|
||||
|
||||
### 披露流程
|
||||
|
||||
1. 我们确认收到报告,并评估其有效性与严重程度。
|
||||
2. 我们在安全公告中私下开发修复,可能请你协助验证。
|
||||
3. 我们先发布包含修复的版本,再公开 GitHub 安全公告(适用时申请 CVE),并致谢报告者(如你希望匿名则不署名)。
|
||||
|
||||
在公告发布前请勿公开漏洞细节。本项目由志愿者维护,响应时间尽力而为。
|
||||
@@ -827,6 +827,31 @@ Administrators can override automatic media eligibility through the account crea
|
||||
|
||||
---
|
||||
|
||||
## TypeSafe / Jev Support
|
||||
|
||||
Sub2API supports TypeSafe API-key accounts through Jev's native, non-streaming System One protocol.
|
||||
|
||||
- Platform: `typesafe`; account type: API Key
|
||||
- Default upstream: `https://api.typesafe.ai`
|
||||
- Public endpoint: `POST /v1/systemone`
|
||||
- Model: `jev-latest`, also returned by `/v1/models` for TypeSafe groups
|
||||
- Questions: `noul`, `choice`, and `score`
|
||||
|
||||
Requests and successful responses retain the native System One JSON structure. This endpoint is not compatible with Chat Completions, Responses, Anthropic Messages, or streaming clients.
|
||||
|
||||
Question validation follows the TypeSafe OpenAPI wire schema (also used by SDK v0.5.7). `instructions` may be omitted or `null` for all question types. Noul `criteria` may be omitted or `null`; its `true`/`false` descriptions and Choice descriptions accept strings, objects, arrays, or `null`. Score `criteria` must be a non-empty array of string, object, or array descriptions; a single level is valid. SDK integer-keyed Score maps are normalized to arrays by the SDK before sending.
|
||||
|
||||
```bash
|
||||
curl https://your-sub2api.example.com/v1/systemone \
|
||||
-H 'Authorization: Bearer sk-your-sub2api-key' \
|
||||
-H 'Content-Type: application/json' \
|
||||
--data '{"model":"jev-latest","state":"Text to evaluate","questions":{"safety":{"type":"noul","instructions":"Evaluate whether the text is unsafe"}}}'
|
||||
```
|
||||
|
||||
The built-in `jev-latest` price is `$0.042` per million input tokens and `$0` for output tokens. Channel pricing can override both values. Credential, billing, permission, rate-limit, overload, server, and network failures (`401`, `402`, `403`, `429`, `529`, `5xx`, transport errors) use the existing account error policy (including custom error codes and temporary-unschedulable rules) and fail over to another account; request errors (`400`, `413`, and `422`) are returned without retrying another account and never change account state. TypeSafe groups (and Composite requests routed to TypeSafe) reject Messages, Chat Completions, Responses, and count_tokens requests with `404`.
|
||||
|
||||
---
|
||||
|
||||
## Antigravity Support
|
||||
|
||||
Sub2API supports [Antigravity](https://antigravity.so/) accounts. After authorization, dedicated endpoints are available for Claude and Gemini models.
|
||||
|
||||
@@ -743,6 +743,31 @@ OAuth / Setup Token 图片请求使用 Responses 主控模型调用 `image_gener
|
||||
|
||||
---
|
||||
|
||||
## TypeSafe / Jev 使用说明
|
||||
|
||||
Sub2API 支持使用 TypeSafe API Key 账户,通过 Jev 原生、非流式的 System One 协议调用模型。
|
||||
|
||||
- 平台:`typesafe`;账号类型:API Key
|
||||
- 默认上游:`https://api.typesafe.ai`
|
||||
- 对外端点:`POST /v1/systemone`
|
||||
- 模型:`jev-latest`,TypeSafe 分组的 `/v1/models` 也会返回该模型
|
||||
- 问题类型:`noul`、`choice`、`score`
|
||||
|
||||
请求和成功响应保持 System One 原生 JSON 结构。该端点不兼容 Chat Completions、Responses、Anthropic Messages 或流式客户端。
|
||||
|
||||
问题校验遵循 TypeSafe OpenAPI 的线上协议 schema(SDK v0.5.7 也使用该 schema)。所有问题的 `instructions` 都可以省略或为 `null`。Noul 的 `criteria` 可以省略或为 `null`,其中 `true`/`false` 的描述和 Choice 描述支持字符串、对象、数组或 `null`。Score 的 `criteria` 必须是至少包含一档描述的数组,每档支持字符串、对象或数组;单档也合法。SDK 的整数键 Score 映射会由 SDK 在发送前转换为数组。
|
||||
|
||||
```bash
|
||||
curl https://your-sub2api.example.com/v1/systemone \
|
||||
-H 'Authorization: Bearer sk-your-sub2api-key' \
|
||||
-H 'Content-Type: application/json' \
|
||||
--data '{"model":"jev-latest","state":"待评估文本","questions":{"safety":{"type":"noul","instructions":"评估文本是否不安全"}}}'
|
||||
```
|
||||
|
||||
`jev-latest` 内置价格为输入 `$0.042/百万 tokens`、输出 `$0`,渠道定价可以覆盖。凭据、欠费、权限、限流、过载、服务端和网络错误(`401`、`402`、`403`、`429`、`529`、`5xx`、传输错误)沿用现有账号错误策略(含自定义错误码与临时不可调度规则)并切换账号;请求错误(`400`、`413`、`422`)不会切换账号重试,也不会改变账号状态。TypeSafe 分组(以及路由到 TypeSafe 的 Composite 请求)调用 Messages、Chat Completions、Responses、count_tokens 时返回 `404`。
|
||||
|
||||
---
|
||||
|
||||
## Antigravity 使用说明
|
||||
|
||||
Sub2API 支持 [Antigravity](https://antigravity.so/) 账户,授权后可通过专用端点访问 Claude 和 Gemini 模型。
|
||||
|
||||
@@ -1 +1 @@
|
||||
0.2.8
|
||||
0.2.12
|
||||
|
||||
@@ -288,14 +288,16 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
auditLogHandler := admin.NewAuditLogHandler(auditLogService, totpService)
|
||||
upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db)
|
||||
openCodeGoUsageService := service.ProvideOpenCodeGoUsageService(accountRepository, httpUpstream, settingService, leaderLockCache, db)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService)
|
||||
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
|
||||
claudeResetCreditService := service.ProvideClaudeResetCreditService(accountRepository, claudeTokenProvider, proxyRepository, settingService, idempotencyCoordinator, leaderLockCache)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, claudeResetCreditService)
|
||||
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
|
||||
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
|
||||
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)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator)
|
||||
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)
|
||||
passkeyRepository := repository.NewPasskeyRepository(db)
|
||||
@@ -321,7 +323,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig)
|
||||
batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig)
|
||||
batchImageHandler := handler.ProvideBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService, openAIGatewayHandler)
|
||||
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
|
||||
idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig)
|
||||
openAIQuotaAutoResetService := service.ProvideOpenAIQuotaAutoResetService(accountRepository, openAIQuotaService, rateLimitService, idempotencyCoordinator, auditLogService, settingService, leaderLockCache)
|
||||
handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, channelMonitorV2Handler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, passkeyHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, modelPlazaHandler, asyncImageHandler, batchImageHandler, idempotencyCoordinator, idempotencyCleanupService, openAIQuotaAutoResetService)
|
||||
|
||||
@@ -1128,6 +1128,7 @@ var (
|
||||
{Name: "amount", Type: field.TypeFloat64, SchemaType: map[string]string{"postgres": "decimal(20,2)"}},
|
||||
{Name: "pay_amount", Type: field.TypeFloat64, SchemaType: map[string]string{"postgres": "decimal(20,2)"}},
|
||||
{Name: "fee_rate", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(10,4)"}},
|
||||
{Name: "bonus_amount", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,2)"}},
|
||||
{Name: "recharge_code", Type: field.TypeString, Size: 64},
|
||||
{Name: "out_trade_no", Type: field.TypeString, Size: 64, Default: ""},
|
||||
{Name: "payment_type", Type: field.TypeString, Size: 30},
|
||||
@@ -1170,7 +1171,7 @@ var (
|
||||
ForeignKeys: []*schema.ForeignKey{
|
||||
{
|
||||
Symbol: "payment_orders_users_payment_orders",
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[39]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[40]},
|
||||
RefColumns: []*schema.Column{UsersColumns[0]},
|
||||
OnDelete: schema.NoAction,
|
||||
},
|
||||
@@ -1179,7 +1180,7 @@ var (
|
||||
{
|
||||
Name: "paymentorder_out_trade_no",
|
||||
Unique: true,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[8]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[9]},
|
||||
Annotation: &entsql.IndexAnnotation{
|
||||
Where: "out_trade_no <> ''",
|
||||
},
|
||||
@@ -1187,37 +1188,37 @@ var (
|
||||
{
|
||||
Name: "paymentorder_user_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[39]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[40]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_status",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[21]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[22]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_expires_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[29]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[30]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[37]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[38]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_paid_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[30]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[31]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_payment_type_paid_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[9], PaymentOrdersColumns[30]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[10], PaymentOrdersColumns[31]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_order_type",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[14]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[15]},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
+88
-1
@@ -30113,6 +30113,8 @@ type PaymentOrderMutation struct {
|
||||
addpay_amount *float64
|
||||
fee_rate *float64
|
||||
addfee_rate *float64
|
||||
bonus_amount *float64
|
||||
addbonus_amount *float64
|
||||
recharge_code *string
|
||||
out_trade_no *string
|
||||
payment_type *string
|
||||
@@ -30580,6 +30582,62 @@ func (m *PaymentOrderMutation) ResetFeeRate() {
|
||||
m.addfee_rate = nil
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (m *PaymentOrderMutation) SetBonusAmount(f float64) {
|
||||
m.bonus_amount = &f
|
||||
m.addbonus_amount = nil
|
||||
}
|
||||
|
||||
// BonusAmount returns the value of the "bonus_amount" field in the mutation.
|
||||
func (m *PaymentOrderMutation) BonusAmount() (r float64, exists bool) {
|
||||
v := m.bonus_amount
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldBonusAmount returns the old "bonus_amount" field's value of the PaymentOrder entity.
|
||||
// If the PaymentOrder 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 *PaymentOrderMutation) OldBonusAmount(ctx context.Context) (v float64, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldBonusAmount is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldBonusAmount requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldBonusAmount: %w", err)
|
||||
}
|
||||
return oldValue.BonusAmount, nil
|
||||
}
|
||||
|
||||
// AddBonusAmount adds f to the "bonus_amount" field.
|
||||
func (m *PaymentOrderMutation) AddBonusAmount(f float64) {
|
||||
if m.addbonus_amount != nil {
|
||||
*m.addbonus_amount += f
|
||||
} else {
|
||||
m.addbonus_amount = &f
|
||||
}
|
||||
}
|
||||
|
||||
// AddedBonusAmount returns the value that was added to the "bonus_amount" field in this mutation.
|
||||
func (m *PaymentOrderMutation) AddedBonusAmount() (r float64, exists bool) {
|
||||
v := m.addbonus_amount
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// ResetBonusAmount resets all changes to the "bonus_amount" field.
|
||||
func (m *PaymentOrderMutation) ResetBonusAmount() {
|
||||
m.bonus_amount = nil
|
||||
m.addbonus_amount = nil
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (m *PaymentOrderMutation) SetRechargeCode(s string) {
|
||||
m.recharge_code = &s
|
||||
@@ -32123,7 +32181,7 @@ func (m *PaymentOrderMutation) Type() string {
|
||||
// order to get all numeric fields that were incremented/decremented, call
|
||||
// AddedFields().
|
||||
func (m *PaymentOrderMutation) Fields() []string {
|
||||
fields := make([]string, 0, 39)
|
||||
fields := make([]string, 0, 40)
|
||||
if m.user != nil {
|
||||
fields = append(fields, paymentorder.FieldUserID)
|
||||
}
|
||||
@@ -32145,6 +32203,9 @@ func (m *PaymentOrderMutation) Fields() []string {
|
||||
if m.fee_rate != nil {
|
||||
fields = append(fields, paymentorder.FieldFeeRate)
|
||||
}
|
||||
if m.bonus_amount != nil {
|
||||
fields = append(fields, paymentorder.FieldBonusAmount)
|
||||
}
|
||||
if m.recharge_code != nil {
|
||||
fields = append(fields, paymentorder.FieldRechargeCode)
|
||||
}
|
||||
@@ -32263,6 +32324,8 @@ func (m *PaymentOrderMutation) Field(name string) (ent.Value, bool) {
|
||||
return m.PayAmount()
|
||||
case paymentorder.FieldFeeRate:
|
||||
return m.FeeRate()
|
||||
case paymentorder.FieldBonusAmount:
|
||||
return m.BonusAmount()
|
||||
case paymentorder.FieldRechargeCode:
|
||||
return m.RechargeCode()
|
||||
case paymentorder.FieldOutTradeNo:
|
||||
@@ -32350,6 +32413,8 @@ func (m *PaymentOrderMutation) OldField(ctx context.Context, name string) (ent.V
|
||||
return m.OldPayAmount(ctx)
|
||||
case paymentorder.FieldFeeRate:
|
||||
return m.OldFeeRate(ctx)
|
||||
case paymentorder.FieldBonusAmount:
|
||||
return m.OldBonusAmount(ctx)
|
||||
case paymentorder.FieldRechargeCode:
|
||||
return m.OldRechargeCode(ctx)
|
||||
case paymentorder.FieldOutTradeNo:
|
||||
@@ -32472,6 +32537,13 @@ func (m *PaymentOrderMutation) SetField(name string, value ent.Value) error {
|
||||
}
|
||||
m.SetFeeRate(v)
|
||||
return nil
|
||||
case paymentorder.FieldBonusAmount:
|
||||
v, ok := value.(float64)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetBonusAmount(v)
|
||||
return nil
|
||||
case paymentorder.FieldRechargeCode:
|
||||
v, ok := value.(string)
|
||||
if !ok {
|
||||
@@ -32713,6 +32785,9 @@ func (m *PaymentOrderMutation) AddedFields() []string {
|
||||
if m.addfee_rate != nil {
|
||||
fields = append(fields, paymentorder.FieldFeeRate)
|
||||
}
|
||||
if m.addbonus_amount != nil {
|
||||
fields = append(fields, paymentorder.FieldBonusAmount)
|
||||
}
|
||||
if m.addplan_id != nil {
|
||||
fields = append(fields, paymentorder.FieldPlanID)
|
||||
}
|
||||
@@ -32739,6 +32814,8 @@ func (m *PaymentOrderMutation) AddedField(name string) (ent.Value, bool) {
|
||||
return m.AddedPayAmount()
|
||||
case paymentorder.FieldFeeRate:
|
||||
return m.AddedFeeRate()
|
||||
case paymentorder.FieldBonusAmount:
|
||||
return m.AddedBonusAmount()
|
||||
case paymentorder.FieldPlanID:
|
||||
return m.AddedPlanID()
|
||||
case paymentorder.FieldSubscriptionGroupID:
|
||||
@@ -32777,6 +32854,13 @@ func (m *PaymentOrderMutation) AddField(name string, value ent.Value) error {
|
||||
}
|
||||
m.AddFeeRate(v)
|
||||
return nil
|
||||
case paymentorder.FieldBonusAmount:
|
||||
v, ok := value.(float64)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.AddBonusAmount(v)
|
||||
return nil
|
||||
case paymentorder.FieldPlanID:
|
||||
v, ok := value.(int64)
|
||||
if !ok {
|
||||
@@ -32976,6 +33060,9 @@ func (m *PaymentOrderMutation) ResetField(name string) error {
|
||||
case paymentorder.FieldFeeRate:
|
||||
m.ResetFeeRate()
|
||||
return nil
|
||||
case paymentorder.FieldBonusAmount:
|
||||
m.ResetBonusAmount()
|
||||
return nil
|
||||
case paymentorder.FieldRechargeCode:
|
||||
m.ResetRechargeCode()
|
||||
return nil
|
||||
|
||||
@@ -33,6 +33,8 @@ type PaymentOrder struct {
|
||||
PayAmount float64 `json:"pay_amount,omitempty"`
|
||||
// FeeRate holds the value of the "fee_rate" field.
|
||||
FeeRate float64 `json:"fee_rate,omitempty"`
|
||||
// BonusAmount holds the value of the "bonus_amount" field.
|
||||
BonusAmount float64 `json:"bonus_amount,omitempty"`
|
||||
// RechargeCode holds the value of the "recharge_code" field.
|
||||
RechargeCode string `json:"recharge_code,omitempty"`
|
||||
// OutTradeNo holds the value of the "out_trade_no" field.
|
||||
@@ -132,7 +134,7 @@ func (*PaymentOrder) scanValues(columns []string) ([]any, error) {
|
||||
values[i] = new([]byte)
|
||||
case paymentorder.FieldForceRefund:
|
||||
values[i] = new(sql.NullBool)
|
||||
case paymentorder.FieldAmount, paymentorder.FieldPayAmount, paymentorder.FieldFeeRate, paymentorder.FieldRefundAmount:
|
||||
case paymentorder.FieldAmount, paymentorder.FieldPayAmount, paymentorder.FieldFeeRate, paymentorder.FieldBonusAmount, paymentorder.FieldRefundAmount:
|
||||
values[i] = new(sql.NullFloat64)
|
||||
case paymentorder.FieldID, paymentorder.FieldUserID, paymentorder.FieldPlanID, paymentorder.FieldSubscriptionGroupID, paymentorder.FieldSubscriptionDays:
|
||||
values[i] = new(sql.NullInt64)
|
||||
@@ -204,6 +206,12 @@ func (_m *PaymentOrder) assignValues(columns []string, values []any) error {
|
||||
} else if value.Valid {
|
||||
_m.FeeRate = value.Float64
|
||||
}
|
||||
case paymentorder.FieldBonusAmount:
|
||||
if value, ok := values[i].(*sql.NullFloat64); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field bonus_amount", values[i])
|
||||
} else if value.Valid {
|
||||
_m.BonusAmount = value.Float64
|
||||
}
|
||||
case paymentorder.FieldRechargeCode:
|
||||
if value, ok := values[i].(*sql.NullString); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field recharge_code", values[i])
|
||||
@@ -480,6 +488,9 @@ func (_m *PaymentOrder) String() string {
|
||||
builder.WriteString("fee_rate=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.FeeRate))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("bonus_amount=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.BonusAmount))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("recharge_code=")
|
||||
builder.WriteString(_m.RechargeCode)
|
||||
builder.WriteString(", ")
|
||||
|
||||
@@ -28,6 +28,8 @@ const (
|
||||
FieldPayAmount = "pay_amount"
|
||||
// FieldFeeRate holds the string denoting the fee_rate field in the database.
|
||||
FieldFeeRate = "fee_rate"
|
||||
// FieldBonusAmount holds the string denoting the bonus_amount field in the database.
|
||||
FieldBonusAmount = "bonus_amount"
|
||||
// FieldRechargeCode holds the string denoting the recharge_code field in the database.
|
||||
FieldRechargeCode = "recharge_code"
|
||||
// FieldOutTradeNo holds the string denoting the out_trade_no field in the database.
|
||||
@@ -115,6 +117,7 @@ var Columns = []string{
|
||||
FieldAmount,
|
||||
FieldPayAmount,
|
||||
FieldFeeRate,
|
||||
FieldBonusAmount,
|
||||
FieldRechargeCode,
|
||||
FieldOutTradeNo,
|
||||
FieldPaymentType,
|
||||
@@ -166,6 +169,8 @@ var (
|
||||
UserNameValidator func(string) error
|
||||
// DefaultFeeRate holds the default value on creation for the "fee_rate" field.
|
||||
DefaultFeeRate float64
|
||||
// DefaultBonusAmount holds the default value on creation for the "bonus_amount" field.
|
||||
DefaultBonusAmount float64
|
||||
// RechargeCodeValidator is a validator for the "recharge_code" field. It is called by the builders before save.
|
||||
RechargeCodeValidator func(string) error
|
||||
// DefaultOutTradeNo holds the default value on creation for the "out_trade_no" field.
|
||||
@@ -249,6 +254,11 @@ func ByFeeRate(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldFeeRate, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByBonusAmount orders the results by the bonus_amount field.
|
||||
func ByBonusAmount(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldBonusAmount, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByRechargeCode orders the results by the recharge_code field.
|
||||
func ByRechargeCode(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldRechargeCode, opts...).ToFunc()
|
||||
|
||||
@@ -90,6 +90,11 @@ func FeeRate(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldFeeRate, v))
|
||||
}
|
||||
|
||||
// BonusAmount applies equality check predicate on the "bonus_amount" field. It's identical to BonusAmountEQ.
|
||||
func BonusAmount(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// RechargeCode applies equality check predicate on the "recharge_code" field. It's identical to RechargeCodeEQ.
|
||||
func RechargeCode(v string) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldRechargeCode, v))
|
||||
@@ -590,6 +595,46 @@ func FeeRateLTE(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldLTE(FieldFeeRate, v))
|
||||
}
|
||||
|
||||
// BonusAmountEQ applies the EQ predicate on the "bonus_amount" field.
|
||||
func BonusAmountEQ(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountNEQ applies the NEQ predicate on the "bonus_amount" field.
|
||||
func BonusAmountNEQ(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldNEQ(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountIn applies the In predicate on the "bonus_amount" field.
|
||||
func BonusAmountIn(vs ...float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldIn(FieldBonusAmount, vs...))
|
||||
}
|
||||
|
||||
// BonusAmountNotIn applies the NotIn predicate on the "bonus_amount" field.
|
||||
func BonusAmountNotIn(vs ...float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldNotIn(FieldBonusAmount, vs...))
|
||||
}
|
||||
|
||||
// BonusAmountGT applies the GT predicate on the "bonus_amount" field.
|
||||
func BonusAmountGT(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldGT(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountGTE applies the GTE predicate on the "bonus_amount" field.
|
||||
func BonusAmountGTE(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldGTE(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountLT applies the LT predicate on the "bonus_amount" field.
|
||||
func BonusAmountLT(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldLT(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountLTE applies the LTE predicate on the "bonus_amount" field.
|
||||
func BonusAmountLTE(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldLTE(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// RechargeCodeEQ applies the EQ predicate on the "recharge_code" field.
|
||||
func RechargeCodeEQ(v string) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldRechargeCode, v))
|
||||
|
||||
@@ -81,6 +81,20 @@ func (_c *PaymentOrderCreate) SetNillableFeeRate(v *float64) *PaymentOrderCreate
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (_c *PaymentOrderCreate) SetBonusAmount(v float64) *PaymentOrderCreate {
|
||||
_c.mutation.SetBonusAmount(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetNillableBonusAmount sets the "bonus_amount" field if the given value is not nil.
|
||||
func (_c *PaymentOrderCreate) SetNillableBonusAmount(v *float64) *PaymentOrderCreate {
|
||||
if v != nil {
|
||||
_c.SetBonusAmount(*v)
|
||||
}
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (_c *PaymentOrderCreate) SetRechargeCode(v string) *PaymentOrderCreate {
|
||||
_c.mutation.SetRechargeCode(v)
|
||||
@@ -517,6 +531,10 @@ func (_c *PaymentOrderCreate) defaults() {
|
||||
v := paymentorder.DefaultFeeRate
|
||||
_c.mutation.SetFeeRate(v)
|
||||
}
|
||||
if _, ok := _c.mutation.BonusAmount(); !ok {
|
||||
v := paymentorder.DefaultBonusAmount
|
||||
_c.mutation.SetBonusAmount(v)
|
||||
}
|
||||
if _, ok := _c.mutation.OutTradeNo(); !ok {
|
||||
v := paymentorder.DefaultOutTradeNo
|
||||
_c.mutation.SetOutTradeNo(v)
|
||||
@@ -577,6 +595,9 @@ func (_c *PaymentOrderCreate) check() error {
|
||||
if _, ok := _c.mutation.FeeRate(); !ok {
|
||||
return &ValidationError{Name: "fee_rate", err: errors.New(`ent: missing required field "PaymentOrder.fee_rate"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.BonusAmount(); !ok {
|
||||
return &ValidationError{Name: "bonus_amount", err: errors.New(`ent: missing required field "PaymentOrder.bonus_amount"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.RechargeCode(); !ok {
|
||||
return &ValidationError{Name: "recharge_code", err: errors.New(`ent: missing required field "PaymentOrder.recharge_code"`)}
|
||||
}
|
||||
@@ -725,6 +746,10 @@ func (_c *PaymentOrderCreate) createSpec() (*PaymentOrder, *sqlgraph.CreateSpec)
|
||||
_spec.SetField(paymentorder.FieldFeeRate, field.TypeFloat64, value)
|
||||
_node.FeeRate = value
|
||||
}
|
||||
if value, ok := _c.mutation.BonusAmount(); ok {
|
||||
_spec.SetField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
_node.BonusAmount = value
|
||||
}
|
||||
if value, ok := _c.mutation.RechargeCode(); ok {
|
||||
_spec.SetField(paymentorder.FieldRechargeCode, field.TypeString, value)
|
||||
_node.RechargeCode = value
|
||||
@@ -1030,6 +1055,24 @@ func (u *PaymentOrderUpsert) AddFeeRate(v float64) *PaymentOrderUpsert {
|
||||
return u
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsert) SetBonusAmount(v float64) *PaymentOrderUpsert {
|
||||
u.Set(paymentorder.FieldBonusAmount, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateBonusAmount sets the "bonus_amount" field to the value that was provided on create.
|
||||
func (u *PaymentOrderUpsert) UpdateBonusAmount() *PaymentOrderUpsert {
|
||||
u.SetExcluded(paymentorder.FieldBonusAmount)
|
||||
return u
|
||||
}
|
||||
|
||||
// AddBonusAmount adds v to the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsert) AddBonusAmount(v float64) *PaymentOrderUpsert {
|
||||
u.Add(paymentorder.FieldBonusAmount, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (u *PaymentOrderUpsert) SetRechargeCode(v string) *PaymentOrderUpsert {
|
||||
u.Set(paymentorder.FieldRechargeCode, v)
|
||||
@@ -1711,6 +1754,27 @@ func (u *PaymentOrderUpsertOne) UpdateFeeRate() *PaymentOrderUpsertOne {
|
||||
})
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsertOne) SetBonusAmount(v float64) *PaymentOrderUpsertOne {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.SetBonusAmount(v)
|
||||
})
|
||||
}
|
||||
|
||||
// AddBonusAmount adds v to the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsertOne) AddBonusAmount(v float64) *PaymentOrderUpsertOne {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.AddBonusAmount(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateBonusAmount sets the "bonus_amount" field to the value that was provided on create.
|
||||
func (u *PaymentOrderUpsertOne) UpdateBonusAmount() *PaymentOrderUpsertOne {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.UpdateBonusAmount()
|
||||
})
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (u *PaymentOrderUpsertOne) SetRechargeCode(v string) *PaymentOrderUpsertOne {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
@@ -2643,6 +2707,27 @@ func (u *PaymentOrderUpsertBulk) UpdateFeeRate() *PaymentOrderUpsertBulk {
|
||||
})
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsertBulk) SetBonusAmount(v float64) *PaymentOrderUpsertBulk {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.SetBonusAmount(v)
|
||||
})
|
||||
}
|
||||
|
||||
// AddBonusAmount adds v to the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsertBulk) AddBonusAmount(v float64) *PaymentOrderUpsertBulk {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.AddBonusAmount(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateBonusAmount sets the "bonus_amount" field to the value that was provided on create.
|
||||
func (u *PaymentOrderUpsertBulk) UpdateBonusAmount() *PaymentOrderUpsertBulk {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.UpdateBonusAmount()
|
||||
})
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (u *PaymentOrderUpsertBulk) SetRechargeCode(v string) *PaymentOrderUpsertBulk {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
|
||||
@@ -154,6 +154,27 @@ func (_u *PaymentOrderUpdate) AddFeeRate(v float64) *PaymentOrderUpdate {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (_u *PaymentOrderUpdate) SetBonusAmount(v float64) *PaymentOrderUpdate {
|
||||
_u.mutation.ResetBonusAmount()
|
||||
_u.mutation.SetBonusAmount(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableBonusAmount sets the "bonus_amount" field if the given value is not nil.
|
||||
func (_u *PaymentOrderUpdate) SetNillableBonusAmount(v *float64) *PaymentOrderUpdate {
|
||||
if v != nil {
|
||||
_u.SetBonusAmount(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// AddBonusAmount adds value to the "bonus_amount" field.
|
||||
func (_u *PaymentOrderUpdate) AddBonusAmount(v float64) *PaymentOrderUpdate {
|
||||
_u.mutation.AddBonusAmount(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (_u *PaymentOrderUpdate) SetRechargeCode(v string) *PaymentOrderUpdate {
|
||||
_u.mutation.SetRechargeCode(v)
|
||||
@@ -881,6 +902,12 @@ func (_u *PaymentOrderUpdate) sqlSave(ctx context.Context) (_node int, err error
|
||||
if value, ok := _u.mutation.AddedFeeRate(); ok {
|
||||
_spec.AddField(paymentorder.FieldFeeRate, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.BonusAmount(); ok {
|
||||
_spec.SetField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.AddedBonusAmount(); ok {
|
||||
_spec.AddField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.RechargeCode(); ok {
|
||||
_spec.SetField(paymentorder.FieldRechargeCode, field.TypeString, value)
|
||||
}
|
||||
@@ -1217,6 +1244,27 @@ func (_u *PaymentOrderUpdateOne) AddFeeRate(v float64) *PaymentOrderUpdateOne {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (_u *PaymentOrderUpdateOne) SetBonusAmount(v float64) *PaymentOrderUpdateOne {
|
||||
_u.mutation.ResetBonusAmount()
|
||||
_u.mutation.SetBonusAmount(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableBonusAmount sets the "bonus_amount" field if the given value is not nil.
|
||||
func (_u *PaymentOrderUpdateOne) SetNillableBonusAmount(v *float64) *PaymentOrderUpdateOne {
|
||||
if v != nil {
|
||||
_u.SetBonusAmount(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// AddBonusAmount adds value to the "bonus_amount" field.
|
||||
func (_u *PaymentOrderUpdateOne) AddBonusAmount(v float64) *PaymentOrderUpdateOne {
|
||||
_u.mutation.AddBonusAmount(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (_u *PaymentOrderUpdateOne) SetRechargeCode(v string) *PaymentOrderUpdateOne {
|
||||
_u.mutation.SetRechargeCode(v)
|
||||
@@ -1974,6 +2022,12 @@ func (_u *PaymentOrderUpdateOne) sqlSave(ctx context.Context) (_node *PaymentOrd
|
||||
if value, ok := _u.mutation.AddedFeeRate(); ok {
|
||||
_spec.AddField(paymentorder.FieldFeeRate, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.BonusAmount(); ok {
|
||||
_spec.SetField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.AddedBonusAmount(); ok {
|
||||
_spec.AddField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.RechargeCode(); ok {
|
||||
_spec.SetField(paymentorder.FieldRechargeCode, field.TypeString, value)
|
||||
}
|
||||
|
||||
@@ -1323,70 +1323,74 @@ func init() {
|
||||
paymentorderDescFeeRate := paymentorderFields[6].Descriptor()
|
||||
// paymentorder.DefaultFeeRate holds the default value on creation for the fee_rate field.
|
||||
paymentorder.DefaultFeeRate = paymentorderDescFeeRate.Default.(float64)
|
||||
// paymentorderDescBonusAmount is the schema descriptor for bonus_amount field.
|
||||
paymentorderDescBonusAmount := paymentorderFields[7].Descriptor()
|
||||
// paymentorder.DefaultBonusAmount holds the default value on creation for the bonus_amount field.
|
||||
paymentorder.DefaultBonusAmount = paymentorderDescBonusAmount.Default.(float64)
|
||||
// paymentorderDescRechargeCode is the schema descriptor for recharge_code field.
|
||||
paymentorderDescRechargeCode := paymentorderFields[7].Descriptor()
|
||||
paymentorderDescRechargeCode := paymentorderFields[8].Descriptor()
|
||||
// paymentorder.RechargeCodeValidator is a validator for the "recharge_code" field. It is called by the builders before save.
|
||||
paymentorder.RechargeCodeValidator = paymentorderDescRechargeCode.Validators[0].(func(string) error)
|
||||
// paymentorderDescOutTradeNo is the schema descriptor for out_trade_no field.
|
||||
paymentorderDescOutTradeNo := paymentorderFields[8].Descriptor()
|
||||
paymentorderDescOutTradeNo := paymentorderFields[9].Descriptor()
|
||||
// paymentorder.DefaultOutTradeNo holds the default value on creation for the out_trade_no field.
|
||||
paymentorder.DefaultOutTradeNo = paymentorderDescOutTradeNo.Default.(string)
|
||||
// paymentorder.OutTradeNoValidator is a validator for the "out_trade_no" field. It is called by the builders before save.
|
||||
paymentorder.OutTradeNoValidator = paymentorderDescOutTradeNo.Validators[0].(func(string) error)
|
||||
// paymentorderDescPaymentType is the schema descriptor for payment_type field.
|
||||
paymentorderDescPaymentType := paymentorderFields[9].Descriptor()
|
||||
paymentorderDescPaymentType := paymentorderFields[10].Descriptor()
|
||||
// paymentorder.PaymentTypeValidator is a validator for the "payment_type" field. It is called by the builders before save.
|
||||
paymentorder.PaymentTypeValidator = paymentorderDescPaymentType.Validators[0].(func(string) error)
|
||||
// paymentorderDescPaymentTradeNo is the schema descriptor for payment_trade_no field.
|
||||
paymentorderDescPaymentTradeNo := paymentorderFields[10].Descriptor()
|
||||
paymentorderDescPaymentTradeNo := paymentorderFields[11].Descriptor()
|
||||
// paymentorder.PaymentTradeNoValidator is a validator for the "payment_trade_no" field. It is called by the builders before save.
|
||||
paymentorder.PaymentTradeNoValidator = paymentorderDescPaymentTradeNo.Validators[0].(func(string) error)
|
||||
// paymentorderDescOrderType is the schema descriptor for order_type field.
|
||||
paymentorderDescOrderType := paymentorderFields[14].Descriptor()
|
||||
paymentorderDescOrderType := paymentorderFields[15].Descriptor()
|
||||
// paymentorder.DefaultOrderType holds the default value on creation for the order_type field.
|
||||
paymentorder.DefaultOrderType = paymentorderDescOrderType.Default.(string)
|
||||
// paymentorder.OrderTypeValidator is a validator for the "order_type" field. It is called by the builders before save.
|
||||
paymentorder.OrderTypeValidator = paymentorderDescOrderType.Validators[0].(func(string) error)
|
||||
// paymentorderDescProviderInstanceID is the schema descriptor for provider_instance_id field.
|
||||
paymentorderDescProviderInstanceID := paymentorderFields[18].Descriptor()
|
||||
paymentorderDescProviderInstanceID := paymentorderFields[19].Descriptor()
|
||||
// paymentorder.ProviderInstanceIDValidator is a validator for the "provider_instance_id" field. It is called by the builders before save.
|
||||
paymentorder.ProviderInstanceIDValidator = paymentorderDescProviderInstanceID.Validators[0].(func(string) error)
|
||||
// paymentorderDescProviderKey is the schema descriptor for provider_key field.
|
||||
paymentorderDescProviderKey := paymentorderFields[19].Descriptor()
|
||||
paymentorderDescProviderKey := paymentorderFields[20].Descriptor()
|
||||
// paymentorder.ProviderKeyValidator is a validator for the "provider_key" field. It is called by the builders before save.
|
||||
paymentorder.ProviderKeyValidator = paymentorderDescProviderKey.Validators[0].(func(string) error)
|
||||
// paymentorderDescStatus is the schema descriptor for status field.
|
||||
paymentorderDescStatus := paymentorderFields[21].Descriptor()
|
||||
paymentorderDescStatus := paymentorderFields[22].Descriptor()
|
||||
// paymentorder.DefaultStatus holds the default value on creation for the status field.
|
||||
paymentorder.DefaultStatus = paymentorderDescStatus.Default.(string)
|
||||
// paymentorder.StatusValidator is a validator for the "status" field. It is called by the builders before save.
|
||||
paymentorder.StatusValidator = paymentorderDescStatus.Validators[0].(func(string) error)
|
||||
// paymentorderDescRefundAmount is the schema descriptor for refund_amount field.
|
||||
paymentorderDescRefundAmount := paymentorderFields[22].Descriptor()
|
||||
paymentorderDescRefundAmount := paymentorderFields[23].Descriptor()
|
||||
// paymentorder.DefaultRefundAmount holds the default value on creation for the refund_amount field.
|
||||
paymentorder.DefaultRefundAmount = paymentorderDescRefundAmount.Default.(float64)
|
||||
// paymentorderDescForceRefund is the schema descriptor for force_refund field.
|
||||
paymentorderDescForceRefund := paymentorderFields[25].Descriptor()
|
||||
paymentorderDescForceRefund := paymentorderFields[26].Descriptor()
|
||||
// paymentorder.DefaultForceRefund holds the default value on creation for the force_refund field.
|
||||
paymentorder.DefaultForceRefund = paymentorderDescForceRefund.Default.(bool)
|
||||
// paymentorderDescRefundRequestedBy is the schema descriptor for refund_requested_by field.
|
||||
paymentorderDescRefundRequestedBy := paymentorderFields[28].Descriptor()
|
||||
paymentorderDescRefundRequestedBy := paymentorderFields[29].Descriptor()
|
||||
// paymentorder.RefundRequestedByValidator is a validator for the "refund_requested_by" field. It is called by the builders before save.
|
||||
paymentorder.RefundRequestedByValidator = paymentorderDescRefundRequestedBy.Validators[0].(func(string) error)
|
||||
// paymentorderDescClientIP is the schema descriptor for client_ip field.
|
||||
paymentorderDescClientIP := paymentorderFields[34].Descriptor()
|
||||
paymentorderDescClientIP := paymentorderFields[35].Descriptor()
|
||||
// paymentorder.ClientIPValidator is a validator for the "client_ip" field. It is called by the builders before save.
|
||||
paymentorder.ClientIPValidator = paymentorderDescClientIP.Validators[0].(func(string) error)
|
||||
// paymentorderDescSrcHost is the schema descriptor for src_host field.
|
||||
paymentorderDescSrcHost := paymentorderFields[35].Descriptor()
|
||||
paymentorderDescSrcHost := paymentorderFields[36].Descriptor()
|
||||
// paymentorder.SrcHostValidator is a validator for the "src_host" field. It is called by the builders before save.
|
||||
paymentorder.SrcHostValidator = paymentorderDescSrcHost.Validators[0].(func(string) error)
|
||||
// paymentorderDescCreatedAt is the schema descriptor for created_at field.
|
||||
paymentorderDescCreatedAt := paymentorderFields[37].Descriptor()
|
||||
paymentorderDescCreatedAt := paymentorderFields[38].Descriptor()
|
||||
// paymentorder.DefaultCreatedAt holds the default value on creation for the created_at field.
|
||||
paymentorder.DefaultCreatedAt = paymentorderDescCreatedAt.Default.(func() time.Time)
|
||||
// paymentorderDescUpdatedAt is the schema descriptor for updated_at field.
|
||||
paymentorderDescUpdatedAt := paymentorderFields[38].Descriptor()
|
||||
paymentorderDescUpdatedAt := paymentorderFields[39].Descriptor()
|
||||
// paymentorder.DefaultUpdatedAt holds the default value on creation for the updated_at field.
|
||||
paymentorder.DefaultUpdatedAt = paymentorderDescUpdatedAt.Default.(func() time.Time)
|
||||
// paymentorder.UpdateDefaultUpdatedAt holds the default value on update for the updated_at field.
|
||||
|
||||
@@ -50,6 +50,10 @@ func (PaymentOrder) Fields() []ent.Field {
|
||||
field.Float("fee_rate").
|
||||
SchemaType(map[string]string{dialect.Postgres: "decimal(10,4)"}).
|
||||
Default(0),
|
||||
// 充值赠送额度(USD)。已计入 amount;单独记录用于订单展示与推广返利基数剔除。
|
||||
field.Float("bonus_amount").
|
||||
SchemaType(map[string]string{dialect.Postgres: "decimal(20,2)"}).
|
||||
Default(0),
|
||||
field.String("recharge_code").
|
||||
MaxLen(64),
|
||||
|
||||
|
||||
@@ -42,7 +42,7 @@ func (UserPlatformQuota) Fields() []ent.Field {
|
||||
// 此处为 ent 构建期约束,需与 service.AllowedQuotaPlatforms 保持同步。
|
||||
switch s {
|
||||
case "anthropic", "openai", "gemini", "antigravity", "grok",
|
||||
"kimi", "zhipu", "deepseek", "minimax", "opencode_go":
|
||||
"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe":
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("platform %q is not allowed", s)
|
||||
|
||||
@@ -89,6 +89,7 @@ type Config struct {
|
||||
Pricing PricingConfig `mapstructure:"pricing"`
|
||||
Gateway GatewayConfig `mapstructure:"gateway"`
|
||||
APIKeyAuth APIKeyAuthCacheConfig `mapstructure:"api_key_auth_cache"`
|
||||
APIKeyCreate APIKeyCreateConfig `mapstructure:"api_key_create"`
|
||||
SubscriptionCache SubscriptionCacheConfig `mapstructure:"subscription_cache"`
|
||||
SubscriptionMaintenance SubscriptionMaintenanceConfig `mapstructure:"subscription_maintenance"`
|
||||
Dashboard DashboardCacheConfig `mapstructure:"dashboard_cache"`
|
||||
@@ -921,6 +922,29 @@ type BillingConfig struct {
|
||||
// UserPlatformQuotaSentinelTTLSeconds sentinel(无 limit 占位)entry 的 TTL,
|
||||
// 显著短于 quota cache 默认 86400s 以控 Redis 内存;默认 3600=1h。
|
||||
UserPlatformQuotaSentinelTTLSeconds int `mapstructure:"user_platform_quota_sentinel_ttl_seconds"`
|
||||
// InflightReservation 余额模式在途请求预留(Redis),防止并发请求在预检时看到同一份余额而集体透支。
|
||||
InflightReservation InflightReservationConfig `mapstructure:"inflight_reservation"`
|
||||
}
|
||||
|
||||
// InflightReservationConfig 余额模式在途预留配置。
|
||||
// 准入时按 输入估算 + 输出单价 × max_tokens 估算单请求费用,在 Redis 中原子地
|
||||
// 校验 缓存余额 - 在途预留合计 >= 估算 后登记预留,请求结束(任意路径)释放。
|
||||
// 估算失败或 Redis 不可用时 fail-open,退回旧的仅余额 > 阈值检查。
|
||||
type InflightReservationConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
// TTLSeconds 单条预留的最长存活时间;进程崩溃等泄漏的预留到期自动失效。
|
||||
TTLSeconds int `mapstructure:"ttl_seconds"`
|
||||
// DefaultMaxTokens 请求未携带 max_tokens 时用于估算的输出 token 数。
|
||||
DefaultMaxTokens int `mapstructure:"default_max_tokens"`
|
||||
// MaxOutputTokens 估算输出 token 的上限(max_tokens 超出时截断)。
|
||||
MaxOutputTokens int `mapstructure:"max_output_tokens"`
|
||||
// MaxInputTokens 输入 token 估算(请求体字节数 / 4)的上限。
|
||||
MaxInputTokens int `mapstructure:"max_input_tokens"`
|
||||
// MaxReservationUSD 单请求预留金额上限;0 表示不设上限。
|
||||
MaxReservationUSD float64 `mapstructure:"max_reservation_usd"`
|
||||
// FailClosedOnUnpriced 无法为请求估算费用(模型/分组/渠道均无定价)时是否拒绝请求。
|
||||
// 默认 false:放行且不预留(fail-open,节流告警日志)。
|
||||
FailClosedOnUnpriced bool `mapstructure:"fail_closed_on_unpriced"`
|
||||
}
|
||||
|
||||
type CircuitBreakerConfig struct {
|
||||
@@ -1696,6 +1720,14 @@ type APIKeyAuthCacheConfig struct {
|
||||
InvalidAbuse InvalidAuthAbuseConfig `mapstructure:"invalid_abuse"`
|
||||
}
|
||||
|
||||
// APIKeyCreateConfig 用户创建 API Key 的防滥用限制(0 表示不限制)
|
||||
type APIKeyCreateConfig struct {
|
||||
// MaxActivePerUser 单个用户同时存在(未删除)的 API Key 上限
|
||||
MaxActivePerUser int `mapstructure:"max_active_per_user"`
|
||||
// MaxPerUserPerHour 单个用户每小时可创建的 API Key 次数(删除不返还次数)
|
||||
MaxPerUserPerHour int `mapstructure:"max_per_user_per_hour"`
|
||||
}
|
||||
|
||||
type InvalidAuthAbuseConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
Threshold int `mapstructure:"threshold"`
|
||||
@@ -2093,6 +2125,13 @@ func setDefaults() {
|
||||
viper.SetDefault("billing.minimum_balance_reserve", 0.000001)
|
||||
viper.SetDefault("billing.user_platform_quota_cache_ttl_seconds", 86400)
|
||||
viper.SetDefault("billing.user_platform_quota_sentinel_ttl_seconds", 3600)
|
||||
viper.SetDefault("billing.inflight_reservation.enabled", true)
|
||||
viper.SetDefault("billing.inflight_reservation.ttl_seconds", 900)
|
||||
viper.SetDefault("billing.inflight_reservation.default_max_tokens", 8192)
|
||||
viper.SetDefault("billing.inflight_reservation.max_output_tokens", 128000)
|
||||
viper.SetDefault("billing.inflight_reservation.max_input_tokens", 200000)
|
||||
viper.SetDefault("billing.inflight_reservation.max_reservation_usd", 0)
|
||||
viper.SetDefault("billing.inflight_reservation.fail_closed_on_unpriced", false)
|
||||
|
||||
// Turnstile
|
||||
viper.SetDefault("turnstile.required", false)
|
||||
@@ -2332,6 +2371,8 @@ func setDefaults() {
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.window_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.block_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.capacity", 16384)
|
||||
viper.SetDefault("api_key_create.max_active_per_user", 200)
|
||||
viper.SetDefault("api_key_create.max_per_user_per_hour", 60)
|
||||
|
||||
// Subscription auth L1 cache
|
||||
viper.SetDefault("subscription_cache.l1_size", 16384)
|
||||
@@ -2711,6 +2752,12 @@ func (c *Config) Validate() error {
|
||||
return fmt.Errorf("server.h2c.max_upload_buffer_per_stream must be positive")
|
||||
}
|
||||
}
|
||||
if c.APIKeyCreate.MaxActivePerUser < 0 {
|
||||
return fmt.Errorf("api_key_create.max_active_per_user must be non-negative")
|
||||
}
|
||||
if c.APIKeyCreate.MaxPerUserPerHour < 0 {
|
||||
return fmt.Errorf("api_key_create.max_per_user_per_hour must be non-negative")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.Enabled {
|
||||
if c.APIKeyAuth.InvalidAbuse.Threshold < 10 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.threshold must be at least 10")
|
||||
@@ -3062,6 +3109,11 @@ func (c *Config) Validate() error {
|
||||
if c.Billing.MinimumBalanceReserve < 0 {
|
||||
return fmt.Errorf("billing.minimum_balance_reserve must be non-negative")
|
||||
}
|
||||
if c.Billing.InflightReservation.TTLSeconds < 0 || c.Billing.InflightReservation.DefaultMaxTokens < 0 ||
|
||||
c.Billing.InflightReservation.MaxOutputTokens < 0 || c.Billing.InflightReservation.MaxInputTokens < 0 ||
|
||||
c.Billing.InflightReservation.MaxReservationUSD < 0 {
|
||||
return fmt.Errorf("billing.inflight_reservation values must be non-negative")
|
||||
}
|
||||
if c.Database.MaxOpenConns <= 0 {
|
||||
return fmt.Errorf("database.max_open_conns must be positive")
|
||||
}
|
||||
|
||||
@@ -28,6 +28,7 @@ const (
|
||||
PlatformZhipu = "zhipu" // 智谱 GLM (bigmodel)
|
||||
PlatformDeepseek = "deepseek" // DeepSeek
|
||||
PlatformMiniMax = "minimax" // MiniMax (M 系列)
|
||||
PlatformTypeSafe = "typesafe" // TypeSafe AI System One (Jev)
|
||||
// PlatformOpenCodeGo 是 OpenCode 平台(账号类型 Zen 按量 / Go 订阅)。
|
||||
// 值保持 opencode_go 以兼容已落库的分组、配额与 Composite 路由 CHECK。
|
||||
PlatformOpenCodeGo = "opencode_go"
|
||||
@@ -190,6 +191,8 @@ var DefaultBedrockModelMapping = map[string]string{
|
||||
"claude-opus-4-1": "us.anthropic.claude-opus-4-1-20250805-v1:0",
|
||||
"claude-opus-4-20250514": "us.anthropic.claude-opus-4-20250514-v1:0",
|
||||
// Claude Sonnet
|
||||
// Sonnet 5.5 is available on bedrock-runtime through Global inference only.
|
||||
"claude-sonnet-5-5": "global.anthropic.claude-sonnet-5-5",
|
||||
"claude-sonnet-5": "us.anthropic.claude-sonnet-5-v1",
|
||||
"claude-sonnet-4-6-thinking": "us.anthropic.claude-sonnet-4-6",
|
||||
"claude-sonnet-4-6": "us.anthropic.claude-sonnet-4-6",
|
||||
|
||||
@@ -109,9 +109,10 @@ func TestDefaultBedrockModelMapping_ContainsNewClaudeModels(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := map[string]string{
|
||||
"claude-fable-5-1": "anthropic.claude-fable-5-1",
|
||||
"claude-fable-5": "anthropic.claude-fable-5",
|
||||
"claude-opus-4-8": "us.anthropic.claude-opus-4-8-v1",
|
||||
"claude-fable-5-1": "anthropic.claude-fable-5-1",
|
||||
"claude-fable-5": "anthropic.claude-fable-5",
|
||||
"claude-opus-4-8": "us.anthropic.claude-opus-4-8-v1",
|
||||
"claude-sonnet-5-5": "global.anthropic.claude-sonnet-5-5",
|
||||
}
|
||||
for from, want := range cases {
|
||||
got, ok := DefaultBedrockModelMapping[from]
|
||||
|
||||
@@ -27,6 +27,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/typesafe"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
@@ -48,6 +49,7 @@ func NewOAuthHandler(oauthService *service.OAuthService) *OAuthHandler {
|
||||
|
||||
// AccountHandler handles admin account management
|
||||
type AccountHandler struct {
|
||||
claudeResetCredits claudeResetReader
|
||||
adminService service.AdminService
|
||||
oauthService *service.OAuthService
|
||||
openaiOAuthService *service.OpenAIOAuthService
|
||||
@@ -2949,6 +2951,12 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// TypeSafe accounts serve only the native System One model.
|
||||
if account.IsTypeSafe() {
|
||||
response.Success(c, []claude.Model{{ID: typesafe.JevLatestModel, Type: "model", DisplayName: typesafe.JevLatestModel}})
|
||||
return
|
||||
}
|
||||
|
||||
// Handle Claude/Anthropic accounts
|
||||
// For OAuth and Setup-Token accounts: return default models
|
||||
if account.IsOAuth() {
|
||||
|
||||
@@ -526,3 +526,36 @@ func TestAccountHandlerSyncUpstreamModels_MetadataEnrichmentFailureReturnsWarnin
|
||||
require.Len(t, resp.Data.Warnings, 1)
|
||||
require.Equal(t, "upstream_model_metadata_incomplete", resp.Data.Warnings[0].Code)
|
||||
}
|
||||
|
||||
func TestAccountHandlerGetAvailableModels_TypeSafeOnlyReturnsJev(t *testing.T) {
|
||||
for _, credentials := range []map[string]any{
|
||||
{"api_key": "ts-secret"},
|
||||
{"api_key": "ts-secret", "model_mapping": map[string]any{"jev-latest": "jev-latest"}},
|
||||
} {
|
||||
svc := &availableModelsAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
account: service.Account{
|
||||
ID: 46,
|
||||
Name: "typesafe",
|
||||
Platform: service.PlatformTypeSafe,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive,
|
||||
Credentials: credentials,
|
||||
},
|
||||
}
|
||||
router := setupAvailableModelsRouter(svc)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/46/models", nil))
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var resp struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Len(t, resp.Data, 1)
|
||||
require.Equal(t, "jev-latest", resp.Data[0].ID)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -644,6 +644,7 @@ var platformToLiteLLMProvider = map[string]string{
|
||||
service.PlatformDeepseek: "deepseek",
|
||||
service.PlatformMiniMax: "minimax",
|
||||
service.PlatformOpenCodeGo: "opencode-go",
|
||||
service.PlatformTypeSafe: "typesafe",
|
||||
}
|
||||
|
||||
// SyncPricingModels 返回 LiteLLM 定价目录中指定平台的最新模型列表
|
||||
|
||||
@@ -565,7 +565,7 @@ func TestSyncPricingModels_ValidPlatform_EmptyService(t *testing.T) {
|
||||
svc := service.NewPricingService(nil, nil)
|
||||
router := setupSyncPricingModelsRouter(svc)
|
||||
|
||||
for _, platform := range []string{"anthropic", "openai", "gemini", "antigravity", "grok", "kimi", "zhipu", "deepseek", "minimax"} {
|
||||
for _, platform := range []string{"anthropic", "openai", "gemini", "antigravity", "grok", "kimi", "zhipu", "deepseek", "minimax", "typesafe"} {
|
||||
req := httptest.NewRequest(http.MethodGet, "/channels/pricing/sync-models?platform="+platform, nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
type claudeResetReader interface {
|
||||
Query(context.Context, int64) (*service.ClaudeResetCredits, error)
|
||||
Redeem(context.Context, int64, string) (*service.ClaudeResetOutcome, error)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) SetClaudeResetCreditService(s *service.ClaudeResetCreditService) {
|
||||
h.claudeResetCredits = s
|
||||
}
|
||||
func (h *AccountHandler) ClaudeResetCredits(c *gin.Context) {
|
||||
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h.claudeResetCredits == nil {
|
||||
response.Error(c, 503, "Claude reset service unavailable")
|
||||
return
|
||||
}
|
||||
status, err := h.claudeResetCredits.Query(c.Request.Context(), id)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, status)
|
||||
}
|
||||
|
||||
// RedeemClaudeResetCredit consumes the upstream next reset credit. The request has
|
||||
// no body: the server picks the grant; the Idempotency-Key header identifies one
|
||||
// operator confirmation and replays its outcome.
|
||||
func (h *AccountHandler) RedeemClaudeResetCredit(c *gin.Context) {
|
||||
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h.claudeResetCredits == nil {
|
||||
response.Error(c, 503, "Claude reset service unavailable")
|
||||
return
|
||||
}
|
||||
result, err := h.claudeResetCredits.Redeem(c.Request.Context(), id, c.GetHeader("Idempotency-Key"))
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type claudeResetHandlerStub struct {
|
||||
key string
|
||||
calls int
|
||||
}
|
||||
|
||||
func (s *claudeResetHandlerStub) Query(context.Context, int64) (*service.ClaudeResetCredits, error) {
|
||||
return &service.ClaudeResetCredits{Credits: []service.ClaudeResetCredit{}}, nil
|
||||
}
|
||||
|
||||
func (s *claudeResetHandlerStub) Redeem(_ context.Context, _ int64, key string) (*service.ClaudeResetOutcome, error) {
|
||||
s.calls++
|
||||
s.key = key
|
||||
if key == "" {
|
||||
return nil, service.ErrIdempotencyKeyRequired
|
||||
}
|
||||
return &service.ClaudeResetOutcome{Outcome: "reset", Cleared: []string{"five_hour"}}, nil
|
||||
}
|
||||
|
||||
func TestClaudeResetHandlerRedeemContract(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
stub := &claudeResetHandlerStub{}
|
||||
h := &AccountHandler{claudeResetCredits: stub}
|
||||
r := gin.New()
|
||||
r.POST("/accounts/:id/claude/reset-credits/redeem", h.RedeemClaudeResetCredit)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/accounts/1/claude/reset-credits/redeem", nil)
|
||||
req.Header.Set("Idempotency-Key", "same-operation")
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, "same-operation", stub.key)
|
||||
require.Contains(t, w.Body.String(), `"outcome":"reset"`)
|
||||
|
||||
w = httptest.NewRecorder()
|
||||
r.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/accounts/1/claude/reset-credits/redeem", nil))
|
||||
require.Equal(t, http.StatusBadRequest, w.Code)
|
||||
require.Contains(t, w.Body.String(), "IDEMPOTENCY_KEY_REQUIRED")
|
||||
|
||||
w = httptest.NewRecorder()
|
||||
r.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/accounts/abc/claude/reset-credits/redeem", nil))
|
||||
require.Equal(t, http.StatusBadRequest, w.Code)
|
||||
require.Equal(t, 2, stub.calls)
|
||||
}
|
||||
@@ -500,7 +500,11 @@ func (h *DashboardHandler) GetUserUsageTrend(c *gin.Context) {
|
||||
limit = 12
|
||||
}
|
||||
|
||||
trend, hit, err := h.getUserUsageTrendCached(c.Request.Context(), startTime, endTime, granularity, limit)
|
||||
metric := c.DefaultQuery("metric", "tokens")
|
||||
if metric != "tokens" && metric != "actual_cost" {
|
||||
metric = "tokens"
|
||||
}
|
||||
trend, hit, err := h.getUserUsageTrendCached(c.Request.Context(), startTime, endTime, granularity, limit, metric)
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get user usage trend")
|
||||
return
|
||||
|
||||
@@ -45,6 +45,7 @@ func (r *dashboardUsageRepoCacheProbe) GetUserUsageTrend(
|
||||
startTime, endTime time.Time,
|
||||
granularity string,
|
||||
limit int,
|
||||
metric string,
|
||||
) ([]usagestats.UserUsageTrendPoint, error) {
|
||||
r.usersTrendCalls.Add(1)
|
||||
return []usagestats.UserUsageTrendPoint{{
|
||||
@@ -115,4 +116,20 @@ func TestDashboardHandler_GetUserUsageTrend_UsesCache(t *testing.T) {
|
||||
require.Equal(t, http.StatusOK, rec2.Code)
|
||||
require.Equal(t, "hit", rec2.Header().Get("X-Snapshot-Cache"))
|
||||
require.Equal(t, int32(1), repo.usersTrendCalls.Load())
|
||||
|
||||
for _, tc := range []struct {
|
||||
metric, cache string
|
||||
calls int32
|
||||
}{
|
||||
{"actual_cost", "miss", 2},
|
||||
{"actual_cost", "hit", 2},
|
||||
{"tokens", "hit", 2},
|
||||
} {
|
||||
req := httptest.NewRequest(http.MethodGet, "/admin/dashboard/users-trend?start_date=2026-03-01&end_date=2026-03-07&granularity=day&limit=8&metric="+tc.metric, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, tc.cache, rec.Header().Get("X-Snapshot-Cache"))
|
||||
require.Equal(t, tc.calls, repo.usersTrendCalls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -53,6 +53,7 @@ type dashboardEntityTrendCacheKey struct {
|
||||
EndTime string `json:"end_time"`
|
||||
Granularity string `json:"granularity"`
|
||||
Limit int `json:"limit"`
|
||||
Metric string `json:"metric,omitempty"`
|
||||
}
|
||||
|
||||
func cacheStatusValue(hit bool) string {
|
||||
@@ -213,15 +214,16 @@ func (h *DashboardHandler) getAPIKeyUsageTrendCached(ctx context.Context, startT
|
||||
return trend, hit, err
|
||||
}
|
||||
|
||||
func (h *DashboardHandler) getUserUsageTrendCached(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.UserUsageTrendPoint, bool, error) {
|
||||
func (h *DashboardHandler) getUserUsageTrendCached(ctx context.Context, startTime, endTime time.Time, granularity string, limit int, metric string) ([]usagestats.UserUsageTrendPoint, bool, error) {
|
||||
key := mustMarshalDashboardCacheKey(dashboardEntityTrendCacheKey{
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
Granularity: granularity,
|
||||
Limit: limit,
|
||||
Metric: metric,
|
||||
})
|
||||
entry, hit, err := dashboardUsersTrendCache.GetOrLoad(key, func() (any, error) {
|
||||
return h.dashboardService.GetUserUsageTrend(ctx, startTime, endTime, granularity, limit)
|
||||
return h.dashboardService.GetUserUsageTrend(ctx, startTime, endTime, granularity, limit, metric)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, hit, err
|
||||
|
||||
@@ -243,7 +243,7 @@ func (h *DashboardHandler) buildSnapshotV2Response(
|
||||
}
|
||||
|
||||
if includeUsersTrend {
|
||||
usersTrend, _, err := h.getUserUsageTrendCached(ctx, startTime, endTime, granularity, usersTrendLimit)
|
||||
usersTrend, _, err := h.getUserUsageTrendCached(ctx, startTime, endTime, granularity, usersTrendLimit, "tokens")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to get user usage trend")
|
||||
}
|
||||
|
||||
@@ -184,7 +184,7 @@ func sanitizeUpdateGroupRequestForSimpleMode(req *UpdateGroupRequest) {
|
||||
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 composite"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go typesafe composite"`
|
||||
RateMultiplier float64 `json:"rate_multiplier"`
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
|
||||
@@ -258,7 +258,7 @@ 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 composite"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go typesafe composite"`
|
||||
RateMultiplier *float64 `json:"rate_multiplier"`
|
||||
IsExclusive *bool `json:"is_exclusive"`
|
||||
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
|
||||
@@ -332,7 +332,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"`
|
||||
TargetPlatform string `json:"target_platform" binding:"required,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go typesafe"`
|
||||
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"`
|
||||
|
||||
@@ -28,6 +28,7 @@ func TestGroupPlatformBinding_AllowedPlatforms(t *testing.T) {
|
||||
allowed := []string{
|
||||
"anthropic", "openai", "gemini", "antigravity", "grok",
|
||||
"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "composite",
|
||||
"typesafe",
|
||||
}
|
||||
for _, platform := range allowed {
|
||||
t.Run("create_"+platform, func(t *testing.T) {
|
||||
@@ -71,8 +72,8 @@ func TestGroupPlatformBinding_RejectsInvalidPlatforms(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompositeRouteTargetPlatform_AllowsCNProviders(t *testing.T) {
|
||||
for _, platform := range []string{"kimi", "zhipu", "deepseek", "minimax", "opencode_go"} {
|
||||
func TestCompositeRouteTargetPlatform_AllowsConcreteProviders(t *testing.T) {
|
||||
for _, platform := range []string{"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe"} {
|
||||
var req CompositeRouteRequest
|
||||
body := fmt.Sprintf(`{"public_model":"m","target_platform":%q}`, platform)
|
||||
require.NoError(t, bindGroupPlatformJSON(t, &req, body))
|
||||
|
||||
@@ -125,6 +125,7 @@ type AdminPaymentOrderResult struct {
|
||||
Amount float64 `json:"amount"`
|
||||
PayAmount float64 `json:"pay_amount"`
|
||||
FeeRate float64 `json:"fee_rate"`
|
||||
BonusAmount float64 `json:"bonus_amount"`
|
||||
Currency string `json:"currency"`
|
||||
RechargeCode string `json:"recharge_code,omitempty"`
|
||||
OutTradeNo string `json:"out_trade_no"`
|
||||
@@ -182,6 +183,7 @@ func sanitizeAdminPaymentOrderForResponse(order *dbent.PaymentOrder) *AdminPayme
|
||||
Amount: order.Amount,
|
||||
PayAmount: order.PayAmount,
|
||||
FeeRate: order.FeeRate,
|
||||
BonusAmount: order.BonusAmount,
|
||||
Currency: service.PaymentOrderCurrency(order),
|
||||
RechargeCode: order.RechargeCode,
|
||||
OutTradeNo: order.OutTradeNo,
|
||||
|
||||
@@ -265,6 +265,7 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
DefaultBalance: settings.DefaultBalance,
|
||||
RiskControlEnabled: settings.RiskControlEnabled,
|
||||
CyberSessionBlockEnabled: settings.CyberSessionBlockEnabled,
|
||||
CyberPolicyUserAllowlist: settings.CyberPolicyUserAllowlist,
|
||||
CyberSessionBlockTTLSeconds: settings.CyberSessionBlockTTLSeconds,
|
||||
AffiliateRebateRate: settings.AffiliateRebateRate,
|
||||
AffiliateRebateFreezeHours: settings.AffiliateRebateFreezeHours,
|
||||
@@ -361,6 +362,9 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
PaymentBalanceRechargeMultiplier: paymentCfg.BalanceRechargeMultiplier,
|
||||
PaymentSubscriptionUSDToCNYRate: paymentCfg.SubscriptionUSDToCNYRate,
|
||||
PaymentRechargeFeeRate: paymentCfg.RechargeFeeRate,
|
||||
PaymentRechargeBonusTiers: rechargeBonusTiersToDTO(paymentCfg.RechargeBonusTiers),
|
||||
PaymentRechargeBonusMode: rechargeBonusModeToDTO(paymentCfg.RechargeBonusMode),
|
||||
PaymentRechargeBonusNotice: paymentCfg.RechargeBonusNotice,
|
||||
PaymentLoadBalanceStrat: paymentCfg.LoadBalanceStrategy,
|
||||
PaymentProductNamePrefix: paymentCfg.ProductNamePrefix,
|
||||
PaymentProductNameSuffix: paymentCfg.ProductNameSuffix,
|
||||
|
||||
@@ -603,6 +603,9 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.RiskControlEnabled != after.RiskControlEnabled {
|
||||
changed = append(changed, "risk_control_enabled")
|
||||
}
|
||||
if before.CyberPolicyUserAllowlist != after.CyberPolicyUserAllowlist {
|
||||
changed = append(changed, "cyber_policy_user_allowlist")
|
||||
}
|
||||
if before.CyberSessionBlockEnabled != after.CyberSessionBlockEnabled {
|
||||
changed = append(changed, "cyber_session_block_enabled")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
// rechargeBonusTiersFromDTO 请求 nil 表示未携带该字段(保持现值);空数组表示清空阶梯。
|
||||
func rechargeBonusTiersFromDTO(items *[]dto.RechargeBonusTier) *[]service.RechargeBonusTier {
|
||||
if items == nil {
|
||||
return nil
|
||||
}
|
||||
out := make([]service.RechargeBonusTier, 0, len(*items))
|
||||
for _, item := range *items {
|
||||
out = append(out, service.RechargeBonusTier{MinAmount: item.MinAmount, BonusPercent: item.BonusPercent})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// rechargeBonusModeToDTO 输出已归一化的模式(空/非法按 bonus)。
|
||||
func rechargeBonusModeToDTO(mode string) string {
|
||||
normalized, _ := service.NormalizeRechargeBonusMode(mode)
|
||||
return normalized
|
||||
}
|
||||
|
||||
// rechargeBonusTiersToDTO 始终返回非 nil 切片,空配置输出 []。
|
||||
func rechargeBonusTiersToDTO(items []service.RechargeBonusTier) []dto.RechargeBonusTier {
|
||||
out := make([]dto.RechargeBonusTier, 0, len(items))
|
||||
for _, item := range items {
|
||||
out = append(out, dto.RechargeBonusTier{MinAmount: item.MinAmount, BonusPercent: item.BonusPercent})
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -312,11 +312,15 @@ type UpdateSettingsRequest struct {
|
||||
PaymentBalanceRechargeMultiplier *float64 `json:"payment_balance_recharge_multiplier"`
|
||||
PaymentSubscriptionUSDToCNYRate *float64 `json:"payment_subscription_usd_to_cny_rate"`
|
||||
PaymentRechargeFeeRate *float64 `json:"payment_recharge_fee_rate"`
|
||||
PaymentLoadBalanceStrat *string `json:"payment_load_balance_strategy"`
|
||||
PaymentProductNamePrefix *string `json:"payment_product_name_prefix"`
|
||||
PaymentProductNameSuffix *string `json:"payment_product_name_suffix"`
|
||||
PaymentHelpImageURL *string `json:"payment_help_image_url"`
|
||||
PaymentHelpText *string `json:"payment_help_text"`
|
||||
// nil 表示不更新;空数组表示清空阶梯
|
||||
PaymentRechargeBonusTiers *[]dto.RechargeBonusTier `json:"payment_recharge_bonus_tiers"`
|
||||
PaymentRechargeBonusMode *string `json:"payment_recharge_bonus_mode"`
|
||||
PaymentRechargeBonusNotice *string `json:"payment_recharge_bonus_notice"`
|
||||
PaymentLoadBalanceStrat *string `json:"payment_load_balance_strategy"`
|
||||
PaymentProductNamePrefix *string `json:"payment_product_name_prefix"`
|
||||
PaymentProductNameSuffix *string `json:"payment_product_name_suffix"`
|
||||
PaymentHelpImageURL *string `json:"payment_help_image_url"`
|
||||
PaymentHelpText *string `json:"payment_help_text"`
|
||||
|
||||
// Cancel rate limit
|
||||
PaymentCancelRateLimitEnabled *bool `json:"payment_cancel_rate_limit_enabled"`
|
||||
@@ -364,8 +368,9 @@ type UpdateSettingsRequest struct {
|
||||
RiskControlEnabled *bool `json:"risk_control_enabled"`
|
||||
|
||||
// cyber 会话屏蔽开关 + TTL
|
||||
CyberSessionBlockEnabled *bool `json:"cyber_session_block_enabled"`
|
||||
CyberSessionBlockTTLSeconds *int `json:"cyber_session_block_ttl_seconds"`
|
||||
CyberSessionBlockEnabled *bool `json:"cyber_session_block_enabled"`
|
||||
CyberPolicyUserAllowlist *string `json:"cyber_policy_user_allowlist"`
|
||||
CyberSessionBlockTTLSeconds *int `json:"cyber_session_block_ttl_seconds"`
|
||||
|
||||
// OpenAI fast/flex policy (optional, only updated when provided)
|
||||
OpenAIFastPolicySettings *dto.OpenAIFastPolicySettings `json:"openai_fast_policy_settings,omitempty"`
|
||||
@@ -1503,6 +1508,13 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
if req.CyberPolicyUserAllowlist != nil {
|
||||
if _, err := service.ParseCyberPolicyUserAllowlist(*req.CyberPolicyUserAllowlist); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// cyber 会话屏蔽 TTL 校验:提供时必须 > 0
|
||||
if req.CyberSessionBlockTTLSeconds != nil && *req.CyberSessionBlockTTLSeconds <= 0 {
|
||||
response.BadRequest(c, "cyber_session_block_ttl_seconds must be > 0")
|
||||
@@ -2008,6 +2020,12 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
return previousSettings.RiskControlEnabled
|
||||
}(),
|
||||
CyberPolicyUserAllowlist: func() string {
|
||||
if req.CyberPolicyUserAllowlist != nil {
|
||||
return *req.CyberPolicyUserAllowlist
|
||||
}
|
||||
return previousSettings.CyberPolicyUserAllowlist
|
||||
}(),
|
||||
CyberSessionBlockEnabled: func() bool {
|
||||
if req.CyberSessionBlockEnabled != nil {
|
||||
return *req.CyberSessionBlockEnabled
|
||||
@@ -2114,6 +2132,9 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
BalanceRechargeMultiplier: req.PaymentBalanceRechargeMultiplier,
|
||||
SubscriptionUSDToCNYRate: req.PaymentSubscriptionUSDToCNYRate,
|
||||
RechargeFeeRate: req.PaymentRechargeFeeRate,
|
||||
RechargeBonusTiers: rechargeBonusTiersFromDTO(req.PaymentRechargeBonusTiers),
|
||||
RechargeBonusMode: req.PaymentRechargeBonusMode,
|
||||
RechargeBonusNotice: req.PaymentRechargeBonusNotice,
|
||||
LoadBalanceStrategy: req.PaymentLoadBalanceStrat,
|
||||
ProductNamePrefix: req.PaymentProductNamePrefix,
|
||||
ProductNameSuffix: req.PaymentProductNameSuffix,
|
||||
@@ -2393,6 +2414,9 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
PaymentBalanceRechargeMultiplier: updatedPaymentCfg.BalanceRechargeMultiplier,
|
||||
PaymentSubscriptionUSDToCNYRate: updatedPaymentCfg.SubscriptionUSDToCNYRate,
|
||||
PaymentRechargeFeeRate: updatedPaymentCfg.RechargeFeeRate,
|
||||
PaymentRechargeBonusTiers: rechargeBonusTiersToDTO(updatedPaymentCfg.RechargeBonusTiers),
|
||||
PaymentRechargeBonusMode: rechargeBonusModeToDTO(updatedPaymentCfg.RechargeBonusMode),
|
||||
PaymentRechargeBonusNotice: updatedPaymentCfg.RechargeBonusNotice,
|
||||
PaymentLoadBalanceStrat: updatedPaymentCfg.LoadBalanceStrategy,
|
||||
PaymentProductNamePrefix: updatedPaymentCfg.ProductNamePrefix,
|
||||
PaymentProductNameSuffix: updatedPaymentCfg.ProductNameSuffix,
|
||||
@@ -2429,6 +2453,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
|
||||
RiskControlEnabled: updatedSettings.RiskControlEnabled,
|
||||
CyberSessionBlockEnabled: updatedSettings.CyberSessionBlockEnabled,
|
||||
CyberPolicyUserAllowlist: updatedSettings.CyberPolicyUserAllowlist,
|
||||
CyberSessionBlockTTLSeconds: updatedSettings.CyberSessionBlockTTLSeconds,
|
||||
AccountSchedulingThresholds: updatedSettings.AccountSchedulingThresholds,
|
||||
AllowUserViewErrorRequests: updatedSettings.AllowUserViewErrorRequests,
|
||||
@@ -2468,6 +2493,7 @@ func hasPaymentFields(req UpdateSettingsRequest) bool {
|
||||
req.PaymentEnabledTypes != nil || req.PaymentBalanceDisabled != nil ||
|
||||
req.PaymentBalanceRechargeMultiplier != nil || req.PaymentSubscriptionUSDToCNYRate != nil ||
|
||||
req.PaymentRechargeFeeRate != nil ||
|
||||
req.PaymentRechargeBonusTiers != nil || req.PaymentRechargeBonusMode != nil || req.PaymentRechargeBonusNotice != nil ||
|
||||
req.PaymentLoadBalanceStrat != nil || req.PaymentProductNamePrefix != nil ||
|
||||
req.PaymentProductNameSuffix != nil || req.PaymentHelpImageURL != nil ||
|
||||
req.PaymentHelpText != nil || req.PaymentCancelRateLimitEnabled != nil ||
|
||||
|
||||
@@ -3593,3 +3593,20 @@ func (oauthPendingFlowTotpEncryptorStub) Encrypt(plaintext string) (string, erro
|
||||
func (oauthPendingFlowTotpEncryptorStub) Decrypt(ciphertext string) (string, error) {
|
||||
return ciphertext, nil
|
||||
}
|
||||
|
||||
func (s *oauthPendingFlowEmailCacheStub) IncrVerificationCodeAttempts(_ context.Context, email string) (int, error) {
|
||||
data := s.verificationCodes[email]
|
||||
if data == nil {
|
||||
return 0, errors.New("verification code not found")
|
||||
}
|
||||
data.Attempts++
|
||||
return data.Attempts, nil
|
||||
}
|
||||
|
||||
func (s *oauthPendingFlowEmailCacheStub) IncrNotifyVerifyCodeAttempts(context.Context, string) (int, error) {
|
||||
return 0, errors.New("notify verification code not found")
|
||||
}
|
||||
|
||||
func (s *oauthPendingFlowEmailCacheStub) ConsumePasswordResetToken(context.Context, string, string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
package dto
|
||||
|
||||
// RechargeBonusTier 充值赠送档位:余额充值支付金额 ≥ MinAmount 时,在到账基数上赠送 BonusPercent%。
|
||||
type RechargeBonusTier struct {
|
||||
MinAmount float64 `json:"min_amount"`
|
||||
BonusPercent float64 `json:"bonus_percent"`
|
||||
}
|
||||
@@ -279,11 +279,15 @@ type SystemSettings struct {
|
||||
PaymentBalanceRechargeMultiplier float64 `json:"payment_balance_recharge_multiplier"`
|
||||
PaymentSubscriptionUSDToCNYRate float64 `json:"payment_subscription_usd_to_cny_rate"`
|
||||
PaymentRechargeFeeRate float64 `json:"payment_recharge_fee_rate"`
|
||||
PaymentLoadBalanceStrat string `json:"payment_load_balance_strategy"`
|
||||
PaymentProductNamePrefix string `json:"payment_product_name_prefix"`
|
||||
PaymentProductNameSuffix string `json:"payment_product_name_suffix"`
|
||||
PaymentHelpImageURL string `json:"payment_help_image_url"`
|
||||
PaymentHelpText string `json:"payment_help_text"`
|
||||
// 充值赠送阶梯与活动文案
|
||||
PaymentRechargeBonusTiers []RechargeBonusTier `json:"payment_recharge_bonus_tiers"`
|
||||
PaymentRechargeBonusMode string `json:"payment_recharge_bonus_mode"`
|
||||
PaymentRechargeBonusNotice string `json:"payment_recharge_bonus_notice"`
|
||||
PaymentLoadBalanceStrat string `json:"payment_load_balance_strategy"`
|
||||
PaymentProductNamePrefix string `json:"payment_product_name_prefix"`
|
||||
PaymentProductNameSuffix string `json:"payment_product_name_suffix"`
|
||||
PaymentHelpImageURL string `json:"payment_help_image_url"`
|
||||
PaymentHelpText string `json:"payment_help_text"`
|
||||
|
||||
// Cancel rate limit
|
||||
PaymentCancelRateLimitEnabled bool `json:"payment_cancel_rate_limit_enabled"`
|
||||
@@ -335,8 +339,9 @@ type SystemSettings struct {
|
||||
RiskControlEnabled bool `json:"risk_control_enabled"`
|
||||
|
||||
// cyber 会话屏蔽开关 + TTL
|
||||
CyberSessionBlockEnabled bool `json:"cyber_session_block_enabled"`
|
||||
CyberSessionBlockTTLSeconds int `json:"cyber_session_block_ttl_seconds"`
|
||||
CyberSessionBlockEnabled bool `json:"cyber_session_block_enabled"`
|
||||
CyberPolicyUserAllowlist string `json:"cyber_policy_user_allowlist"`
|
||||
CyberSessionBlockTTLSeconds int `json:"cyber_session_block_ttl_seconds"`
|
||||
|
||||
// Affiliate (邀请返利) feature switch
|
||||
AffiliateEnabled bool `json:"affiliate_enabled"`
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
|
||||
const (
|
||||
EndpointMessages = "/v1/messages"
|
||||
EndpointSystemOne = "/v1/systemone"
|
||||
EndpointChatCompletions = "/v1/chat/completions"
|
||||
EndpointEmbeddings = "/v1/embeddings"
|
||||
EndpointAlphaSearch = "/v1/alpha/search"
|
||||
@@ -94,6 +95,8 @@ func NormalizeInboundEndpoint(path string) string {
|
||||
return EndpointChatCompletions
|
||||
case strings.Contains(path, EndpointMessages):
|
||||
return EndpointMessages
|
||||
case strings.Contains(path, EndpointSystemOne):
|
||||
return EndpointSystemOne
|
||||
case strings.Contains(path, EndpointImagesGenerations) || strings.Contains(path, "/images/generations"):
|
||||
return EndpointImagesGenerations
|
||||
case strings.Contains(path, EndpointImagesEdits) || strings.Contains(path, "/images/edits"):
|
||||
@@ -223,6 +226,9 @@ func DeriveUpstreamEndpoint(inbound, rawRequestPath, platform string) string {
|
||||
case service.PlatformAnthropic:
|
||||
return EndpointMessages
|
||||
|
||||
case service.PlatformTypeSafe:
|
||||
return EndpointSystemOne
|
||||
|
||||
case service.PlatformGemini:
|
||||
return EndpointGeminiModels
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ func TestNormalizeInboundEndpoint(t *testing.T) {
|
||||
}{
|
||||
// Direct canonical paths.
|
||||
{"/v1/messages", EndpointMessages},
|
||||
{"/v1/systemone", EndpointSystemOne},
|
||||
{"/v1/chat/completions", EndpointChatCompletions},
|
||||
{"/v1/embeddings", EndpointEmbeddings},
|
||||
{"/v1/alpha/search", EndpointAlphaSearch},
|
||||
@@ -97,6 +98,7 @@ func TestDeriveUpstreamEndpoint(t *testing.T) {
|
||||
}{
|
||||
// Anthropic.
|
||||
{"anthropic messages", EndpointMessages, "/v1/messages", service.PlatformAnthropic, EndpointMessages},
|
||||
{"typesafe system one", EndpointSystemOne, "/v1/systemone", service.PlatformTypeSafe, EndpointSystemOne},
|
||||
|
||||
// Gemini.
|
||||
{"gemini models", EndpointGeminiModels, "/v1beta/models/gemini:gen", service.PlatformGemini, EndpointGeminiModels},
|
||||
|
||||
@@ -24,6 +24,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/typesafe"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
@@ -222,6 +223,9 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||||
return
|
||||
}
|
||||
if rejectSystemOneOnlyPlatform(c, apiKey, h.errorResponse) {
|
||||
return
|
||||
}
|
||||
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.anthropicSecurityAuditError(c, decision)
|
||||
@@ -263,6 +267,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gateway.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// 设置请求所属分组 ID(用于渠道级功能判断,如 WebSearch 模拟)
|
||||
parsedReq.GroupID = apiKey.GroupID
|
||||
|
||||
@@ -1149,7 +1166,7 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
}
|
||||
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID)
|
||||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID, true)
|
||||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
if len(source) == 0 {
|
||||
@@ -1193,6 +1210,10 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
writeGrokModelsList(c, xai.DefaultModelIDs())
|
||||
return
|
||||
}
|
||||
if platform == service.PlatformTypeSafe {
|
||||
writeModelsList(c, platform, []string{typesafe.JevLatestModel})
|
||||
return
|
||||
}
|
||||
|
||||
writeModelsListResponse(c, claude.DefaultModels)
|
||||
}
|
||||
@@ -1244,7 +1265,7 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
platform = group.Platform
|
||||
}
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(ctx, groupID)
|
||||
availableModels := h.compositeAvailableModels(ctx, groupID, false)
|
||||
fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite)
|
||||
if group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
@@ -1270,14 +1291,20 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
return fallbackModels
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64) []string {
|
||||
// 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 {
|
||||
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} {
|
||||
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} {
|
||||
if platform == service.PlatformTypeSafe && !includeSystemOne {
|
||||
continue
|
||||
}
|
||||
platformModels := h.gatewayService.GetAvailableModels(ctx, groupID, platform)
|
||||
if len(platformModels) == 0 {
|
||||
// CN 供应商没有静态默认模型列表(defaultModelIDsForPlatform 的
|
||||
@@ -1461,9 +1488,14 @@ func defaultModelIDsForPlatform(platform string) []string {
|
||||
return xai.DefaultModelIDs()
|
||||
case service.PlatformOpenCodeGo:
|
||||
return service.DefaultOpenCodeGoModelIDs()
|
||||
case service.PlatformTypeSafe:
|
||||
return []string{"jev-latest"}
|
||||
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 _, id := range defaultModelIDsForPlatform(concretePlatform) {
|
||||
if _, ok := seen[id]; ok {
|
||||
@@ -2149,6 +2181,9 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||||
return
|
||||
}
|
||||
if rejectSystemOneOnlyPlatform(c, apiKey, h.errorResponse) {
|
||||
return
|
||||
}
|
||||
|
||||
setOpsRequestContext(c, parsedReq.Model, parsedReq.Stream)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(parsedReq.Stream, false)))
|
||||
@@ -2514,9 +2549,12 @@ func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task serv
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, abandon := wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped {
|
||||
if mode.Dropped() {
|
||||
abandon()
|
||||
}
|
||||
return
|
||||
}
|
||||
// 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。
|
||||
@@ -2544,7 +2582,7 @@ func (h *GatewayHandler) submitMandatoryUsageRecordTask(parent context.Context,
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, _ = wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() {
|
||||
return
|
||||
|
||||
@@ -82,6 +82,9 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||||
return
|
||||
}
|
||||
if rejectSystemOneOnlyPlatform(c, apiKey, h.chatCompletionsErrorResponse) {
|
||||
return
|
||||
}
|
||||
reqStream, ok := parseOpenAICompatibleStream(body)
|
||||
if !ok {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", invalidStreamFieldTypeMessage)
|
||||
@@ -101,8 +104,11 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/chat/completions is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
@@ -144,6 +150,19 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gateway.cc.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.chatCompletionsErrorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// Parse request for session hash
|
||||
bodyRef := service.NewRequestBodyRef(body)
|
||||
parsedReq, _ := service.ParseGatewayRequest(bodyRef, "chat_completions")
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"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"
|
||||
)
|
||||
|
||||
// groupScopedSchedulerCache 按桶的分组只返回该分组的成员账号,并记录被查询过的分组。
|
||||
type groupScopedSchedulerCache struct {
|
||||
*fakeSchedulerCache
|
||||
mu sync.Mutex
|
||||
groupIDs []int64
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) GetSnapshot(_ context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
|
||||
c.mu.Lock()
|
||||
c.groupIDs = append(c.groupIDs, bucket.GroupID)
|
||||
c.mu.Unlock()
|
||||
var members []*service.Account
|
||||
for _, account := range c.accounts {
|
||||
for _, ag := range account.AccountGroups {
|
||||
if ag.GroupID == bucket.GroupID {
|
||||
members = append(members, account)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return members, true, nil
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) queriedGroupIDs() []int64 {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return append([]int64(nil), c.groupIDs...)
|
||||
}
|
||||
|
||||
type groupMapRepo struct {
|
||||
*fakeGroupRepo
|
||||
groups map[int64]*service.Group
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByID(_ context.Context, id int64) (*service.Group, error) {
|
||||
if group, ok := r.groups[id]; ok {
|
||||
return group, nil
|
||||
}
|
||||
return nil, service.ErrGroupNotFound
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByIDLite(ctx context.Context, id int64) (*service.Group, error) {
|
||||
return r.GetByID(ctx, id)
|
||||
}
|
||||
|
||||
func TestGatewayOpenAICompatibleHandlersClaudeCodeOnlyFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
const (
|
||||
primaryGroupID = int64(9300)
|
||||
fallbackGroupID = int64(9301)
|
||||
primaryAccountID = int64(9310)
|
||||
fallbackAccountID = int64(9311)
|
||||
)
|
||||
|
||||
// 账号不带 api_key:转发在取令牌时失败,请求不会触达上游,断言只看选号结果。
|
||||
newAccount := func(id, groupID int64) *service.Account {
|
||||
return &service.Account{
|
||||
ID: id, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 1,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: id, GroupID: groupID}},
|
||||
}
|
||||
}
|
||||
|
||||
endpoints := []struct {
|
||||
name string
|
||||
path string
|
||||
body string
|
||||
call func(*GatewayHandler, *gin.Context)
|
||||
}{
|
||||
{
|
||||
name: "responses", path: "/v1/responses",
|
||||
body: `{"model":"claude-sonnet-4-5","input":"hello","stream":false}`,
|
||||
call: (*GatewayHandler).Responses,
|
||||
},
|
||||
{
|
||||
name: "chat completions", path: "/v1/chat/completions",
|
||||
body: `{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}],"stream":false}`,
|
||||
call: (*GatewayHandler).ChatCompletions,
|
||||
},
|
||||
}
|
||||
|
||||
for _, ep := range endpoints {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
hasFallback bool
|
||||
}{
|
||||
{name: "with fallback group", hasFallback: true},
|
||||
{name: "without fallback group", hasFallback: false},
|
||||
} {
|
||||
t.Run(ep.name+"/"+tc.name, func(t *testing.T) {
|
||||
primary := &service.Group{
|
||||
ID: primaryGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive, ClaudeCodeOnly: true,
|
||||
}
|
||||
if tc.hasFallback {
|
||||
fallbackID := fallbackGroupID
|
||||
primary.FallbackGroupID = &fallbackID
|
||||
}
|
||||
fallback := &service.Group{
|
||||
ID: fallbackGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
schedulerCache := &groupScopedSchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{
|
||||
newAccount(primaryAccountID, primaryGroupID),
|
||||
newAccount(fallbackAccountID, fallbackGroupID),
|
||||
}}}
|
||||
gatewayService := service.NewGatewayService(
|
||||
nil, &groupMapRepo{fakeGroupRepo: &fakeGroupRepo{}, groups: map[int64]*service.Group{
|
||||
primaryGroupID: primary,
|
||||
fallbackGroupID: fallback,
|
||||
}}, nil, nil, nil, nil, nil, nil, nil,
|
||||
service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil),
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
billingCacheService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCacheService.Stop)
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gatewayService,
|
||||
billingCacheService: billingCacheService,
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0),
|
||||
maxAccountSwitches: 1,
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
primaryGroupIDRef := primaryGroupID
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9320, UserID: 9330, GroupID: &primaryGroupIDRef, Group: primary, Status: service.StatusActive,
|
||||
User: &service.User{ID: 9330, Concurrency: 10, Balance: 100},
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
ctx := context.WithValue(context.Background(), ctxkey.Group, primary)
|
||||
req := httptest.NewRequest(http.MethodPost, ep.path, bytes.NewBufferString(ep.body)).WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
|
||||
ep.call(h, c)
|
||||
|
||||
selected, reachedSelection := c.Get(opsAccountIDKey)
|
||||
if !tc.hasFallback {
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "This group is restricted to Claude Code clients")
|
||||
require.False(t, reachedSelection, "a Claude Code only group without fallback must be rejected before account selection")
|
||||
require.Empty(t, schedulerCache.queriedGroupIDs())
|
||||
return
|
||||
}
|
||||
require.NotContains(t, recorder.Body.String(), "restricted to Claude Code clients")
|
||||
require.True(t, reachedSelection, "a Claude Code only group with fallback must reach account selection")
|
||||
require.Equal(t, fallbackAccountID, selected)
|
||||
queried := schedulerCache.queriedGroupIDs()
|
||||
require.Contains(t, queried, fallbackGroupID)
|
||||
require.NotContains(t, queried, primaryGroupID)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -82,6 +82,9 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
h.responsesErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||||
return
|
||||
}
|
||||
if rejectSystemOneOnlyPlatform(c, apiKey, h.responsesErrorResponse) {
|
||||
return
|
||||
}
|
||||
reqStream, ok := parseOpenAICompatibleStream(body)
|
||||
if !ok {
|
||||
h.responsesErrorResponse(c, http.StatusBadRequest, "invalid_request_error", invalidStreamFieldTypeMessage)
|
||||
@@ -105,13 +108,11 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(requestCtx, apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction:
|
||||
// /v1/responses is never a Claude Code endpoint.
|
||||
// When claude_code_only is enabled, this endpoint is rejected.
|
||||
// The existing service-layer checkClaudeCodeRestriction handles degradation
|
||||
// to fallback groups when the Forward path calls SelectAccountForModelWithExclusions.
|
||||
// Here we just reject at handler level since /v1/responses clients can't be Claude Code.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/responses is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.responsesErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
@@ -153,6 +154,19 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gateway.responses.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.responsesErrorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// Parse request for session hash
|
||||
bodyRef := service.NewRequestBodyRef(body)
|
||||
parsedReq, _ := service.ParseGatewayRequest(bodyRef, "responses")
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// inflightReservationEstimator 估算单请求在途预留金额(USD);false 表示无法定价。
|
||||
type inflightReservationEstimator interface {
|
||||
EstimateInflightReservation(ctx context.Context, apiKey *service.APIKey, req service.InflightEstimateRequest) (float64, bool)
|
||||
}
|
||||
|
||||
// requestMaxOutputTokens 从请求体中提取输出 token 上限(兼容 Anthropic / OpenAI Chat / Responses / Gemini)。
|
||||
func requestMaxOutputTokens(body []byte) int {
|
||||
for _, path := range []string{"max_tokens", "max_completion_tokens", "max_output_tokens", "generationConfig.maxOutputTokens", "generation_config.max_output_tokens"} {
|
||||
if v := gjson.GetBytes(body, path); v.Exists() && v.Type == gjson.Number && v.Int() > 0 {
|
||||
return int(v.Int())
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// tokenInflightEstimate 文本类请求的估算输入。
|
||||
func tokenInflightEstimate(model string, body []byte) service.InflightEstimateRequest {
|
||||
return service.InflightEstimateRequest{
|
||||
Model: model,
|
||||
BodyBytes: len(body),
|
||||
MaxTokens: requestMaxOutputTokens(body),
|
||||
Kind: service.InflightEstimateToken,
|
||||
}
|
||||
}
|
||||
|
||||
func inflightNoop() {}
|
||||
|
||||
// reserveInflightBalance 在 CheckBillingEligibility 之后为余额模式请求登记在途预留。
|
||||
//
|
||||
// 成功时把预留句柄挂到 c.Request 的 context 上:之后通过 submit*UsageRecordTask 提交的
|
||||
// 计费任务会接管一个引用,直到余额缓存实际扣减后才释放。返回的 done 必须 defer 调用
|
||||
// (停止续期并归还 handler 引用;没有提交计费任务时即刻释放)。
|
||||
// 开关关闭、订阅模式、Redis 故障时返回 no-op(fail-open);无法定价时默认 fail-open,
|
||||
// 配置 fail_closed_on_unpriced=true 时返回 ErrInsufficientBalance。
|
||||
func reserveInflightBalance(
|
||||
c *gin.Context,
|
||||
billing *service.BillingCacheService,
|
||||
estimator inflightReservationEstimator,
|
||||
apiKey *service.APIKey,
|
||||
subscription *service.UserSubscription,
|
||||
req service.InflightEstimateRequest,
|
||||
) (func(), error) {
|
||||
if c == nil || c.Request == nil {
|
||||
return inflightNoop, nil
|
||||
}
|
||||
ctx, done, err := reserveInflightBalanceCtx(c.Request.Context(), billing, estimator, apiKey, subscription, req)
|
||||
if err != nil {
|
||||
return inflightNoop, err
|
||||
}
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
return done, nil
|
||||
}
|
||||
|
||||
// reserveInflightBalanceCtx 同 reserveInflightBalance,但返回携带预留句柄的新 context
|
||||
// (供 WebSocket 等自管 context 的路径)。
|
||||
func reserveInflightBalanceCtx(
|
||||
ctx context.Context,
|
||||
billing *service.BillingCacheService,
|
||||
estimator inflightReservationEstimator,
|
||||
apiKey *service.APIKey,
|
||||
subscription *service.UserSubscription,
|
||||
req service.InflightEstimateRequest,
|
||||
) (context.Context, func(), error) {
|
||||
if billing == nil || estimator == nil || apiKey == nil || apiKey.User == nil || !billing.InflightReservationEnabled() {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() && subscription != nil {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
estimate, priced := estimator.EstimateInflightReservation(ctx, apiKey, req)
|
||||
if !priced && billing.InflightReservationFailClosedOnUnpriced() {
|
||||
return ctx, inflightNoop, service.ErrInsufficientBalance
|
||||
}
|
||||
if estimate <= 0 {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
res, err := billing.ReserveInflight(ctx, apiKey.User, apiKey.Group, subscription, estimate)
|
||||
if err != nil {
|
||||
return ctx, inflightNoop, err
|
||||
}
|
||||
if res == nil {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
return service.WithInflightReservation(ctx, res), res.HandlerDone, nil
|
||||
}
|
||||
|
||||
// grokMediaInflightEstimate 媒体生成请求的估算输入;状态/内容查询返回空模型(不预留:
|
||||
// 查询会为已生成的媒体计费,不能因余额预留而拦截用户取回已付费结果)。
|
||||
func grokMediaInflightEstimate(endpoint service.GrokMediaEndpoint, model string, info service.GrokMediaRequestInfo, body []byte) service.InflightEstimateRequest {
|
||||
if !endpoint.IsGenerationRequest() {
|
||||
return service.InflightEstimateRequest{}
|
||||
}
|
||||
switch endpoint {
|
||||
case service.GrokMediaEndpointImagesGenerations, service.GrokMediaEndpointImagesEdits:
|
||||
return service.InflightEstimateRequest{Model: model, BodyBytes: len(body), Kind: service.InflightEstimateImage, Units: info.N}
|
||||
default:
|
||||
return service.InflightEstimateRequest{
|
||||
Model: model,
|
||||
BodyBytes: len(body),
|
||||
Kind: service.InflightEstimateVideo,
|
||||
Units: 1,
|
||||
VideoResolution: info.Resolution,
|
||||
VideoDurationSeconds: info.DurationSeconds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// grokVoiceSTTBytesPerSecond STT 时长粗估(~128kbps 压缩音频)。
|
||||
const grokVoiceSTTBytesPerSecond = 16000
|
||||
|
||||
// grokVoiceInflightEstimate 语音 HTTP 接口估算:TTS 按输入字符数(百万字符),STT 按音频字节粗估时长(小时)。
|
||||
// 其他接口(custom-voices)无音频计量,返回的估算为 0(不预留)。
|
||||
func grokVoiceInflightEstimate(endpoint string, body []byte) service.InflightEstimateRequest {
|
||||
req := service.InflightEstimateRequest{Model: endpoint, Kind: service.InflightEstimateAudio, AudioMode: endpoint}
|
||||
switch endpoint {
|
||||
case "tts":
|
||||
req.AudioUnits = float64(len([]rune(extractGrokTTSInputText(body)))) / 1e6
|
||||
case "stt":
|
||||
req.AudioUnits = float64(len(body)) / grokVoiceSTTBytesPerSecond / 3600
|
||||
}
|
||||
return req
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRequestMaxOutputTokens(t *testing.T) {
|
||||
require.Equal(t, 1024, requestMaxOutputTokens([]byte(`{"max_tokens":1024}`)))
|
||||
require.Equal(t, 2048, requestMaxOutputTokens([]byte(`{"max_completion_tokens":2048}`)))
|
||||
require.Equal(t, 4096, requestMaxOutputTokens([]byte(`{"max_output_tokens":4096}`)))
|
||||
require.Equal(t, 512, requestMaxOutputTokens([]byte(`{"generationConfig":{"maxOutputTokens":512}}`)))
|
||||
require.Equal(t, 0, requestMaxOutputTokens([]byte(`{"max_tokens":"x"}`)))
|
||||
require.Equal(t, 0, requestMaxOutputTokens([]byte(`{}`)))
|
||||
}
|
||||
|
||||
type countingEstimator struct {
|
||||
calls int
|
||||
cost float64
|
||||
priced bool
|
||||
}
|
||||
|
||||
func (e *countingEstimator) EstimateInflightReservation(context.Context, *service.APIKey, service.InflightEstimateRequest) (float64, bool) {
|
||||
e.calls++
|
||||
return e.cost, e.priced
|
||||
}
|
||||
|
||||
func newInflightTestGinContext() *gin.Context {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
return c
|
||||
}
|
||||
|
||||
func TestReserveInflightBalance_SkipsWhenDisabledOrSubscription(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
billing := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billing.Stop)
|
||||
est := &countingEstimator{cost: 1, priced: true}
|
||||
apiKey := &service.APIKey{User: &service.User{ID: 1}}
|
||||
|
||||
done, err := reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("m", []byte(`{}`)))
|
||||
require.NoError(t, err)
|
||||
done()
|
||||
require.Equal(t, 0, est.calls, "disabled switch must not even estimate")
|
||||
|
||||
cfg.Billing.InflightReservation.Enabled = true
|
||||
apiKey.Group = &service.Group{SubscriptionType: service.SubscriptionTypeSubscription}
|
||||
done, err = reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, &service.UserSubscription{}, tokenInflightEstimate("m", []byte(`{}`)))
|
||||
require.NoError(t, err)
|
||||
done()
|
||||
require.Equal(t, 0, est.calls, "subscription mode must be unaffected")
|
||||
}
|
||||
|
||||
func TestReserveInflightBalance_UnpricedFailOpenByDefaultFailClosedOptIn(t *testing.T) {
|
||||
cache := newHandlerInflightCache(10)
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
billing := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billing.Stop)
|
||||
apiKey := &service.APIKey{User: &service.User{ID: 1}}
|
||||
est := &countingEstimator{priced: false}
|
||||
|
||||
done, err := reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("unknown", nil))
|
||||
require.NoError(t, err)
|
||||
done()
|
||||
|
||||
cfg.Billing.InflightReservation.FailClosedOnUnpriced = true
|
||||
_, err = reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("unknown", nil))
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance)
|
||||
}
|
||||
|
||||
// handlerInflightCache 内存版余额缓存 + 在途预留(语义同 Redis Lua)。
|
||||
type handlerInflightCache struct {
|
||||
service.BillingCache
|
||||
mu sync.Mutex
|
||||
balance float64
|
||||
res map[string]float64
|
||||
}
|
||||
|
||||
func newHandlerInflightCache(balance float64) *handlerInflightCache {
|
||||
return &handlerInflightCache{balance: balance, res: map[string]float64{}}
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) GetUserBalance(context.Context, int64) (float64, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return m.balance, nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) GetUserPlatformQuotaCache(context.Context, int64, string) (*service.UserPlatformQuotaCacheEntry, bool, error) {
|
||||
return nil, false, nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) ReserveInflightBalance(_ context.Context, _ int64, id string, amount, balance float64, _ time.Duration) (bool, float64, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
sum := 0.0
|
||||
for _, v := range m.res {
|
||||
sum += v
|
||||
}
|
||||
if len(m.res) > 0 && balance-sum < amount {
|
||||
return false, sum, nil
|
||||
}
|
||||
m.res[id] = amount
|
||||
return true, sum, nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) ReleaseInflightBalance(_ context.Context, _ int64, id string) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
delete(m.res, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) count() int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return len(m.res)
|
||||
}
|
||||
|
||||
func TestWrapUsageRecordTaskContext_HandsReservationToBillingTask(t *testing.T) {
|
||||
cache := newHandlerInflightCache(1)
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
billing := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billing.Stop)
|
||||
apiKey := &service.APIKey{User: &service.User{ID: 5}}
|
||||
|
||||
c := newInflightTestGinContext()
|
||||
done, err := reserveInflightBalance(c, billing, &countingEstimator{cost: 0.9, priced: true}, apiKey, nil, tokenInflightEstimate("m", nil))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, cache.count())
|
||||
|
||||
ran := false
|
||||
task, abandon := wrapUsageRecordTaskContext(c.Request.Context(), func(context.Context) { ran = true })
|
||||
done() // handler returns; billing still pending
|
||||
require.Equal(t, 1, cache.count(), "reservation held until the billing task finishes")
|
||||
task(context.Background())
|
||||
require.True(t, ran)
|
||||
require.Equal(t, 0, cache.count())
|
||||
abandon() // idempotent with the task's own done
|
||||
|
||||
// Dropped task: the submitter abandons it and the reservation is released.
|
||||
c2 := newInflightTestGinContext()
|
||||
done2, err := reserveInflightBalance(c2, billing, &countingEstimator{cost: 0.9, priced: true}, apiKey, nil, tokenInflightEstimate("m", nil))
|
||||
require.NoError(t, err)
|
||||
_, abandon2 := wrapUsageRecordTaskContext(c2.Request.Context(), func(context.Context) {})
|
||||
done2()
|
||||
require.Equal(t, 1, cache.count())
|
||||
abandon2()
|
||||
require.Equal(t, 0, cache.count())
|
||||
}
|
||||
|
||||
// 新接入的端点(独立 web_search):在途预留超过余额时拒绝,且不残留预留。
|
||||
func TestWebSearch_RejectsWhenInflightExceedsBalance(t *testing.T) {
|
||||
cache := newHandlerInflightCache(1.5)
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
billingCache := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCache.Stop)
|
||||
billing := service.NewBillingService(cfg, nil)
|
||||
gw := service.NewGatewayService(
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, cfg, nil, nil, billing, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, service.NewModelPricingResolver(nil, billing), nil, nil, nil,
|
||||
)
|
||||
h := &GatewayHandler{gatewayService: gw, billingCacheService: billingCache}
|
||||
|
||||
groupID := int64(3)
|
||||
perK := 1000.0 // $1 per search
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9, User: &service.User{ID: 42, Balance: 1.5}, GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, RateMultiplier: 1, SearchPricePer1k: &perK},
|
||||
}
|
||||
|
||||
// Another in-flight request of this user already holds $1.
|
||||
held, err := billingCache.ReserveInflight(context.Background(), apiKey.User, apiKey.Group, nil, 1.0)
|
||||
require.NoError(t, err)
|
||||
defer held.HandlerDone()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/web_search", bytes.NewBufferString(`{"query":"sub2api"}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), apiKey)
|
||||
|
||||
h.WebSearch(c)
|
||||
|
||||
require.NotEqual(t, http.StatusOK, w.Code)
|
||||
require.Contains(t, w.Body.String(), "balance")
|
||||
require.Equal(t, 1, cache.count(), "rejected request must not leave a reservation behind")
|
||||
}
|
||||
@@ -27,7 +27,7 @@ func requestModelForTest(h *GatewayHandler, group *service.Group, modelID, etag
|
||||
|
||||
func TestRetrieveModelMatchesVisibleCatalogue(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, platform := range []string{service.PlatformOpenAI, service.PlatformAnthropic, service.PlatformGemini, service.PlatformGrok, service.PlatformComposite} {
|
||||
for _, platform := range []string{service.PlatformOpenAI, service.PlatformAnthropic, service.PlatformGemini, service.PlatformGrok, service.PlatformTypeSafe, service.PlatformComposite} {
|
||||
for _, mapped := range []bool{false, true} {
|
||||
name := platform + "/fallback"
|
||||
if mapped {
|
||||
|
||||
@@ -932,6 +932,10 @@ func TestDefaultCodexModelIDsForPlatform_DeepSeekUsesDeepSeekModels(t *testing.T
|
||||
require.Equal(t, defaultModelIDsForPlatform(service.PlatformAnthropic), defaultCodexModelIDsForPlatform(service.PlatformAnthropic))
|
||||
}
|
||||
|
||||
func TestDefaultModelIDsForPlatform_TypeSafeUsesJev(t *testing.T) {
|
||||
require.Equal(t, []string{"jev-latest"}, defaultModelIDsForPlatform(service.PlatformTypeSafe))
|
||||
}
|
||||
|
||||
func TestGatewayCodexModels_DeepSeekWithoutMappingUsesDeepSeekDefaults(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 130
|
||||
@@ -1499,9 +1503,9 @@ func TestGatewayModels_GPT6SolLunaDiscoveryRespectsGroupAndAccountRestrictions(t
|
||||
restricted bool
|
||||
want []string
|
||||
}{
|
||||
{"selected and ordered", []string{"gpt-6-luna", "gpt-6-sol"}, false, []string{"gpt-6-luna", "gpt-6-sol"}},
|
||||
{"selected and ordered", []string{"gpt-6.1-sol", "gpt-6-luna", "gpt-6-sol"}, false, []string{"gpt-6.1-sol", "gpt-6-luna", "gpt-6-sol"}},
|
||||
{"group excludes new models", []string{"gpt-5.6-sol"}, false, []string{"gpt-5.6-sol"}},
|
||||
{"account restricts new models", []string{"gpt-6-sol", "gpt-6-luna", "gpt-5.6-sol"}, true, []string{"gpt-5.6-sol"}},
|
||||
{"account restricts new models", []string{"gpt-6.1-sol", "gpt-6-sol", "gpt-6-luna", "gpt-5.6-sol"}, true, []string{"gpt-5.6-sol"}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
groupID := int64(25)
|
||||
@@ -1522,3 +1526,47 @@ func TestGatewayModels_GPT6SolLunaDiscoveryRespectsGroupAndAccountRestrictions(t
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Scenario: jev-latest only works through /v1/systemone, so Composite groups list
|
||||
// it in /v1/models only when they can serve it, and never in the Codex manifest.
|
||||
func TestGatewayModels_CompositeTypeSafeListingScope(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(66)
|
||||
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{
|
||||
groupID: {{ID: 1, Platform: service.PlatformAnthropic}, {ID: 2, Platform: service.PlatformTypeSafe, Type: service.AccountTypeAPIKey}},
|
||||
},
|
||||
})
|
||||
newContext := func(path string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, path, nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
|
||||
})
|
||||
return c, rec
|
||||
}
|
||||
|
||||
c, rec := newContext("/v1/models")
|
||||
h.Models(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var models gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &models))
|
||||
require.Contains(t, modelIDsForTest(models.Data), "jev-latest")
|
||||
require.Contains(t, modelIDsForTest(models.Data), "claude-opus-4-6")
|
||||
|
||||
c, rec = newContext("/models?client_version=0.147.0")
|
||||
h.CodexModels(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var manifest codexModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &manifest))
|
||||
slugs := codexModelSlugsForTest(manifest.Models)
|
||||
require.Contains(t, slugs, "claude-opus-4-6")
|
||||
require.NotContains(t, slugs, "jev-latest")
|
||||
}
|
||||
|
||||
func TestDefaultModelIDsForPlatform_CompositeFallbackExcludesTypeSafe(t *testing.T) {
|
||||
require.NotContains(t, defaultModelIDsForPlatform(service.PlatformComposite), "jev-latest")
|
||||
require.NotContains(t, defaultCodexModelIDsForPlatform(service.PlatformComposite), "jev-latest")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,287 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/typesafe"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const systemOneOnlyPlatformMessage = "TypeSafe models are only available through POST /v1/systemone"
|
||||
|
||||
// rejectSystemOneOnlyPlatform stops TypeSafe traffic from entering a non System
|
||||
// One protocol chain. TypeSafe accounts only speak the native System One
|
||||
// protocol; letting them reach the Anthropic/OpenAI converters would send
|
||||
// foreign payloads (and the account key) to the wrong upstream path and feed
|
||||
// the resulting auth failures back into account state.
|
||||
func rejectSystemOneOnlyPlatform(c *gin.Context, apiKey *service.APIKey, writeError func(*gin.Context, int, string, string)) bool {
|
||||
if _, forced := middleware2.GetForcePlatformFromContext(c); forced {
|
||||
return false
|
||||
}
|
||||
if effectiveAPIKeyPlatform(c, apiKey) != service.PlatformTypeSafe {
|
||||
return false
|
||||
}
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
writeError(c, http.StatusNotFound, "not_found_error", systemOneOnlyPlatformMessage)
|
||||
return true
|
||||
}
|
||||
|
||||
// SystemOne proxies TypeSafe's native, non-streaming System One protocol.
|
||||
func (h *GatewayHandler) SystemOne(c *gin.Context) {
|
||||
requestStart := time.Now()
|
||||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||||
if !ok || apiKey == nil || apiKey.Group == nil {
|
||||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||||
return
|
||||
}
|
||||
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
||||
if !ok {
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
||||
return
|
||||
}
|
||||
reqLog := requestLogger(c, "handler.gateway.systemone",
|
||||
zap.Int64("user_id", subject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
)
|
||||
|
||||
body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg)
|
||||
if err != nil {
|
||||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||||
h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
||||
return
|
||||
}
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
|
||||
return
|
||||
}
|
||||
model, err := typesafe.ValidateSystemOneRequest(body)
|
||||
if err != nil {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
ensureCompositeTargetPlatform(c, apiKey, model)
|
||||
if apiKey.Group.Platform != service.PlatformTypeSafe &&
|
||||
(apiKey.Group.Platform != service.PlatformComposite || !compositeTargetPlatformAllowed(c, apiKey, model, service.PlatformTypeSafe)) {
|
||||
h.errorResponse(c, http.StatusNotFound, "not_found_error", "System One is only available for TypeSafe and compatible Composite groups")
|
||||
return
|
||||
}
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolTypeSafeSystemOne, model, body); decision != nil && !decision.AllowNextStage {
|
||||
h.anthropicSecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
setOpsRequestContext(c, model, false)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeSync))
|
||||
service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds())
|
||||
pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context())
|
||||
c.Request = c.Request.WithContext(pricingCtx)
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, model)
|
||||
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
||||
|
||||
streamStarted := false
|
||||
userRelease, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, subject.UserID, subject.Concurrency, false, &streamStarted)
|
||||
if err != nil {
|
||||
reqLog.Warn("systemone.user_slot_acquire_failed", zap.Error(err))
|
||||
h.handleConcurrencyError(c, err, "user", false)
|
||||
return
|
||||
}
|
||||
// 在请求结束或 Context 取消时确保释放槽位,避免客户端断开造成泄漏。
|
||||
userRelease = wrapReleaseOnDone(c.Request.Context(), userRelease)
|
||||
if userRelease != nil {
|
||||
defer userRelease()
|
||||
}
|
||||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
|
||||
reqLog.Info("systemone.billing_eligibility_check_failed", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(model, body))
|
||||
if err != nil {
|
||||
reqLog.Info("systemone.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
fs := NewFailoverState(h.maxAccountSwitches, false)
|
||||
for {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, "", model, fs.FailedAccountIDs, "", subject.UserID)
|
||||
if err == nil && (selection == nil || selection.Account == nil) {
|
||||
err = service.ErrNoAvailableAccounts
|
||||
}
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("systemone.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, model, model, service.PlatformTypeSafe)
|
||||
cls = classifySelectionFailureError(err, cls)
|
||||
if !cls.ModelNotFound {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
}
|
||||
reqLog.Warn("systemone.select_account_no_available", zap.Bool("model_not_found", cls.ModelNotFound), zap.Error(err))
|
||||
h.errorResponse(c, cls.Status, cls.ErrType, cls.Message)
|
||||
return
|
||||
}
|
||||
switch fs.HandleSelectionExhausted(c.Request.Context()) {
|
||||
case FailoverContinue:
|
||||
continue
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
default:
|
||||
if fs.LastFailoverErr != nil {
|
||||
h.handleFailoverExhausted(c, fs.LastFailoverErr, service.PlatformTypeSafe, false)
|
||||
} else {
|
||||
h.handleFailoverExhaustedSimple(c, http.StatusBadGateway, false)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
account := selection.Account
|
||||
|
||||
accountRelease := selection.ReleaseFunc
|
||||
if !selection.Acquired {
|
||||
if selection.WaitPlan == nil {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "No available accounts")
|
||||
return
|
||||
}
|
||||
accountRelease, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(c, account.ID, selection.WaitPlan.MaxConcurrency, selection.WaitPlan.Timeout, false, &streamStarted)
|
||||
if err != nil {
|
||||
reqLog.Warn("systemone.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
h.handleConcurrencyError(c, err, "account", false)
|
||||
return
|
||||
}
|
||||
}
|
||||
// 准入终检:与其他网关入口一致,利润控制否决的账号不得承接本次请求。
|
||||
admissionCtx := service.ContextWithSelectionProfitGate(c.Request.Context(), selection)
|
||||
latest, vetoed, reason := h.gatewayService.GatewayProfitControlVetoLatest(admissionCtx, account)
|
||||
if vetoed {
|
||||
if accountRelease != nil {
|
||||
accountRelease()
|
||||
}
|
||||
reqLog.Debug("systemone.account_slot_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason))
|
||||
if fs.RecordProfitVeto(account.ID) == FailoverExhausted {
|
||||
reqLog.Warn("systemone.profit_veto_attempts_exhausted", zap.Int("profit_veto_count", fs.ProfitVetoCount()))
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", profitVetoExhaustedMessage)
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
account = latest
|
||||
accountRelease = wrapReleaseOnDone(c.Request.Context(), accountRelease)
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
service.SetOpsUpstreamModel(c, model)
|
||||
|
||||
forwardStart := time.Now()
|
||||
result, forwardErr := h.gatewayService.ForwardSystemOne(c.Request.Context(), c, account, body)
|
||||
if accountRelease != nil {
|
||||
accountRelease()
|
||||
}
|
||||
service.SetOpsLatencyMs(c, service.OpsResponseLatencyMsKey, time.Since(forwardStart).Milliseconds())
|
||||
|
||||
if forwardErr != nil {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(forwardErr, &failoverErr) {
|
||||
switch fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr) {
|
||||
case FailoverContinue:
|
||||
reqLog.Warn("systemone.upstream_failover_switching",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
zap.Int("switch_count", fs.SwitchCount),
|
||||
)
|
||||
continue
|
||||
case FailoverExhausted:
|
||||
h.handleFailoverExhausted(c, fs.LastFailoverErr, service.PlatformTypeSafe, false)
|
||||
return
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
}
|
||||
}
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
var upstreamErr *service.SystemOneUpstreamError
|
||||
if errors.As(forwardErr, &upstreamErr) {
|
||||
status := upstreamErr.StatusCode
|
||||
if !service.IsSystemOneRequestErrorStatus(status) {
|
||||
status = http.StatusBadGateway
|
||||
}
|
||||
h.errorResponse(c, status, "upstream_error", "TypeSafe rejected the request")
|
||||
return
|
||||
}
|
||||
reqLog.Warn("systemone.forward_failed", zap.Int64("account_id", account.ID), zap.Error(forwardErr))
|
||||
if errors.Is(forwardErr, typesafe.ErrSystemOneResponseTooLarge) {
|
||||
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "TypeSafe response exceeds the gateway size limit")
|
||||
return
|
||||
}
|
||||
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "TypeSafe upstream request failed")
|
||||
return
|
||||
}
|
||||
|
||||
c.Data(result.StatusCode, result.ContentType, result.Body)
|
||||
h.recordSystemOneUsage(c, apiKey, account, subscription, channelMapping, model, body, result, subject.UserID, pricingAt)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) recordSystemOneUsage(c *gin.Context, apiKey *service.APIKey, account *service.Account, subscription *service.UserSubscription, mapping service.ChannelMappingResult, model string, body []byte, result *service.SystemOneForwardResult, userID int64, pricingAt time.Time) {
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||||
|
||||
h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: &result.ForwardResult,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
PricingAt: pricingAt,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
SessionID: sessionID,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, mapping, model, result.UpstreamModel),
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.gateway.systemone"),
|
||||
zap.Int64("user_id", userID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Int64("account_id", account.ID),
|
||||
).Error("systemone.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const validSystemOneHandlerBody = `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul","instructions":"Evaluate"}}}`
|
||||
|
||||
func newSystemOneHandlerContext(body string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/systemone", strings.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
return c, recorder
|
||||
}
|
||||
|
||||
func TestSystemOneRequiresAuthentication(t *testing.T) {
|
||||
c, recorder := newSystemOneHandlerContext(validSystemOneHandlerBody)
|
||||
(&GatewayHandler{}).SystemOne(c)
|
||||
require.Equal(t, http.StatusUnauthorized, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "authentication_error")
|
||||
}
|
||||
|
||||
func TestSystemOneRejectsNonTypeSafeGroupBeforeScheduling(t *testing.T) {
|
||||
c, recorder := newSystemOneHandlerContext(validSystemOneHandlerBody)
|
||||
groupID := int64(3)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ID: 4, UserID: 5, GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI}})
|
||||
c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 5, Concurrency: 1})
|
||||
|
||||
(&GatewayHandler{cfg: &config.Config{Gateway: config.GatewayConfig{MaxBodySize: 1 << 20}}}).SystemOne(c)
|
||||
require.Equal(t, http.StatusNotFound, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "only available for TypeSafe")
|
||||
}
|
||||
|
||||
func newTypeSafeGroupContext(t *testing.T, path, body, groupPlatform string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
t.Helper()
|
||||
c, recorder := newSystemOneHandlerContext(body)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, path, strings.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
groupID := int64(9)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ID: 4, UserID: 5, GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: groupPlatform}})
|
||||
c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 5, Concurrency: 1})
|
||||
return c, recorder
|
||||
}
|
||||
|
||||
func TestRejectSystemOneOnlyPlatform(t *testing.T) {
|
||||
write := func(c *gin.Context, status int, errType, message string) {
|
||||
c.JSON(status, gin.H{"type": errType, "message": message})
|
||||
}
|
||||
|
||||
c, recorder := newTypeSafeGroupContext(t, "/v1/messages", `{}`, service.PlatformTypeSafe)
|
||||
require.True(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
require.Equal(t, http.StatusNotFound, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), systemOneOnlyPlatformMessage)
|
||||
|
||||
c, _ = newTypeSafeGroupContext(t, "/v1/messages", `{}`, service.PlatformComposite)
|
||||
require.False(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
c.Request = c.Request.WithContext(service.WithResolvedTargetPlatform(c.Request.Context(), service.PlatformTypeSafe))
|
||||
require.True(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
|
||||
c, _ = newTypeSafeGroupContext(t, "/v1/messages", `{}`, service.PlatformAnthropic)
|
||||
require.False(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
|
||||
c, _ = newTypeSafeGroupContext(t, "/antigravity/v1/messages", `{}`, service.PlatformTypeSafe)
|
||||
c.Set(string(middleware2.ContextKeyForcePlatform), service.PlatformAntigravity)
|
||||
require.False(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
}
|
||||
|
||||
func mustAPIKey(t *testing.T, c *gin.Context) *service.APIKey {
|
||||
t.Helper()
|
||||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||||
require.True(t, ok)
|
||||
return apiKey
|
||||
}
|
||||
|
||||
func TestTypeSafeGroupsRejectNonSystemOneProtocolsBeforeScheduling(t *testing.T) {
|
||||
h := &GatewayHandler{
|
||||
cfg: &config.Config{Gateway: config.GatewayConfig{MaxBodySize: 1 << 20}},
|
||||
gatewayService: &service.GatewayService{},
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
path string
|
||||
body string
|
||||
handler func(*gin.Context)
|
||||
}{
|
||||
{"messages", "/v1/messages", `{"model":"jev-latest","max_tokens":1,"messages":[{"role":"user","content":"hi"}]}`, h.Messages},
|
||||
{"count_tokens", "/v1/messages/count_tokens", `{"model":"jev-latest","messages":[{"role":"user","content":"hi"}]}`, h.CountTokens},
|
||||
{"chat_completions", "/v1/chat/completions", `{"model":"jev-latest","messages":[{"role":"user","content":"hi"}]}`, h.ChatCompletions},
|
||||
{"responses", "/v1/responses", `{"model":"jev-latest","input":"hi"}`, h.Responses},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, recorder := newTypeSafeGroupContext(t, tc.path, tc.body, service.PlatformTypeSafe)
|
||||
tc.handler(c)
|
||||
require.Equal(t, http.StatusNotFound, recorder.Code, recorder.Body.String())
|
||||
require.Contains(t, recorder.Body.String(), systemOneOnlyPlatformMessage)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -88,6 +88,18 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: "grok-" + strings.ReplaceAll(searchLabel, "_", "-"), Kind: service.InflightEstimatePerRequest, SearchCalls: 1})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
c.JSON(status, gin.H{"error": gin.H{"type": code, "message": message}})
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
subject, _ := middleware2.GetAuthSubjectFromContext(c)
|
||||
reqLog := requestLogger(c, "handler.gateway.web_search")
|
||||
// Audit user search query before upstream Grok web_search traffic.
|
||||
|
||||
@@ -390,6 +390,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
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)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleError(c, status, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// 3) select account (sticky session based on request body)
|
||||
// 优先使用 Gemini CLI 的会话标识(privileged-user-id + tmp 目录哈希)
|
||||
sessionHash := extractGeminiCLISessionHash(c, body)
|
||||
|
||||
@@ -49,6 +49,18 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
|
||||
if strings.TrimSpace(model) == "" {
|
||||
model = "grok-voice-latest"
|
||||
}
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: model, Kind: service.InflightEstimateAudio, AudioMode: "realtime", AudioUnits: 1})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
// Keep the HTTP response uncommitted while selecting and probing an account.
|
||||
// Realtime is not an HTTP streaming response; using reqStream=true here would
|
||||
// let the wait queue flush an SSE ping before the WebSocket handshake succeeds.
|
||||
@@ -212,6 +224,18 @@ func (h *OpenAIGatewayHandler) GrokVoice(c *gin.Context, endpoint string) {
|
||||
return
|
||||
}
|
||||
}
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, grokVoiceInflightEstimate(endpoint, body))
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
contentType := c.GetHeader("Content-Type")
|
||||
if strings.TrimSpace(contentType) == "" {
|
||||
contentType = "application/json"
|
||||
|
||||
@@ -173,6 +173,18 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, grokMediaInflightEstimate(endpoint, routingModel, requestInfo, body))
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
sessionSeed := body
|
||||
if len(sessionSeed) == 0 && strings.TrimSpace(requestID) != "" {
|
||||
sessionSeed = []byte(requestID)
|
||||
|
||||
@@ -109,6 +109,18 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: requestedModel, BodyBytes: len(body), MaxTokens: requestMaxOutputTokens(body), Kind: service.InflightEstimateToken, SearchCalls: 1})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
searchID := strings.TrimSpace(gjson.GetBytes(body, "id").String())
|
||||
sessionHash := h.gatewayService.GenerateSessionHashWithFallback(c, nil, searchID)
|
||||
profitVetoCount := 0
|
||||
|
||||
@@ -146,6 +146,19 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("openai_chat_completions.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c, body)
|
||||
promptCacheKey := h.gatewayService.ExtractSessionID(c, body)
|
||||
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func (h *OpenAIGatewayHandler) cyberPolicyLogOnly(c *gin.Context, apiKey *service.APIKey) bool {
|
||||
return h != nil && c != nil && c.Request != nil && h.gatewayService.CyberPolicyLogOnly(c.Request.Context(), apiKey)
|
||||
}
|
||||
|
||||
// Both HTTP and WebSocket admission skip existing blocks for trusted users.
|
||||
func (h *OpenAIGatewayHandler) findBlockedCyberSessionForAPIKey(c *gin.Context, apiKey *service.APIKey, body []byte) string {
|
||||
if apiKey == nil || h.cyberPolicyLogOnly(c, apiKey) {
|
||||
return ""
|
||||
}
|
||||
return findBlockedCyberSessionKey(c.Request.Context(), h.gatewayService, apiKey.ID, c, body)
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestCyberAllowlistedUserBypassesExistingBlocksAndContinuesWebSocket(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := coderws.Accept(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.CloseNow() }()
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
|
||||
defer cancel()
|
||||
for _, response := range []string{
|
||||
`{"type":"response.failed","response":{"id":"resp_cyber","model":"gpt-5.1","error":{"code":"cyber_policy","message":"blocked by upstream"}}}`,
|
||||
`{"type":"response.completed","response":{"id":"resp_ok","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`,
|
||||
} {
|
||||
if _, _, err := conn.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
if err := conn.Write(ctx, coderws.MessageText, []byte(response)); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
// Let the downstream close so the completed frame can drain normally.
|
||||
_, _, _ = conn.Read(ctx)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
harness := newOpenAIWSPassthroughHandlerHarness(t, upstream.URL, map[string]string{
|
||||
service.SettingKeyCyberPolicyUserAllowlist: "1751",
|
||||
})
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5.1","prompt_cache_key":"trusted-session","input":"test"}`)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(string(payload)))
|
||||
explicitKey := service.CyberSessionExplicitBlockKey(harness.apiKey.ID, c, payload)
|
||||
store, ok := harness.gatewayCache.(service.CyberSessionBlockStore)
|
||||
require.True(t, ok)
|
||||
require.NoError(t, store.SetCyberSessionBlocked(context.Background(), "", []string{explicitKey}, time.Minute))
|
||||
for _, format := range []cyberSessionBlockFormat{cyberBlockFormatResponses, cyberBlockFormatChat, cyberBlockFormatAnthropic} {
|
||||
require.False(t, harness.handler.rejectIfCyberSessionBlocked(c, harness.apiKey, payload, "gpt-5.1", format))
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
require.NoError(t, harness.clientConn.Write(ctx, coderws.MessageText, payload))
|
||||
_, event, err := harness.clientConn.Read(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "cyber_policy", gjson.GetBytes(event, "response.error.code").String())
|
||||
require.Eventually(t, func() bool {
|
||||
logs := harness.moderationRepo.logSnapshot()
|
||||
return len(logs) == 1 && logs[0].Mode == service.ContentModerationModeCyberLogOnly
|
||||
}, 3*time.Second, 10*time.Millisecond)
|
||||
matched, err := store.FindCyberSessionBlocked(ctx, service.CyberSessionTranscriptBlockKeys(harness.apiKey.ID, payload))
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, matched, "the cyber response must not write new transcript blocks")
|
||||
require.NoError(t, harness.clientConn.Write(ctx, coderws.MessageText, payload))
|
||||
_, event, err = harness.clientConn.Read(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
|
||||
_ = harness.clientConn.CloseNow()
|
||||
select {
|
||||
case <-harness.handlerDone:
|
||||
case <-ctx.Done():
|
||||
t.Fatal("handler did not exit")
|
||||
}
|
||||
}
|
||||
@@ -108,6 +108,18 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: reqModel, BodyBytes: len(body), MaxTokens: 1, Kind: service.InflightEstimateToken})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
profitVetoCount := 0
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
|
||||
@@ -34,6 +34,7 @@ import (
|
||||
|
||||
// OpenAIGatewayHandler handles OpenAI API gateway requests
|
||||
type OpenAIGatewayHandler struct {
|
||||
compositeResolver *service.CompositeRouteResolver
|
||||
gatewayService *service.OpenAIGatewayService
|
||||
billingCacheService *service.BillingCacheService
|
||||
apiKeyService *service.APIKeyService
|
||||
@@ -257,13 +258,21 @@ func usageRecordContext(parent context.Context, base context.Context) context.Co
|
||||
return base
|
||||
}
|
||||
|
||||
func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) service.UsageRecordTask {
|
||||
// wrapUsageRecordTaskContext 包装计费任务:复制请求级 context 值,并接管请求的在途余额预留引用。
|
||||
// 返回的 abandon 在任务未被执行(被丢弃)时必须调用以归还预留引用;任务执行结束时(含 panic)
|
||||
// 自动归还,此时余额缓存已在计费路径中同步扣减。
|
||||
func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) (service.UsageRecordTask, func()) {
|
||||
if task == nil {
|
||||
return nil
|
||||
return nil, func() {}
|
||||
}
|
||||
done := func() {}
|
||||
if parent != nil {
|
||||
done = service.InflightReservationFromContext(parent).Acquire()
|
||||
}
|
||||
return func(ctx context.Context) {
|
||||
defer done()
|
||||
task(usageRecordContext(parent, ctx))
|
||||
}
|
||||
}, done
|
||||
}
|
||||
|
||||
func openAICompatibleRequestPlatform(ctx context.Context, apiKey *service.APIKey) string {
|
||||
@@ -606,6 +615,19 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("openai.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// Generate session hash (header first; fallback to prompt_cache_key)
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c, sessionHashBody)
|
||||
if h.rejectIfCyberSessionBlocked(c, apiKey, sessionHashBody, reqModel, cyberBlockFormatResponses) {
|
||||
@@ -1242,6 +1264,18 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.anthropicStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c, body)
|
||||
promptCacheKey := h.gatewayService.ExtractSessionID(c, body)
|
||||
sessionHash, promptCacheKey = resolveOpenAIMessagesMetadataSession(c, sessionHash, promptCacheKey, reqModel, body)
|
||||
@@ -2391,7 +2425,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
// 分组级模型白名单:首帧校验客户端模型,不通过则关闭连接并标记运维原因。
|
||||
// 必须在 ensureCompositeTargetPlatform(合成路由改写)之前执行。
|
||||
// 必须在合成路由解析和上游模型映射之前执行。
|
||||
// 与 HTTP 准入一致:帧内重复 model 键/大小写变体可能被上游按末值绑定,
|
||||
// 全部候选值逐一校验,任一未命中即拒绝。
|
||||
if blocked := blockedModelAllowlistCandidate(apiKey.Group, requestmodel.FromBodyCandidates("", "application/json", firstMessage)); blocked != "" {
|
||||
@@ -2400,7 +2434,21 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, fmt.Sprintf("Model %q is not available for this group", blocked))
|
||||
return
|
||||
}
|
||||
ensureCompositeTargetPlatform(c, apiKey, reqModel)
|
||||
// Keep the client model in the frame for admission, response identity and usage.
|
||||
// Apply the resolved upstream model through MapRequestModel on every turn.
|
||||
wsRouteModel := reqModel
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
decision, resolveErr := h.compositeResolver.Resolve(c.Request.Context(), apiKey.Group.ID, reqModel, service.CompositeRouteEndpointResponses)
|
||||
if resolveErr != nil {
|
||||
reqLog.Error("openai.websocket_composite_route_failed", zap.Error(resolveErr))
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "Failed to resolve composite model route")
|
||||
return
|
||||
}
|
||||
if decision.Matched {
|
||||
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
|
||||
wsRouteModel = decision.UpstreamModel
|
||||
}
|
||||
}
|
||||
ctx = c.Request.Context()
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
platform, ok := service.ResolvedTargetPlatformFromContext(ctx)
|
||||
@@ -2440,7 +2488,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
|
||||
// The first response.create frame is available here, so explicit IDs are
|
||||
// checked directly and body-derived sessions use the coarse scope gate.
|
||||
if cyberBlockKey := findBlockedCyberSessionKey(c.Request.Context(), h.gatewayService, apiKey.ID, c, firstMessage); cyberBlockKey != "" {
|
||||
if cyberBlockKey := h.findBlockedCyberSessionForAPIKey(c, apiKey, firstMessage); cyberBlockKey != "" {
|
||||
writeCyberSessionBlockedWSError(c.Request.Context(), wsConn)
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "session blocked by cyber-security policy")
|
||||
h.enqueueCyberSessionBlockedOpsEntry(c, apiKey, reqModel, cyberBlockKey)
|
||||
@@ -2464,8 +2512,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 解析渠道级模型映射
|
||||
channelMappingWS, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, reqModel)
|
||||
wsForwardModel := openAIChannelForwardModel(channelMappingWS, reqModel)
|
||||
channelMappingWS, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, wsRouteModel)
|
||||
wsForwardModel := openAIChannelForwardModel(channelMappingWS, wsRouteModel)
|
||||
|
||||
var currentUserRelease func()
|
||||
var currentAccountRelease func()
|
||||
@@ -2526,6 +2574,17 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(会话级):按首帧估算一次,会话期间续期;每轮计费任务接管引用,
|
||||
// 会话结束且所有轮次扣减落地后释放。
|
||||
inflightCtx, inflightDone, inflightErr := reserveInflightBalanceCtx(ctx, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, firstMessage))
|
||||
if inflightErr != nil {
|
||||
reqLog.Info("openai.websocket_inflight_reservation_rejected", zap.Error(inflightErr))
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "billing check failed")
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
ctx = inflightCtx
|
||||
|
||||
// A WebSocket may outlive a key's remaining spending window. Recheck
|
||||
// after acquiring turn slots, including the first account-selection wait.
|
||||
// Restrict this extra check to the opt-in mode so standard-mode RPM checks
|
||||
@@ -2817,7 +2876,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// 连接级 cyber session gate 也在 BeforeRequest 先执行,使 native 与
|
||||
// passthrough ingress 都能在 BeforeTurn 及上游写入前无副作用地拒绝。
|
||||
// BeforeTurn 中保留同一检查作为防御式兜底。
|
||||
if cyberBlockedThisConn {
|
||||
if cyberBlockedThisConn && !h.cyberPolicyLogOnly(c, apiKey) {
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, cyberSessionBlockedClientMsg, nil)
|
||||
}
|
||||
if turn == 1 {
|
||||
@@ -2856,7 +2915,16 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
model = reqModel
|
||||
}
|
||||
setOpsRequestContext(c, model, true)
|
||||
mapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, model)
|
||||
routeModel := model
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
// The account, target platform and route context are connection-scoped.
|
||||
// A different public model needs a new connection and fresh resolution.
|
||||
if model != reqModel {
|
||||
return "", newOpenAIWSUnsupportedModelSwitchError(model)
|
||||
}
|
||||
routeModel = wsRouteModel
|
||||
}
|
||||
mapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, routeModel)
|
||||
mappedModelUnchanged := false
|
||||
if previous := turnChannelMapping.Load(); previous != nil && previous.turn < turn {
|
||||
mappedModelUnchanged = strings.TrimSpace(previous.mapping.MappedModel) == strings.TrimSpace(mapping.MappedModel)
|
||||
@@ -2869,7 +2937,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
},
|
||||
BeforeTurn: func(turn int) error {
|
||||
// turn==1 的会话屏蔽已由握手层检查覆盖;连接内 flag 只拦截后续 turn。
|
||||
if cyberBlockedThisConn {
|
||||
if cyberBlockedThisConn && !h.cyberPolicyLogOnly(c, apiKey) {
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, cyberSessionBlockedClientMsg, nil)
|
||||
}
|
||||
// 长连接跨峰谷/倍率刷新防护:每个 turn 按当前时刻重装门并复核
|
||||
@@ -2949,7 +3017,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
cyberBlockedThisConn, cyberBlockPendingAfterFailover = advanceOpenAIWSCyberBlockState(
|
||||
cyberBlockedThisConn,
|
||||
cyberBlockPendingAfterFailover,
|
||||
cyberMarked,
|
||||
cyberMarked && !h.cyberPolicyLogOnly(c, apiKey),
|
||||
turnErr,
|
||||
)
|
||||
if turnErr != nil {
|
||||
@@ -3266,9 +3334,12 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(parent context.Context, tas
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, abandon := wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped {
|
||||
if mode.Dropped() {
|
||||
abandon()
|
||||
}
|
||||
return
|
||||
}
|
||||
// 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。
|
||||
@@ -3305,7 +3376,7 @@ func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(parent context.Con
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, _ = wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() {
|
||||
return
|
||||
@@ -4045,7 +4116,7 @@ func (h *OpenAIGatewayHandler) rejectIfCyberSessionBlocked(c *gin.Context, apiKe
|
||||
if enabled, _ := h.gatewayService.CyberSessionBlockRuntime(c.Request.Context()); !enabled {
|
||||
return false
|
||||
}
|
||||
key := findBlockedCyberSessionKey(c.Request.Context(), h.gatewayService, apiKey.ID, c, body)
|
||||
key := h.findBlockedCyberSessionForAPIKey(c, apiKey, body)
|
||||
if key == "" {
|
||||
return false
|
||||
}
|
||||
@@ -4242,7 +4313,8 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
|
||||
ClientIP: clientIPStr,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if gwSvc != nil && apiKey != nil {
|
||||
cyberLogOnly := h.cyberPolicyLogOnly(c, apiKey)
|
||||
if gwSvc != nil && apiKey != nil && !cyberLogOnly {
|
||||
plan := buildCyberSessionBlockWritePlan(apiKey.ID, c, cyberBlockBody)
|
||||
if len(plan.keys) > 0 {
|
||||
blockCtx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
|
||||
@@ -4255,6 +4327,7 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
|
||||
defer cancel()
|
||||
if cmSvc != nil {
|
||||
cmSvc.RecordCyberPolicyEvent(ctx, service.CyberPolicyRecordInput{
|
||||
LogOnly: cyberLogOnly,
|
||||
RequestID: requestID,
|
||||
UserID: userID,
|
||||
UserEmail: userEmail,
|
||||
|
||||
@@ -1931,7 +1931,10 @@ func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject
|
||||
|
||||
type openAIResponsesWSUsageLogCase struct {
|
||||
simpleModeRejectAtRead int64
|
||||
compositeResolver *service.CompositeRouteResolver
|
||||
accountPlatform string
|
||||
closeReason string
|
||||
closeStatus coderws.StatusCode
|
||||
firstPayload string
|
||||
// midPayload 在首个 turn 完成后发送(如 session.update),上游桩会为它
|
||||
// 回一个 response.completed,客户端按普通事件读取。
|
||||
@@ -2878,6 +2881,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
upstreamErrCh := make(chan error, 1)
|
||||
var channelSvc *service.ChannelService
|
||||
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if tc.accountPlatform == service.PlatformGrok {
|
||||
payload, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
upstreamErrCh <- err
|
||||
return
|
||||
}
|
||||
upstreamPayloadCh <- payload
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = fmt.Fprintf(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_grok_test\",\"model\":%q,\"usage\":{\"input_tokens\":2,\"output_tokens\":1}}}\n\n", gjson.GetBytes(payload, "model").String())
|
||||
return
|
||||
}
|
||||
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
|
||||
CompressionMode: coderws.CompressionContextTakeover,
|
||||
})
|
||||
@@ -2945,6 +2959,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||||
},
|
||||
}
|
||||
if tc.accountPlatform != "" {
|
||||
account.Platform = tc.accountPlatform
|
||||
}
|
||||
if strings.TrimSpace(tc.ingressMode) != "" {
|
||||
account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
|
||||
}
|
||||
@@ -3000,7 +3017,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
service.NewBillingService(cfg, nil),
|
||||
nil,
|
||||
billingCacheSvc,
|
||||
nil,
|
||||
&compositeWSHTTPUpstream{},
|
||||
&service.DeferredService{},
|
||||
nil,
|
||||
nil,
|
||||
@@ -3021,6 +3038,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
}
|
||||
h := &OpenAIGatewayHandler{
|
||||
cfg: cfg,
|
||||
compositeResolver: tc.compositeResolver,
|
||||
gatewayService: gatewaySvc,
|
||||
billingCacheService: billingCacheSvc,
|
||||
apiKeyService: &service.APIKeyService{},
|
||||
@@ -3069,6 +3087,12 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
|
||||
if tc.closeReason == "" {
|
||||
tc.closeReason = "not available for this group"
|
||||
}
|
||||
if tc.closeStatus == 0 {
|
||||
tc.closeStatus = coderws.StatusPolicyViolation
|
||||
}
|
||||
if tc.firstFrameCloseExpected {
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
_, _, readErr := clientConn.Read(readCtx)
|
||||
@@ -3076,12 +3100,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
require.Error(t, readErr, "first frame should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
status := tc.closeStatus
|
||||
if status == 0 {
|
||||
status = coderws.StatusPolicyViolation
|
||||
}
|
||||
require.Equal(t, status, closeErr.Code)
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
require.Empty(t, upstreamPayloadCh, "rejected first frame must not reach upstream")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3115,12 +3144,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
require.Error(t, readErr, "second turn should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
status := tc.closeStatus
|
||||
if status == 0 {
|
||||
status = coderws.StatusPolicyViolation
|
||||
}
|
||||
require.Equal(t, status, closeErr.Code)
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
require.Len(t, upstreamPayloadCh, turnCount-1, "rejected turn must not reach upstream")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3149,6 +3183,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
}
|
||||
}
|
||||
|
||||
if tc.accountPlatform == service.PlatformGrok {
|
||||
upstreamErrCh <- nil
|
||||
}
|
||||
select {
|
||||
case upstreamErr := <-upstreamErrCh:
|
||||
require.NoError(t, upstreamErr)
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type compositeWSRouteRepo struct {
|
||||
service.CompositeModelRouteRepository
|
||||
routes []service.CompositeModelRoute
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *compositeWSRouteRepo) ListByGroup(context.Context, int64, bool) ([]service.CompositeModelRoute, error) {
|
||||
return r.routes, r.err
|
||||
}
|
||||
|
||||
type compositeWSHTTPUpstream struct{ service.HTTPUpstream }
|
||||
|
||||
func (*compositeWSHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
return http.DefaultClient.Do(req)
|
||||
}
|
||||
|
||||
func compositeWSResolver(platform, endpoint, upstream string) *service.CompositeRouteResolver {
|
||||
return service.NewCompositeRouteResolver(&compositeWSRouteRepo{routes: []service.CompositeModelRoute{{
|
||||
GroupID: 4201, PublicModel: "public-alias", MatchType: service.CompositeRouteMatchExact,
|
||||
TargetPlatform: platform, Endpoint: endpoint, UpstreamModel: upstream, Enabled: true,
|
||||
}}})
|
||||
}
|
||||
func compositeWSGroup(models ...string) *service.Group {
|
||||
g := wsAllowlistGroup(len(models) > 0, models...)
|
||||
g.Platform = service.PlatformComposite
|
||||
return g
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeAlias(t *testing.T) {
|
||||
for _, platform := range []string{service.PlatformOpenAI, service.PlatformGrok} {
|
||||
for _, endpoint := range []string{service.CompositeRouteEndpointResponses, service.CompositeRouteEndpointAny} {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
t.Run(platform+"/"+endpoint+"/"+mode, func(t *testing.T) {
|
||||
upstream := "gpt-5.4"
|
||||
if platform == service.PlatformGrok {
|
||||
upstream = "grok-4.3"
|
||||
}
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias","input":"hi"}`,
|
||||
secondPayload: `{"type":"response.create","input":"again"}`,
|
||||
group: compositeWSGroup("public-alias"), accountPlatform: platform, ingressMode: mode,
|
||||
compositeResolver: compositeWSResolver(platform, endpoint, upstream),
|
||||
})
|
||||
require.Len(t, got.upstreamPayloads, 2)
|
||||
for i, payload := range got.upstreamPayloads {
|
||||
require.Equal(t, upstream, gjson.GetBytes(payload, "model").String())
|
||||
require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[i], "response.model").String())
|
||||
require.Equal(t, "public-alias", got.logs[i].RequestedModel)
|
||||
require.NotNil(t, got.logs[i].UpstreamModel)
|
||||
require.Equal(t, upstream, *got.logs[i].UpstreamModel)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeRouteRejections(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, platform, endpoint, model, reason string
|
||||
repoErr error
|
||||
status coderws.StatusCode
|
||||
}{
|
||||
{name: "disallowed platform", platform: service.PlatformAnthropic, endpoint: "responses", model: "public-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "wrong endpoint", platform: service.PlatformOpenAI, endpoint: "messages", model: "public-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "unknown alias", platform: service.PlatformOpenAI, endpoint: "responses", model: "unknown-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "resolver error", model: "gpt-5.4", repoErr: errors.New("database unavailable"), reason: "Failed to resolve composite model route", status: coderws.StatusInternalError},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
resolver := compositeWSResolver(tc.platform, tc.endpoint, "gpt-5.4")
|
||||
if tc.repoErr != nil {
|
||||
resolver = service.NewCompositeRouteResolver(&compositeWSRouteRepo{err: tc.repoErr})
|
||||
}
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"` + tc.model + `"}`, group: compositeWSGroup(), compositeResolver: resolver,
|
||||
firstFrameCloseExpected: true, closeReason: tc.reason, closeStatus: tc.status,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeAdmissionUsesPublicModel(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`, group: compositeWSGroup("gpt-5.4"),
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"), firstFrameCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeModelSwitchRequiresReconnect(t *testing.T) {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
for _, model := range []string{"gpt-5.4", "grok-4.3", "unknown-alias"} {
|
||||
t.Run(mode+"/"+model, func(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`,
|
||||
secondPayload: `{"type":"response.create","model":"` + model + `"}`, group: compositeWSGroup(), ingressMode: mode,
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"),
|
||||
secondTurnCloseExpected: true, closeReason: "model switch requires reconnect",
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeChannelBilling(t *testing.T) {
|
||||
for _, source := range []string{service.BillingModelSourceRequested, service.BillingModelSourceChannelMapped, service.BillingModelSourceUpstream} {
|
||||
t.Run(source, func(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.6-sol"}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-5.6-sol"}`,
|
||||
group: compositeWSGroup("gpt-5.6-sol"),
|
||||
compositeResolver: service.NewCompositeRouteResolver(&compositeWSRouteRepo{routes: []service.CompositeModelRoute{{PublicModel: "gpt-5.6-sol", MatchType: service.CompositeRouteMatchExact, TargetPlatform: service.PlatformOpenAI, Endpoint: service.CompositeRouteEndpointResponses, UpstreamModel: "route-target"}}}),
|
||||
channelMapping: map[string]string{"route-target": "gpt-5.4"}, billingModelSource: source,
|
||||
accountModelMapping: map[string]any{"gpt-5.4": "gpt-5.4"},
|
||||
})
|
||||
for i, log := range got.logs {
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(got.upstreamPayloads[i], "model").String())
|
||||
require.Equal(t, "gpt-5.6-sol", log.RequestedModel)
|
||||
require.Equal(t, "gpt-5.6-sol", log.Model)
|
||||
if source == service.BillingModelSourceRequested {
|
||||
require.InDelta(t, 40e-6, log.TotalCost, 1e-12)
|
||||
} else {
|
||||
require.InDelta(t, 20e-6, log.TotalCost, 1e-12)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeDetectorFallback(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4"}`, group: compositeWSGroup(),
|
||||
compositeResolver: service.NewCompositeRouteResolver(&compositeWSRouteRepo{}),
|
||||
})
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(got.upstreamFirstPayload, "model").String())
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeSessionModelSwitchRequiresReconnect(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`,
|
||||
midPayload: `{"type":"session.update","session":{"model":"grok-4.3"}}`,
|
||||
secondPayload: `{"type":"response.create"}`, group: compositeWSGroup(),
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"),
|
||||
secondTurnCloseExpected: true, closeReason: "model switch requires reconnect",
|
||||
})
|
||||
}
|
||||
@@ -143,6 +143,18 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: routingModel, BodyBytes: len(body), Kind: service.InflightEstimateImage, Units: parsed.N})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, body)
|
||||
requestCtx := service.WithOpenAIImagesEndpoint(service.WithOpenAIImageGenerationIntent(c.Request.Context()))
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
)
|
||||
|
||||
type openAIWSPassthroughHandlerHarness struct {
|
||||
handler *OpenAIGatewayHandler
|
||||
clientConn *coderws.Conn
|
||||
handlerDone <-chan struct{}
|
||||
moderationRepo *contentModerationHandlerTestRepo
|
||||
@@ -26,7 +27,7 @@ type openAIWSPassthroughHandlerHarness struct {
|
||||
apiKey *service.APIKey
|
||||
}
|
||||
|
||||
func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *openAIWSPassthroughHandlerHarness {
|
||||
func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string, settings ...map[string]string) *openAIWSPassthroughHandlerHarness {
|
||||
t.Helper()
|
||||
gatewayCache := testutil.NewRedisGatewayCache(t)
|
||||
|
||||
@@ -35,6 +36,11 @@ func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *ope
|
||||
service.SettingKeyCyberSessionBlockEnabled: "true",
|
||||
service.SettingKeyCyberSessionBlockTTLSeconds: "60",
|
||||
}}
|
||||
for _, overrides := range settings {
|
||||
for key, value := range overrides {
|
||||
settingRepo.values[key] = value
|
||||
}
|
||||
}
|
||||
moderationRepo := &contentModerationHandlerTestRepo{}
|
||||
moderationSvc := service.NewContentModerationService(settingRepo, moderationRepo, nil, nil, nil, nil, nil, nil)
|
||||
settingSvc := service.NewSettingService(settingRepo, nil)
|
||||
@@ -90,6 +96,7 @@ func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *ope
|
||||
|
||||
apiKey := &service.APIKey{
|
||||
ID: 1851,
|
||||
UserID: 1751,
|
||||
Name: "ws-cyber-key",
|
||||
Key: "sk-handler-cyber-test",
|
||||
GroupID: &groupID,
|
||||
@@ -116,6 +123,7 @@ func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *ope
|
||||
t.Cleanup(func() { _ = clientConn.CloseNow() })
|
||||
|
||||
return &openAIWSPassthroughHandlerHarness{
|
||||
handler: h,
|
||||
clientConn: clientConn,
|
||||
handlerDone: handlerDone,
|
||||
moderationRepo: moderationRepo,
|
||||
|
||||
@@ -149,6 +149,9 @@ func (h *PaymentHandler) GetCheckoutInfo(c *gin.Context) {
|
||||
BalanceRechargeMultiplier: cfg.BalanceRechargeMultiplier,
|
||||
SubscriptionUSDToCNYRate: cfg.SubscriptionUSDToCNYRate,
|
||||
RechargeFeeRate: cfg.RechargeFeeRate,
|
||||
RechargeBonusTiers: cfg.RechargeBonusTiers,
|
||||
RechargeBonusMode: cfg.RechargeBonusMode,
|
||||
RechargeBonusNotice: cfg.RechargeBonusNotice,
|
||||
HelpText: cfg.HelpText,
|
||||
HelpImageURL: cfg.HelpImageURL,
|
||||
StripePublishableKey: cfg.StripePublishableKey,
|
||||
@@ -166,6 +169,9 @@ type checkoutInfoResponse struct {
|
||||
BalanceRechargeMultiplier float64 `json:"balance_recharge_multiplier"`
|
||||
SubscriptionUSDToCNYRate float64 `json:"subscription_usd_to_cny_rate"`
|
||||
RechargeFeeRate float64 `json:"recharge_fee_rate"`
|
||||
RechargeBonusTiers []service.RechargeBonusTier `json:"recharge_bonus_tiers"`
|
||||
RechargeBonusMode string `json:"recharge_bonus_mode"`
|
||||
RechargeBonusNotice string `json:"recharge_bonus_notice"`
|
||||
HelpText string `json:"help_text"`
|
||||
HelpImageURL string `json:"help_image_url"`
|
||||
StripePublishableKey string `json:"stripe_publishable_key"`
|
||||
@@ -482,6 +488,7 @@ type PublicOrderResult struct {
|
||||
Amount float64 `json:"amount"`
|
||||
PayAmount float64 `json:"pay_amount"`
|
||||
FeeRate float64 `json:"fee_rate"`
|
||||
BonusAmount float64 `json:"bonus_amount"`
|
||||
Currency string `json:"currency"`
|
||||
PaymentType string `json:"payment_type"`
|
||||
OrderType string `json:"order_type"`
|
||||
@@ -517,6 +524,7 @@ func buildPublicOrderResult(order *dbent.PaymentOrder) PublicOrderResult {
|
||||
Amount: order.Amount,
|
||||
PayAmount: order.PayAmount,
|
||||
FeeRate: order.FeeRate,
|
||||
BonusAmount: order.BonusAmount,
|
||||
Currency: service.PaymentOrderCurrency(order),
|
||||
PaymentType: order.PaymentType,
|
||||
OrderType: order.OrderType,
|
||||
@@ -625,6 +633,7 @@ type PaymentOrderResult struct {
|
||||
Amount float64 `json:"amount"`
|
||||
PayAmount float64 `json:"pay_amount"`
|
||||
FeeRate float64 `json:"fee_rate"`
|
||||
BonusAmount float64 `json:"bonus_amount"`
|
||||
Currency string `json:"currency"`
|
||||
PaymentType string `json:"payment_type"`
|
||||
OutTradeNo string `json:"out_trade_no"`
|
||||
@@ -663,6 +672,7 @@ func sanitizePaymentOrderForResponse(order *dbent.PaymentOrder) *PaymentOrderRes
|
||||
Amount: order.Amount,
|
||||
PayAmount: order.PayAmount,
|
||||
FeeRate: order.FeeRate,
|
||||
BonusAmount: order.BonusAmount,
|
||||
Currency: service.PaymentOrderCurrency(order),
|
||||
PaymentType: order.PaymentType,
|
||||
OutTradeNo: order.OutTradeNo,
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
@@ -811,3 +812,19 @@ func TestUserHandlerStartIdentityBindingReturnsAuthorizeURL(t *testing.T) {
|
||||
require.Contains(t, resp.Data.AuthorizeURL, "intent=bind_current_user")
|
||||
require.Contains(t, resp.Data.AuthorizeURL, "redirect=%2Fsettings%2Fprofile")
|
||||
}
|
||||
|
||||
func (s *userHandlerEmailCacheStub) IncrVerificationCodeAttempts(context.Context, string) (int, error) {
|
||||
if s.data == nil {
|
||||
return 0, errors.New("verification code not found")
|
||||
}
|
||||
s.data.Attempts++
|
||||
return s.data.Attempts, nil
|
||||
}
|
||||
|
||||
func (s *userHandlerEmailCacheStub) IncrNotifyVerifyCodeAttempts(context.Context, string) (int, error) {
|
||||
return 0, errors.New("notify verification code not found")
|
||||
}
|
||||
|
||||
func (s *userHandlerEmailCacheStub) ConsumePasswordResetToken(context.Context, string, string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
@@ -50,10 +50,12 @@ func ProvideAdminHandlers(
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService,
|
||||
ollamaCloudUsage *service.OllamaCloudUsageService,
|
||||
opencodeGoUsage *service.OpenCodeGoUsageService,
|
||||
claudeResetCredits *service.ClaudeResetCreditService,
|
||||
) *AdminHandlers {
|
||||
accountHandler.SetUpstreamBillingProbeService(upstreamBillingProbe)
|
||||
accountHandler.SetOllamaCloudUsageService(ollamaCloudUsage)
|
||||
accountHandler.SetOpenCodeGoUsageService(opencodeGoUsage)
|
||||
accountHandler.SetClaudeResetCreditService(claudeResetCredits)
|
||||
return &AdminHandlers{
|
||||
Dashboard: dashboardHandler,
|
||||
User: userHandler,
|
||||
@@ -132,10 +134,12 @@ func ProvideOpenAIGatewayHandler(
|
||||
grokQuotaService *service.GrokQuotaService,
|
||||
cfg *config.Config,
|
||||
coordinator *securityaudit.Coordinator,
|
||||
compositeResolver *service.CompositeRouteResolver,
|
||||
) *OpenAIGatewayHandler {
|
||||
gatewayService.SetPluginManager(pluginManager)
|
||||
h := NewOpenAIGatewayHandler(gatewayService, concurrencyService, billingCacheService, apiKeyService,
|
||||
usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, cfg)
|
||||
h.compositeResolver = compositeResolver
|
||||
h.securityAuditCoordinator = coordinator
|
||||
h.grokMediaEligibilityProber = grokQuotaService
|
||||
return h
|
||||
|
||||
@@ -46,6 +46,7 @@ const (
|
||||
PlatformDeepseek = domain.PlatformDeepseek
|
||||
PlatformMiniMax = domain.PlatformMiniMax
|
||||
PlatformOpenCodeGo = domain.PlatformOpenCodeGo
|
||||
PlatformTypeSafe = domain.PlatformTypeSafe
|
||||
)
|
||||
|
||||
// AllPlatforms 返回所有支持的平台列表
|
||||
@@ -61,6 +62,7 @@ func AllPlatforms() []string {
|
||||
PlatformDeepseek,
|
||||
PlatformMiniMax,
|
||||
PlatformOpenCodeGo,
|
||||
PlatformTypeSafe,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -18,5 +18,6 @@ func TestAllPlatformsIncludesEveryConcretePlatform(t *testing.T) {
|
||||
"deepseek",
|
||||
"minimax",
|
||||
"opencode_go",
|
||||
"typesafe",
|
||||
}, AllPlatforms())
|
||||
}
|
||||
|
||||
@@ -1882,6 +1882,64 @@ func TestOpus55SignedThinkingResponsesRoundTrip(t *testing.T) {
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSonnet55ResponsesThinkingAndSampling(t *testing.T) {
|
||||
for _, effort := range []string{"", "low", "medium", "high", "xhigh", "max", "none"} {
|
||||
req := &ResponsesRequest{Model: "claude-sonnet-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, "between_tools", out.Thinking.Type)
|
||||
require.Equal(t, "low", out.OutputConfig.Effort)
|
||||
} else {
|
||||
require.Equal(t, "adaptive", out.Thinking.Type)
|
||||
if effort == "" {
|
||||
effort = "high"
|
||||
}
|
||||
require.Equal(t, effort, out.OutputConfig.Effort)
|
||||
}
|
||||
}
|
||||
for _, choice := range []string{`"required"`, `{"type":"function","name":"lookup"}`} {
|
||||
_, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), ToolChoice: json.RawMessage(choice)})
|
||||
require.ErrorContains(t, err, "forced tool_choice")
|
||||
}
|
||||
temperature := 0.7
|
||||
_, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Temperature: &temperature})
|
||||
require.ErrorContains(t, err, "temperature")
|
||||
topP := 0.5
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), TopP: &topP})
|
||||
require.ErrorContains(t, err, "top_p")
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: "minimal"}})
|
||||
require.ErrorContains(t, err, "reasoning effort")
|
||||
|
||||
temperature, topP = 1, 0.99
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Temperature: &temperature, TopP: &topP})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestSonnet55SignedThinkingResponsesRoundTrip(t *testing.T) {
|
||||
block := AnthropicContentBlock{Type: "thinking", Thinking: "", Signature: "signed-sonnet-block"}
|
||||
response := AnthropicToResponsesResponse(&AnthropicResponse{Model: "claude-sonnet-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-sonnet-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"} {
|
||||
@@ -1921,3 +1979,16 @@ func TestMessageStartSSE_StopReasonIsJSONNull(t *testing.T) {
|
||||
require.Contains(t, sse, `"stop_reason":null`)
|
||||
require.NotContains(t, sse, `"stop_reason":""`)
|
||||
}
|
||||
|
||||
func TestGPT61SolCacheOptionsAndBreakpointsSurviveChatBridge(t *testing.T) {
|
||||
sampling := 0.7
|
||||
for _, effort := range []string{"low", "medium", "high", "xhigh", "max"} {
|
||||
out, err := ChatCompletionsToResponses(&ChatCompletionsRequest{Model: "gpt-6.1-sol", ReasoningEffort: effort, Temperature: &sampling, TopP: &sampling, PromptCacheOptions: json.RawMessage(`{"ttl":"30m","mode":"explicit"}`), Messages: []ChatMessage{{Role: "user", Content: json.RawMessage(`[{"type":"text","text":"prefix","prompt_cache_breakpoint":{"mode":"explicit"}}]`)}}})
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, out.Temperature)
|
||||
require.Nil(t, out.TopP)
|
||||
require.Equal(t, effort, out.Reasoning.Effort)
|
||||
require.JSONEq(t, `{"ttl":"30m","mode":"explicit"}`, string(out.PromptCacheOptions))
|
||||
require.Contains(t, string(out.Input), "prompt_cache_breakpoint")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,9 @@ import (
|
||||
// Chat Completions intermediary round-trip (e.g. thinking, cache_control,
|
||||
// structured system prompts).
|
||||
func AnthropicToResponses(req *AnthropicRequest) (*ResponsesRequest, error) {
|
||||
if err := openai.ValidateGPT61SolReasoningEffort(req.Model, anthropicReasoningEffort(req)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
input, err := convertAnthropicToResponsesInput(req.System, req.Messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -59,6 +62,9 @@ func AnthropicToResponses(req *AnthropicRequest) (*ResponsesRequest, error) {
|
||||
|
||||
// An explicit thinking disable takes precedence over output_config.effort.
|
||||
effort := anthropicReasoningEffort(req)
|
||||
if openai.IsGPT61SolModelSpelling(req.Model) && req.OutputConfig != nil && req.OutputConfig.Effort == "max" && effort != "none" {
|
||||
effort = "max"
|
||||
}
|
||||
out.Reasoning = &ResponsesReasoning{
|
||||
Effort: effort,
|
||||
}
|
||||
|
||||
@@ -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) && (block.Signature != "" || block.Data != "") {
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(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) && block.Text != "" {
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(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)
|
||||
state.PreserveThinkingSignatures = state.PreserveThinkingSignatures || claude.IsOpus55(evt.Message.Model) || claude.IsSonnet55(evt.Message.Model)
|
||||
if state.Model == "" {
|
||||
state.Model = evt.Message.Model
|
||||
}
|
||||
|
||||
@@ -18,6 +18,9 @@ type chatMessageContent struct {
|
||||
// true. store is always false and reasoning.encrypted_content is always
|
||||
// included so that the response translator has full context.
|
||||
func ChatCompletionsToResponses(req *ChatCompletionsRequest) (*ResponsesRequest, error) {
|
||||
if err := openai.ValidateGPT61SolReasoningEffort(req.Model, req.ReasoningEffort); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
input, err := convertChatMessagesToResponsesInput(req.Messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -14,7 +14,9 @@ import (
|
||||
// enables Anthropic platform groups to accept OpenAI Responses API requests
|
||||
// by converting them to the native /v1/messages format before forwarding upstream.
|
||||
func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, error) {
|
||||
system, messages, err := convertResponsesInputToAnthropic(req.Instructions, req.Input, claude.IsOpus55(req.Model))
|
||||
isOpus55 := claude.IsOpus55(req.Model)
|
||||
isSonnet55 := claude.IsSonnet55(req.Model)
|
||||
system, messages, err := convertResponsesInputToAnthropic(req.Instructions, req.Input, isOpus55 || isSonnet55)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -54,9 +56,11 @@ func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, erro
|
||||
out.ToolChoice = tc
|
||||
}
|
||||
|
||||
// Opus 5.5 always uses adaptive thinking. Resolve the upstream model before
|
||||
// conversion, because client aliases need not identify a Claude model.
|
||||
if claude.IsOpus55(req.Model) {
|
||||
// The 5.5 models reject manual thinking and forced tool use. Sonnet 5.5
|
||||
// additionally supports between_tools to disable up-front thinking.
|
||||
// Resolve the upstream model before conversion: client aliases need not
|
||||
// identify a Claude model.
|
||||
if isOpus55 || isSonnet55 {
|
||||
var choice struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
@@ -66,16 +70,36 @@ func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, erro
|
||||
}
|
||||
}
|
||||
if choice.Type == "any" || choice.Type == "tool" {
|
||||
return nil, fmt.Errorf("claude-opus-5-5 does not support forced tool_choice; use auto or none")
|
||||
return nil, fmt.Errorf("%s does not support forced tool_choice; use auto or none", req.Model)
|
||||
}
|
||||
effort := "medium"
|
||||
if isSonnet55 {
|
||||
effort = "high"
|
||||
if req.Temperature != nil && *req.Temperature != 1 {
|
||||
return nil, fmt.Errorf("claude-sonnet-5-5 does not support non-default temperature")
|
||||
}
|
||||
if req.TopP != nil && (*req.TopP < 0.99 || *req.TopP > 1) {
|
||||
return nil, fmt.Errorf("claude-sonnet-5-5 does not support non-default top_p")
|
||||
}
|
||||
}
|
||||
if req.Reasoning != nil && req.Reasoning.Effort != "" {
|
||||
effort = req.Reasoning.Effort
|
||||
}
|
||||
if isSonnet55 && effort == "none" {
|
||||
// OpenAI's no-reasoning request maps to Sonnet 5.5's lowest
|
||||
// thinking mode. between_tools still preserves signed progress
|
||||
// blocks produced during tool use.
|
||||
out.Thinking = &AnthropicThinking{Type: "between_tools"}
|
||||
if out.OutputConfig == nil {
|
||||
out.OutputConfig = &AnthropicOutputConfig{}
|
||||
}
|
||||
out.OutputConfig.Effort = "low"
|
||||
return out, nil
|
||||
}
|
||||
switch effort {
|
||||
case "low", "medium", "high", "xhigh", "max":
|
||||
default:
|
||||
return nil, fmt.Errorf("claude-opus-5-5 does not support reasoning effort %q; use low, medium, high, xhigh or max", effort)
|
||||
return nil, fmt.Errorf("%s does not support reasoning effort %q; use low, medium, high, xhigh or max", req.Model, effort)
|
||||
}
|
||||
out.Thinking = &AnthropicThinking{Type: "adaptive"}
|
||||
out.OutputConfig = &AnthropicOutputConfig{Effort: effort}
|
||||
|
||||
@@ -42,7 +42,7 @@ type AnthropicOutputConfig struct {
|
||||
|
||||
// AnthropicThinking configures extended thinking in the Anthropic API.
|
||||
type AnthropicThinking struct {
|
||||
Type string `json:"type"` // "enabled" | "adaptive" | "disabled"
|
||||
Type string `json:"type"` // "enabled" | "adaptive" | "disabled" | "between_tools"
|
||||
BudgetTokens int `json:"budget_tokens,omitempty"` // max thinking tokens
|
||||
}
|
||||
|
||||
|
||||
@@ -190,6 +190,12 @@ var DefaultModels = []Model{
|
||||
DisplayName: "Claude Opus 5",
|
||||
CreatedAt: "2026-07-25T00:00:00Z",
|
||||
},
|
||||
{
|
||||
ID: "claude-sonnet-5-5",
|
||||
Type: "model",
|
||||
DisplayName: "Claude Sonnet 5.5",
|
||||
CreatedAt: "2026-09-28T00:00:00Z",
|
||||
},
|
||||
{
|
||||
ID: "claude-sonnet-5",
|
||||
Type: "model",
|
||||
|
||||
@@ -30,3 +30,16 @@ func TestDefaultModelsContainsOpus55(t *testing.T) {
|
||||
}
|
||||
t.Fatal("claude-opus-5-5 missing")
|
||||
}
|
||||
|
||||
func TestDefaultModelsContainsSonnet55(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
for _, model := range DefaultModels {
|
||||
if model.ID == "claude-sonnet-5-5" {
|
||||
require.Equal(t, "Claude Sonnet 5.5", model.DisplayName)
|
||||
require.Equal(t, "2026-09-28T00:00:00Z", model.CreatedAt)
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatal("claude-sonnet-5-5 missing from DefaultModels")
|
||||
}
|
||||
|
||||
@@ -19,6 +19,7 @@ var effortFamilies = []struct {
|
||||
{family: "claude-mythos-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-fable-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-sonnet-4-6", levels: effortLowMediumHighMax},
|
||||
{family: "claude-sonnet-5-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-sonnet-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-opus-4-8", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-opus-4-7", levels: effortLowMediumHighXHighMax},
|
||||
@@ -45,19 +46,30 @@ func IsOpus55(model string) bool {
|
||||
return normalizeEffortModelID(model) == "claude-opus-5-5"
|
||||
}
|
||||
|
||||
// IsSonnet55 identifies the fixed Sonnet 5.5 ID after provider/local suffix normalization.
|
||||
func IsSonnet55(model string) bool {
|
||||
return normalizeEffortModelID(model) == "claude-sonnet-5-5"
|
||||
}
|
||||
|
||||
func normalizeEffortModelID(model string) string {
|
||||
id := strings.ToLower(strings.TrimSpace(model))
|
||||
id = strings.TrimPrefix(id, "models/")
|
||||
if slash := strings.IndexByte(id, '/'); slash >= 0 {
|
||||
id = strings.TrimPrefix(strings.TrimSpace(id[slash+1:]), "models/")
|
||||
}
|
||||
for _, prefix := range []string{"us.", "eu.", "apac.", "jp.", "au.", "us-gov.", "global."} {
|
||||
id = strings.TrimPrefix(id, prefix)
|
||||
}
|
||||
id = strings.TrimPrefix(id, "anthropic.")
|
||||
id = strings.TrimSuffix(id, "-thinking")
|
||||
// OpenRouter uses a dotted minor version for this exact Opus 5.5 ID.
|
||||
// Normalize it before effort, thinking, and billing family lookups.
|
||||
// OpenRouter uses dotted minor versions for some models. Normalize them
|
||||
// before effort, thinking, and billing family lookups.
|
||||
if id == "claude-opus-5.5" {
|
||||
id = "claude-opus-5-5"
|
||||
}
|
||||
if id == "claude-sonnet-5.5" {
|
||||
id = "claude-sonnet-5-5"
|
||||
}
|
||||
if mapped, ok := ModelIDReverseOverrides[id]; ok {
|
||||
id = mapped
|
||||
}
|
||||
|
||||
@@ -17,6 +17,8 @@ func TestEffortLevelsForModel(t *testing.T) {
|
||||
{model: "anthropic/claude-sonnet-4-6", want: []string{"low", "medium", "high", "max"}},
|
||||
{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: "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},
|
||||
{model: "gpt-5.6", want: nil},
|
||||
@@ -38,3 +40,20 @@ func TestIsOpus55OpenRouterExactAlias(t *testing.T) {
|
||||
require.False(t, IsOpus55(model), model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsSonnet55(t *testing.T) {
|
||||
t.Parallel()
|
||||
for _, model := range []string{
|
||||
"claude-sonnet-5-5",
|
||||
"anthropic/claude-sonnet-5.5",
|
||||
"anthropic.claude-sonnet-5-5",
|
||||
"us.anthropic.claude-sonnet-5-5",
|
||||
"us-gov.anthropic.claude-sonnet-5-5",
|
||||
"global.anthropic.claude-sonnet-5-5-thinking",
|
||||
} {
|
||||
require.True(t, IsSonnet55(model), model)
|
||||
}
|
||||
for _, model := range []string{"claude-sonnet-5", "claude-sonnet-5-5-preview", "claude-opus-5-5"} {
|
||||
require.False(t, IsSonnet55(model), model)
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -3,6 +3,8 @@ package openai
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
@@ -23,6 +25,7 @@ var DefaultModels = []Model{
|
||||
{ID: "gpt-5.6", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 (Sol)"},
|
||||
{ID: "gpt-5.6-terra", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Terra"},
|
||||
{ID: "gpt-5.6-luna", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Luna"},
|
||||
{ID: "gpt-6.1-sol", Object: "model", Created: 1790640000, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6.1 Sol"},
|
||||
{ID: "gpt-6-sol", Object: "model", Created: 1790035200, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6 Sol"},
|
||||
{ID: "gpt-6-luna", Object: "model", Created: 1790035200, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6 Luna"},
|
||||
{ID: "gpt-6-astra", Object: "model", Created: 1788480000, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6 Astra"},
|
||||
@@ -78,6 +81,12 @@ var instructionsGPT55 string
|
||||
//go:embed instructions_gpt6_astra.txt
|
||||
var instructionsGPT6Astra string
|
||||
|
||||
// CodexGPT61SolMetadata is the complete official descriptor from openai/codex
|
||||
// b1e72963c3b71a9265a551e54beff078384efed9, codex-rs/models-manager/models.json.
|
||||
//
|
||||
//go:embed codex_gpt61_sol.json
|
||||
var CodexGPT61SolMetadata []byte
|
||||
|
||||
// latestCodexInstructions 返回当前已知最新版本的 Codex base instructions,
|
||||
// 当前为 GPT-5.5;若 5.5 prompt 意外为空则回退到 DefaultInstructions 保证非空。
|
||||
func latestCodexInstructions() string {
|
||||
@@ -140,6 +149,16 @@ func CanonicalizeOpenAIModelAliasSpelling(model string) string {
|
||||
func CodexBaseInstructionsForModel(model string) string {
|
||||
canonical := CanonicalizeOpenAIModelAliasSpelling(model)
|
||||
switch {
|
||||
case IsGPT61SolModelSpelling(canonical):
|
||||
var metadata struct {
|
||||
ModelMessages struct {
|
||||
InstructionsTemplate string `json:"instructions_template"`
|
||||
} `json:"model_messages"`
|
||||
}
|
||||
if err := json.Unmarshal(CodexGPT61SolMetadata, &metadata); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return metadata.ModelMessages.InstructionsTemplate
|
||||
case canonical == "gpt-6" || canonical == "gpt-6-astra" || strings.HasPrefix(canonical, "gpt-6-astra-"):
|
||||
if v := strings.TrimSpace(instructionsGPT6Astra); v != "" {
|
||||
return instructionsGPT6Astra
|
||||
@@ -177,3 +196,34 @@ func IsGPT6SolOrLunaModelSpelling(model string) bool {
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// IsGPT61SolModelSpelling recognizes the published model and local effort/compact
|
||||
// spellings. Invalid effort suffixes remain identifiable for request validation.
|
||||
func IsGPT61SolModelSpelling(model string) bool {
|
||||
canonical := CanonicalizeOpenAIModelAliasSpelling(model)
|
||||
if canonical == "gpt-6.1-sol" {
|
||||
return true
|
||||
}
|
||||
suffix, ok := strings.CutPrefix(canonical, "gpt-6.1-sol-")
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
switch suffix {
|
||||
case "none", "minimal", "low", "medium", "high", "xhigh", "max", "openai-compact":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// ValidateGPT61SolReasoningEffort rejects disabled reasoning instead of silently
|
||||
// increasing the client's requested effort on compatibility paths.
|
||||
func ValidateGPT61SolReasoningEffort(model, effort string) error {
|
||||
if IsGPT61SolModelSpelling(model) {
|
||||
switch strings.ToLower(strings.TrimSpace(effort)) {
|
||||
case "none", "minimal":
|
||||
return fmt.Errorf("gpt-6.1-sol does not support reasoning effort %q; use low, medium, high, xhigh or max", effort)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -42,3 +42,22 @@ func TestGPT6SolLunaModelIdentity(t *testing.T) {
|
||||
require.False(t, IsGPT6SolOrLunaModelSpelling("gpt-6-solitude"))
|
||||
require.False(t, IsGPT6SolOrLunaModelSpelling("gpt-6-luna-preview"))
|
||||
}
|
||||
|
||||
func TestGPT61SolIdentityAndEffort(t *testing.T) {
|
||||
require.Contains(t, DefaultModelIDs(), "gpt-6.1-sol")
|
||||
for _, id := range []string{"gpt-6.1-sol", "openai/gpt-6.1-sol-max", "GPT_6.1_SOL", "gpt-6.1-sol-openai-compact"} {
|
||||
require.True(t, IsGPT61SolModelSpelling(id), id)
|
||||
require.False(t, IsGPT6SolOrLunaModelSpelling(id), id)
|
||||
}
|
||||
for _, id := range []string{"gpt-6.1", "gpt-6.1-solitude", "gpt-6.1-sol-preview", "gpt-6-sol"} {
|
||||
require.False(t, IsGPT61SolModelSpelling(id), id)
|
||||
}
|
||||
for _, effort := range []string{"", "low", "medium", "high", "xhigh", "max"} {
|
||||
require.NoError(t, ValidateGPT61SolReasoningEffort("gpt-6.1-sol", effort))
|
||||
}
|
||||
for _, effort := range []string{"none", "minimal"} {
|
||||
require.Error(t, ValidateGPT61SolReasoningEffort("gpt-6.1-sol", effort))
|
||||
require.NoError(t, ValidateGPT61SolReasoningEffort("gpt-6-sol", effort))
|
||||
}
|
||||
require.Contains(t, CodexBaseInstructionsForModel("gpt-6.1-sol"), "based on GPT-6")
|
||||
}
|
||||
|
||||
@@ -13,6 +13,12 @@ import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultBaseURL = "https://api.typesafe.ai"
|
||||
SystemOnePath = "/v1/systemone"
|
||||
JevLatestModel = "jev-latest"
|
||||
)
|
||||
|
||||
type Question struct {
|
||||
Type string `json:"type"`
|
||||
Instructions string `json:"instructions"`
|
||||
@@ -35,22 +41,103 @@ type Usage struct {
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
}
|
||||
|
||||
type SystemOneResponse struct {
|
||||
Body []byte
|
||||
Model string
|
||||
Usage Usage
|
||||
}
|
||||
|
||||
func NewSystemOneRequest(ctx context.Context, baseURL, key string, body []byte) (*http.Request, error) {
|
||||
endpoint, err := url.JoinPath(strings.TrimRight(baseURL, "/"), SystemOnePath)
|
||||
if err != nil {
|
||||
return nil, errors.New("typesafe invalid endpoint")
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, errors.New("typesafe invalid request")
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
return req, nil
|
||||
}
|
||||
|
||||
// MaxSystemOneResponseBytes bounds a buffered System One response body.
|
||||
const MaxSystemOneResponseBytes = 4 << 20
|
||||
|
||||
var ErrSystemOneResponseTooLarge = errors.New("typesafe response exceeds size limit")
|
||||
|
||||
func DecodeSystemOneResponse(r io.Reader) (*SystemOneResponse, error) {
|
||||
body, err := io.ReadAll(io.LimitReader(r, MaxSystemOneResponseBytes+1))
|
||||
if err != nil {
|
||||
return nil, errors.New("typesafe invalid response")
|
||||
}
|
||||
// A truncated body would otherwise surface as a misleading "invalid JSON".
|
||||
if len(body) > MaxSystemOneResponseBytes {
|
||||
return nil, ErrSystemOneResponseTooLarge
|
||||
}
|
||||
if !json.Valid(body) {
|
||||
return nil, errors.New("typesafe invalid response")
|
||||
}
|
||||
var envelope map[string]json.RawMessage
|
||||
if err := json.Unmarshal(body, &envelope); err != nil || envelope == nil {
|
||||
return nil, errors.New("typesafe invalid response")
|
||||
}
|
||||
// The upstream already answered (and charged); an unexpected model or usage
|
||||
// shape must not discard the answer, so both are decoded leniently.
|
||||
var model string
|
||||
_ = json.Unmarshal(envelope["model"], &model)
|
||||
var usage map[string]json.RawMessage
|
||||
_ = json.Unmarshal(envelope["usage"], &usage)
|
||||
return &SystemOneResponse{
|
||||
Body: body,
|
||||
Model: model,
|
||||
Usage: Usage{
|
||||
InputTokens: systemOneTokenCount(usage["input_tokens"]),
|
||||
OutputTokens: systemOneTokenCount(usage["output_tokens"]),
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// maxSystemOneTokenCount bounds a reported token count before int conversion.
|
||||
const maxSystemOneTokenCount = 1 << 40
|
||||
|
||||
// systemOneTokenCount accepts integer, float, or numeric-string token counts.
|
||||
func systemOneTokenCount(raw json.RawMessage) int {
|
||||
raw = bytes.TrimSpace(raw)
|
||||
if len(raw) == 0 {
|
||||
return 0
|
||||
}
|
||||
if raw[0] == '"' {
|
||||
var text string
|
||||
if json.Unmarshal(raw, &text) != nil {
|
||||
return 0
|
||||
}
|
||||
raw = json.RawMessage(strings.TrimSpace(text))
|
||||
}
|
||||
var number json.Number
|
||||
if json.Unmarshal(raw, &number) != nil {
|
||||
return 0
|
||||
}
|
||||
value, err := number.Float64()
|
||||
if err != nil || math.IsNaN(value) || value <= 0 {
|
||||
return 0
|
||||
}
|
||||
if value > maxSystemOneTokenCount {
|
||||
value = maxSystemOneTokenCount
|
||||
}
|
||||
return int(math.Round(value))
|
||||
}
|
||||
|
||||
// Evaluate performs one attempt. The caller owns timeouts, retries and key rotation.
|
||||
func Evaluate(ctx context.Context, client *http.Client, baseURL, key string, input Request) (*Result, int, error) {
|
||||
endpoint, err := url.JoinPath(strings.TrimRight(baseURL, "/"), "/v1/systemone")
|
||||
if err != nil {
|
||||
return nil, 0, errors.New("typesafe invalid endpoint")
|
||||
}
|
||||
body, err := json.Marshal(input)
|
||||
if err != nil {
|
||||
return nil, 0, errors.New("typesafe invalid request")
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
|
||||
req, err := NewSystemOneRequest(ctx, baseURL, key, body)
|
||||
if err != nil {
|
||||
return nil, 0, errors.New("typesafe invalid request")
|
||||
return nil, 0, err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
package typesafe
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var ErrStreamingUnsupported = errors.New("typesafe system one does not support streaming")
|
||||
|
||||
type systemOneEnvelope struct {
|
||||
Model json.RawMessage `json:"model"`
|
||||
State json.RawMessage `json:"state"`
|
||||
Questions json.RawMessage `json:"questions"`
|
||||
Stream json.RawMessage `json:"stream"`
|
||||
}
|
||||
|
||||
var (
|
||||
systemOneRequestFields = []string{"model", "state", "questions", "stream"}
|
||||
systemOneQuestionFields = []string{"type", "instructions", "criteria"}
|
||||
)
|
||||
|
||||
func ValidateSystemOneRequest(body []byte) (string, error) {
|
||||
var envelope systemOneEnvelope
|
||||
if err := json.Unmarshal(body, &envelope); err != nil {
|
||||
return "", errors.New("invalid JSON request")
|
||||
}
|
||||
// encoding/json matches struct fields case-insensitively and keeps the last
|
||||
// duplicate, while the raw body is forwarded upstream unchanged. Reject any
|
||||
// spelling the upstream could read differently from this validator.
|
||||
if err := checkSystemOneObjectKeys(body, "request", systemOneRequestFields); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
model, err := requiredString(envelope.Model, "model")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if model != JevLatestModel {
|
||||
return "", fmt.Errorf("model must be %s", JevLatestModel)
|
||||
}
|
||||
if err := validateStringObjectOrArray(envelope.State, "state"); err != nil {
|
||||
return "", err
|
||||
}
|
||||
questionsRaw := bytes.TrimSpace(envelope.Questions)
|
||||
var questions map[string]json.RawMessage
|
||||
if len(questionsRaw) == 0 || questionsRaw[0] != '{' || json.Unmarshal(questionsRaw, &questions) != nil || len(questions) == 0 {
|
||||
return "", errors.New("questions must be a non-empty object")
|
||||
}
|
||||
if err := checkSystemOneObjectKeys(questionsRaw, "questions", nil); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(envelope.Stream) > 0 && string(envelope.Stream) != "null" {
|
||||
var stream bool
|
||||
if err := json.Unmarshal(envelope.Stream, &stream); err != nil {
|
||||
return "", errors.New("stream must be a boolean")
|
||||
}
|
||||
if stream {
|
||||
return "", ErrStreamingUnsupported
|
||||
}
|
||||
}
|
||||
for id, raw := range questions {
|
||||
if err := validateQuestion(id, raw); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
return model, nil
|
||||
}
|
||||
|
||||
func validateQuestion(id string, raw json.RawMessage) error {
|
||||
raw = bytes.TrimSpace(raw)
|
||||
var question struct {
|
||||
Type json.RawMessage `json:"type"`
|
||||
Instructions json.RawMessage `json:"instructions"`
|
||||
Criteria json.RawMessage `json:"criteria"`
|
||||
}
|
||||
if len(raw) == 0 || raw[0] != '{' || json.Unmarshal(raw, &question) != nil {
|
||||
return fmt.Errorf("question %q must be an object", id)
|
||||
}
|
||||
if err := checkSystemOneObjectKeys(raw, fmt.Sprintf("question %q", id), systemOneQuestionFields); err != nil {
|
||||
return err
|
||||
}
|
||||
typ, err := requiredString(question.Type, "question type")
|
||||
if err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
// Mirror the wire schema generated from https://api.typesafe.ai/openapi.json
|
||||
// in TypeSafe SDK v0.5.7: instructions are optional and nullable for all types.
|
||||
if err := validateOptionalDescription(question.Instructions, "instructions"); err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
|
||||
switch typ {
|
||||
case "noul":
|
||||
criteria := bytes.TrimSpace(question.Criteria)
|
||||
if len(criteria) == 0 || bytes.Equal(criteria, []byte("null")) {
|
||||
return nil
|
||||
}
|
||||
var descriptions map[string]json.RawMessage
|
||||
if criteria[0] != '{' || json.Unmarshal(criteria, &descriptions) != nil {
|
||||
return fmt.Errorf("question %q: noul criteria must be an object", id)
|
||||
}
|
||||
for _, outcome := range []string{"true", "false"} {
|
||||
if err := validateOptionalDescription(descriptions[outcome], "noul criteria "+outcome); err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
}
|
||||
case "choice":
|
||||
var criteria map[string]json.RawMessage
|
||||
if json.Unmarshal(question.Criteria, &criteria) != nil || criteria == nil {
|
||||
return fmt.Errorf("question %q: choice criteria must be an object", id)
|
||||
}
|
||||
for _, value := range criteria {
|
||||
if err := validateOptionalDescription(value, "choice criteria value"); err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
}
|
||||
case "score":
|
||||
var criteria []json.RawMessage
|
||||
if json.Unmarshal(question.Criteria, &criteria) != nil || len(criteria) == 0 {
|
||||
return fmt.Errorf("question %q: score criteria must contain at least one level", id)
|
||||
}
|
||||
for _, value := range criteria {
|
||||
if err := validateStringObjectOrArray(value, "score criteria value"); err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("question %q: unsupported type %q", id, typ)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateOptionalDescription(raw json.RawMessage, name string) error {
|
||||
raw = bytes.TrimSpace(raw)
|
||||
if len(raw) == 0 || bytes.Equal(raw, []byte("null")) {
|
||||
return nil
|
||||
}
|
||||
return validateStringObjectOrArray(raw, name)
|
||||
}
|
||||
|
||||
func requiredString(raw json.RawMessage, name string) (string, error) {
|
||||
var value string
|
||||
if len(raw) == 0 || json.Unmarshal(raw, &value) != nil || strings.TrimSpace(value) == "" {
|
||||
return "", fmt.Errorf("%s must be a non-empty string", name)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func validateStringObjectOrArray(raw json.RawMessage, name string) error {
|
||||
raw = bytes.TrimSpace(raw)
|
||||
if len(raw) == 0 || string(raw) == "null" {
|
||||
return fmt.Errorf("%s is required", name)
|
||||
}
|
||||
switch raw[0] {
|
||||
case '{', '[':
|
||||
return nil
|
||||
case '"':
|
||||
if rawString(raw) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("%s must be a string, object, or array", name)
|
||||
}
|
||||
|
||||
func rawString(raw json.RawMessage) bool {
|
||||
var value string
|
||||
return json.Unmarshal(raw, &value) == nil
|
||||
}
|
||||
|
||||
// checkSystemOneObjectKeys rejects duplicate keys and non-canonical spellings
|
||||
// (any case variant) of the known fields in one JSON object level.
|
||||
func checkSystemOneObjectKeys(raw []byte, scope string, canonical []string) error {
|
||||
decoder := json.NewDecoder(bytes.NewReader(raw))
|
||||
if token, err := decoder.Token(); err != nil || token != json.Delim('{') {
|
||||
return fmt.Errorf("%s must be an object", scope)
|
||||
}
|
||||
seen := make(map[string]struct{})
|
||||
for decoder.More() {
|
||||
token, err := decoder.Token()
|
||||
if err != nil {
|
||||
return errors.New("invalid JSON request")
|
||||
}
|
||||
key, _ := token.(string)
|
||||
if _, duplicate := seen[key]; duplicate {
|
||||
return fmt.Errorf("%s contains duplicate field %q", scope, key)
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
for _, field := range canonical {
|
||||
if key != field && strings.EqualFold(key, field) {
|
||||
return fmt.Errorf("%s field %q must be written as %q", scope, key, field)
|
||||
}
|
||||
}
|
||||
var value json.RawMessage
|
||||
if err := decoder.Decode(&value); err != nil {
|
||||
return errors.New("invalid JSON request")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package typesafe
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestValidateSystemOneRequestValidQuestionTypes(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
}{
|
||||
{"noul string state", `{"model":"jev-latest","state":"sample","questions":{"safety":{"type":"noul","instructions":"Evaluate safety","criteria":{"safe":"No harm"}}}}`},
|
||||
{"choice object state", `{"model":"jev-latest","state":{"text":"sample"},"questions":{"label":{"type":"choice","instructions":{"task":"Classify"},"criteria":{"safe":"Allowed","unsafe":null}}},"stream":false}`},
|
||||
{"score array state", `{"model":"jev-latest","state":["sample"],"questions":{"quality":{"type":"score","instructions":["Rate quality"],"criteria":["poor","good"]}}}`},
|
||||
{"noul omitted instructions", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul"}}}`},
|
||||
{"noul nullable instructions and criteria", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul","instructions":null,"criteria":null}}}`},
|
||||
{"noul structured descriptions", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul","criteria":{"true":{"examples":["yes",null]},"false":["no",null],"extension":42}}}}`},
|
||||
{"choice structured descriptions", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"choice","instructions":null,"criteria":{"a":{"description":"A","extra":null},"b":["B",null],"c":null}}}}`},
|
||||
{"choice empty criteria", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"choice","criteria":{}}}}`},
|
||||
{"score one level", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"score","criteria":["only"]}}}`},
|
||||
{"score object level", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"score","instructions":null,"criteria":[{"description":"only","extra":null}]}}}`},
|
||||
{"score array level", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"score","criteria":[["only",null]]}}}`},
|
||||
{"native extensions preserved", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul","extension":{"kept":true}}},"extension":[1,null]}`},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
model, err := ValidateSystemOneRequest([]byte(tc.body))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, JevLatestModel, model)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateSystemOneRequestRejectsInvalidRequests(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
want string
|
||||
}{
|
||||
{"invalid json", `{`, "invalid JSON"},
|
||||
{"missing model", `{"state":"x","questions":{"q":{"type":"noul","instructions":"x"}}}`, "model"},
|
||||
{"illegal model", `{"model":"jev-old","state":"x","questions":{"q":{"type":"noul","instructions":"x"}}}`, "jev-latest"},
|
||||
{"model with whitespace", `{"model":" jev-latest ","state":"x","questions":{"q":{"type":"noul","instructions":"x"}}}`, "jev-latest"},
|
||||
{"missing state", `{"model":"jev-latest","questions":{"q":{"type":"noul","instructions":"x"}}}`, "state"},
|
||||
{"scalar state", `{"model":"jev-latest","state":42,"questions":{"q":{"type":"noul","instructions":"x"}}}`, "state"},
|
||||
{"empty questions", `{"model":"jev-latest","state":"x","questions":{}}`, "questions"},
|
||||
{"array questions", `{"model":"jev-latest","state":"x","questions":[{"type":"noul"}]}`, "questions must be a non-empty object"},
|
||||
{"null questions", `{"model":"jev-latest","state":"x","questions":null}`, "questions must be a non-empty object"},
|
||||
{"unknown question type", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"boolean","instructions":"x"}}}`, "unsupported type"},
|
||||
{"question type with whitespace", `{"model":"jev-latest","state":"x","questions":{"q":{"type":" noul ","instructions":"x"}}}`, "unsupported type"},
|
||||
{"noul criteria array", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","instructions":"x","criteria":[]}}}`, "noul criteria"},
|
||||
{"noul numeric description", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","criteria":{"true":1}}}}`, "noul criteria"},
|
||||
{"noul boolean description", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","criteria":{"false":false}}}}`, "noul criteria"},
|
||||
{"missing choice criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice"}}}`, "choice criteria"},
|
||||
{"null choice criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice","criteria":null}}}`, "choice criteria"},
|
||||
{"choice criteria array", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice","criteria":[]}}}`, "choice criteria"},
|
||||
{"choice numeric value", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice","instructions":"x","criteria":{"one":1}}}}`, "choice criteria"},
|
||||
{"missing score criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score"}}}`, "score criteria"},
|
||||
{"empty score criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":[]}}}`, "score criteria"},
|
||||
{"null score criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":null}}}`, "score criteria"},
|
||||
{"score map is not wire protocol", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":{"0":"low"}}}}`, "score criteria"},
|
||||
{"numeric score level", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":["low",2]}}}`, "score criteria"},
|
||||
{"null score level", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":[null]}}}`, "score criteria"},
|
||||
{"numeric instructions", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","instructions":1}}}`, "instructions"},
|
||||
{"boolean instructions", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice","instructions":true,"criteria":{}}}}`, "instructions"},
|
||||
{"stream true", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","instructions":"x"}},"stream":true}`, "streaming"},
|
||||
{"case variant model smuggles upstream model", `{"model":"jev-pro","MODEL":"jev-latest","state":"x","questions":{"q":{"type":"noul"}}}`, `must be written as "model"`},
|
||||
{"case variant stream hides streaming", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul"}},"stream":true,"Stream":false}`, `must be written as "stream"`},
|
||||
{"unicode fold variant state", `{"model":"jev-latest","state":"x","ſtate":"y","questions":{"q":{"type":"noul"}}}`, `must be written as "state"`},
|
||||
{"duplicate model", `{"model":"jev-pro","model":"jev-latest","state":"x","questions":{"q":{"type":"noul"}}}`, `duplicate field "model"`},
|
||||
{"escaped duplicate model", `{"\u006dodel":"jev-pro","model":"jev-latest","state":"x","questions":{"q":{"type":"noul"}}}`, `duplicate field "model"`},
|
||||
{"duplicate state", `{"model":"jev-latest","state":"benign","state":"payload","questions":{"q":{"type":"noul"}}}`, `duplicate field "state"`},
|
||||
{"duplicate question id", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul"},"q":{"type":"noul"}}}`, `duplicate field "q"`},
|
||||
{"duplicate question instructions", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","instructions":"a","instructions":"b"}}}`, `duplicate field "instructions"`},
|
||||
{"case variant question type", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","Type":"choice"}}}`, `must be written as "type"`},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, err := ValidateSystemOneRequest([]byte(tc.body))
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), tc.want)
|
||||
if tc.name == "stream true" {
|
||||
require.True(t, errors.Is(err, ErrStreamingUnsupported))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeSystemOneResponseToleratesUsageShapes(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
model string
|
||||
input, output int
|
||||
}{
|
||||
{"integers", `{"model":"jev-1","usage":{"input_tokens":12,"output_tokens":3}}`, "jev-1", 12, 3},
|
||||
{"floats", `{"model":"jev-1","usage":{"input_tokens":12.0,"output_tokens":2.6}}`, "jev-1", 12, 3},
|
||||
{"numeric strings", `{"usage":{"input_tokens":"15","output_tokens":" 4 "}}`, "", 15, 4},
|
||||
{"non numeric usage", `{"model":7,"usage":{"input_tokens":"many","output_tokens":null}}`, "", 0, 0},
|
||||
{"negative usage", `{"usage":{"input_tokens":-5}}`, "", 0, 0},
|
||||
{"usage not object", `{"model":"jev-1","usage":"none","answers":{}}`, "jev-1", 0, 0},
|
||||
{"missing usage", `{"answers":{}}`, "", 0, 0},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
decoded, err := DecodeSystemOneResponse(strings.NewReader(tc.body))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []byte(tc.body), decoded.Body)
|
||||
require.Equal(t, tc.model, decoded.Model)
|
||||
require.Equal(t, tc.input, decoded.Usage.InputTokens)
|
||||
require.Equal(t, tc.output, decoded.Usage.OutputTokens)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeSystemOneResponseRejectsNonObjects(t *testing.T) {
|
||||
for _, body := range []string{`[]`, `"text"`, `null`, `{`, `data: {}`} {
|
||||
_, err := DecodeSystemOneResponse(strings.NewReader(body))
|
||||
require.Error(t, err, body)
|
||||
}
|
||||
_, err := DecodeSystemOneResponse(strings.NewReader(`{"pad":"` + strings.Repeat("a", MaxSystemOneResponseBytes) + `"}`))
|
||||
require.ErrorIs(t, err, ErrSystemOneResponseTooLarge)
|
||||
}
|
||||
@@ -19,10 +19,7 @@ const (
|
||||
// repository and service layers build their own client identity from it, so
|
||||
// one bump here covers OAuth traffic and billing probes together.
|
||||
// Keep in sync with https://x.ai/cli/stable.
|
||||
CLIClientVersion = "0.2.120"
|
||||
// billingCLIUserAgent is the legacy pager/shell UA used by billing probes.
|
||||
// Distinct from CLIUserAgent() in cli_identity.go (workspace-style UA).
|
||||
billingCLIUserAgent = "grok-pager/" + CLIClientVersion + " grok-shell/" + CLIClientVersion + " (macos; aarch64)"
|
||||
CLIClientVersion = "1.0.46"
|
||||
|
||||
BillingWeeklyPath = "/billing?format=credits"
|
||||
BillingMonthlyPath = "/billing"
|
||||
@@ -148,7 +145,8 @@ func ApplyCLIBillingHeaders(req *http.Request, accessToken string) {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set(CLITokenAuthHeader, CLITokenAuthValue)
|
||||
req.Header.Set(CLIClientVersionHeader, CLIClientVersion)
|
||||
req.Header.Set("User-Agent", billingCLIUserAgent)
|
||||
req.Header.Set("User-Agent", CLIUserAgent(CLIClientVersion))
|
||||
req.Header.Set("x-grok-client-mode", CLIClientMode)
|
||||
}
|
||||
|
||||
// ParseBillingPayload unmarshals a billing API response body.
|
||||
|
||||
@@ -38,7 +38,8 @@ func TestApplyCLIBillingHeaders(t *testing.T) {
|
||||
require.Equal(t, "Bearer token", req.Header.Get("Authorization"))
|
||||
require.Equal(t, CLITokenAuthValue, req.Header.Get(CLITokenAuthHeader))
|
||||
require.Equal(t, CLIClientVersion, req.Header.Get(CLIClientVersionHeader))
|
||||
require.Equal(t, "grok-pager/"+CLIClientVersion+" grok-shell/"+CLIClientVersion+" (macos; aarch64)", req.UserAgent())
|
||||
require.Equal(t, CLIUserAgent(CLIClientVersion), req.UserAgent())
|
||||
require.Equal(t, "interactive", req.Header.Get("x-grok-client-mode"))
|
||||
}
|
||||
|
||||
func TestBuildBillingSummaryWeeklyAndMonthly(t *testing.T) {
|
||||
|
||||
@@ -3,6 +3,7 @@ package xai
|
||||
import (
|
||||
"net/http"
|
||||
"os"
|
||||
"runtime"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/mod/semver"
|
||||
@@ -16,7 +17,7 @@ const (
|
||||
CLIProxyHost = "cli-chat-proxy.grok.com"
|
||||
|
||||
// CLIStableVersion is the known-good minimum client version accepted by cli-chat-proxy.
|
||||
CLIStableVersion = "0.2.93"
|
||||
CLIStableVersion = "1.0.13"
|
||||
|
||||
// CLIVersionEnv is the optional operator override for CLIStableVersion.
|
||||
CLIVersionEnv = "XAI_GROK_CLI_VERSION"
|
||||
@@ -24,11 +25,11 @@ const (
|
||||
// CLITokenAuth is required by cli-chat-proxy for Grok Build OAuth tokens.
|
||||
CLITokenAuth = "xai-grok-cli"
|
||||
|
||||
// CLIClientIdentifier is the x-grok-client-identifier value used by Grok shell/CLI.
|
||||
CLIClientIdentifier = "grok-shell"
|
||||
// CLIClientIdentifier 对齐官方交互式 CLI 主请求的客户端标识。
|
||||
CLIClientIdentifier = "grok-pager"
|
||||
|
||||
// CLIClientMode is used by billing / quota probes on the CLI surface.
|
||||
CLIClientMode = "cli"
|
||||
// CLIClientMode 对齐官方 CLI 正常交互模式。
|
||||
CLIClientMode = "interactive"
|
||||
)
|
||||
|
||||
// ResolveCLIVersion returns a supported CLI client version.
|
||||
@@ -54,12 +55,24 @@ func IsSupportedCLIVersion(version string) bool {
|
||||
semver.Compare(canonical, minimum) >= 0
|
||||
}
|
||||
|
||||
// CLIUserAgent builds the workspace-style User-Agent for a CLI client version.
|
||||
// CLIUserAgent 对齐官方交互式 CLI 的 UA,平台名称使用 Rust 的格式。
|
||||
func CLIUserAgent(version string) string {
|
||||
if strings.TrimSpace(version) == "" {
|
||||
version = CLIClientVersion
|
||||
}
|
||||
return "xai-grok-workspace/" + version
|
||||
platform, arch := runtime.GOOS, runtime.GOARCH
|
||||
if platform == "darwin" {
|
||||
platform = "macos"
|
||||
}
|
||||
switch arch {
|
||||
case "amd64":
|
||||
arch = "x86_64"
|
||||
case "arm64":
|
||||
arch = "aarch64"
|
||||
case "386":
|
||||
arch = "x86"
|
||||
}
|
||||
return "grok-pager/" + version + " grok-shell/" + version + " (" + platform + "; " + arch + ")"
|
||||
}
|
||||
|
||||
// ApplyCLIProxyHeaders stamps the fixed Grok CLI identity when the request
|
||||
@@ -75,5 +88,7 @@ func ApplyCLIProxyHeaders(req *http.Request) {
|
||||
req.Header.Set("X-XAI-Token-Auth", CLITokenAuth)
|
||||
req.Header.Set("x-grok-client-version", version)
|
||||
req.Header.Set("x-grok-client-identifier", CLIClientIdentifier)
|
||||
req.Header.Set("x-grok-client-mode", CLIClientMode)
|
||||
req.Header.Set("x-authenticateresponse", "authenticate-response")
|
||||
req.Header.Set("User-Agent", CLIUserAgent(version))
|
||||
}
|
||||
|
||||
@@ -2,11 +2,40 @@ package xai
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"runtime"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/mod/semver"
|
||||
)
|
||||
|
||||
func TestApplyCLIProxyHeadersMeetsUpstreamMinimumVersion(t *testing.T) {
|
||||
// 上游 426 明确要求至少 1.0.13;断言独立于生产版本常量,防止旧覆盖值绕过下限。
|
||||
for _, override := range []string{"", "0.2.93", "0.2.120", "1.0.12", "1.0.13-beta.1", "1.0.13", "1.0.14-alpha.1"} {
|
||||
t.Run("override="+override, func(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, override)
|
||||
req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
ApplyCLIProxyHeaders(req)
|
||||
|
||||
version := req.Header.Get("x-grok-client-version")
|
||||
require.True(t, semver.IsValid("v"+version))
|
||||
require.GreaterOrEqual(t, semver.Compare("v"+version, "v1.0.13"), 0,
|
||||
"Grok Responses 会以 426 拒绝版本 %s,最低要求为 1.0.13", version)
|
||||
require.Contains(t, req.Header.Get("User-Agent"), "grok-shell/"+version+" (")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCLIUserAgentMatchesOfficialInteractiveCapture(t *testing.T) {
|
||||
// 来源:官方 CLI 1.0.46 交互式界面在 Linux x86_64 的主 Responses 请求抓包。
|
||||
if runtime.GOOS != "linux" || runtime.GOARCH != "amd64" {
|
||||
t.Skip("该抓包样本来自 Linux x86_64")
|
||||
}
|
||||
require.Equal(t, "grok-pager/1.0.46 grok-shell/1.0.46 (linux; x86_64)", CLIUserAgent("1.0.46"))
|
||||
}
|
||||
|
||||
func TestResolveCLIVersionDefaultsToPinnedClientVersion(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, "")
|
||||
// Default advertise pin is CLIClientVersion; CLIStableVersion is only the floor.
|
||||
@@ -16,18 +45,18 @@ func TestResolveCLIVersionDefaultsToPinnedClientVersion(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestResolveCLIVersionAcceptsValidOverride(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, "0.2.95-alpha.1")
|
||||
require.Equal(t, "0.2.95-alpha.1", ResolveCLIVersion())
|
||||
t.Setenv(CLIVersionEnv, "1.0.14-alpha.1")
|
||||
require.Equal(t, "1.0.14-alpha.1", ResolveCLIVersion())
|
||||
}
|
||||
|
||||
func TestResolveCLIVersionRejectsUnsafeOrTooOld(t *testing.T) {
|
||||
for _, version := range []string{
|
||||
"0.2.92",
|
||||
"0.2.93-beta.1",
|
||||
"0.2.95\r\nX-Injected: true",
|
||||
"0.2.093",
|
||||
"0.3",
|
||||
"1",
|
||||
"1.0.12",
|
||||
"1.0.13-beta.1",
|
||||
"1.0.14\r\nX-Injected: true",
|
||||
"1.0.014",
|
||||
"1.1",
|
||||
"2",
|
||||
} {
|
||||
t.Run(version, func(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, version)
|
||||
@@ -48,11 +77,13 @@ func TestApplyCLIProxyHeaders(t *testing.T) {
|
||||
require.Equal(t, CLIClientVersion, req.Header.Get("x-grok-client-version"))
|
||||
require.Equal(t, CLIClientIdentifier, req.Header.Get("x-grok-client-identifier"))
|
||||
require.Equal(t, CLITokenAuth, req.Header.Get("X-XAI-Token-Auth"))
|
||||
require.Equal(t, "interactive", req.Header.Get("x-grok-client-mode"))
|
||||
require.Equal(t, "authenticate-response", req.Header.Get("x-authenticateresponse"))
|
||||
require.Equal(t, CLIUserAgent(CLIClientVersion), req.Header.Get("User-Agent"))
|
||||
}
|
||||
|
||||
func TestApplyCLIProxyHeadersLeavesAPIHostUnchanged(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, "0.2.95")
|
||||
t.Setenv(CLIVersionEnv, "1.0.14")
|
||||
|
||||
req, err := http.NewRequest(http.MethodPost, "https://api.x.ai/v1/responses", nil)
|
||||
require.NoError(t, err)
|
||||
@@ -63,5 +94,7 @@ func TestApplyCLIProxyHeadersLeavesAPIHostUnchanged(t *testing.T) {
|
||||
require.Empty(t, req.Header.Get("x-grok-client-version"))
|
||||
require.Empty(t, req.Header.Get("x-grok-client-identifier"))
|
||||
require.Empty(t, req.Header.Get("X-XAI-Token-Auth"))
|
||||
require.Empty(t, req.Header.Get("x-grok-client-mode"))
|
||||
require.Empty(t, req.Header.Get("x-authenticateresponse"))
|
||||
require.Equal(t, "direct-api-client/1.0", req.Header.Get("User-Agent"))
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
const (
|
||||
apiKeyRateLimitKeyPrefix = "apikey:ratelimit:"
|
||||
apiKeyRateLimitDuration = 24 * time.Hour
|
||||
apiKeyCreateCountKeyPrefix = "apikey:create_count:"
|
||||
apiKeyAuthCachePrefix = "apikey:auth:"
|
||||
authCacheInvalidateChannel = "auth:cache:invalidate"
|
||||
)
|
||||
@@ -24,6 +25,11 @@ func apiKeyRateLimitKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", apiKeyRateLimitKeyPrefix, userID)
|
||||
}
|
||||
|
||||
// apiKeyCreateCountKey generates the Redis key for per-user API key creation counting.
|
||||
func apiKeyCreateCountKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", apiKeyCreateCountKeyPrefix, userID)
|
||||
}
|
||||
|
||||
func apiKeyAuthCacheKey(key string) string {
|
||||
return fmt.Sprintf("%s%s", apiKeyAuthCachePrefix, key)
|
||||
}
|
||||
@@ -54,9 +60,18 @@ func (c *apiKeyCache) IncrementCreateAttemptCount(ctx context.Context, userID in
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *apiKeyCache) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
key := apiKeyRateLimitKey(userID)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
// IncrementCreateCount 在固定窗口内累加创建次数并返回累加后的值。
|
||||
// ExpireNX 只在首次创建计数键时设置过期,后续创建不会延长窗口;
|
||||
// MULTI 保证 INCR 与 EXPIRE 同时生效,避免计数键丢失 TTL 后永久封禁。
|
||||
func (c *apiKeyCache) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
key := apiKeyCreateCountKey(userID)
|
||||
pipe := c.rdb.TxPipeline()
|
||||
incr := pipe.Incr(ctx, key)
|
||||
pipe.ExpireNX(ctx, key, window)
|
||||
if _, err := pipe.Exec(ctx); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return incr.Val(), nil
|
||||
}
|
||||
|
||||
func (c *apiKeyCache) IncrementDailyUsage(ctx context.Context, apiKey string) error {
|
||||
|
||||
@@ -51,19 +51,6 @@ func (s *ApiKeyCacheSuite) TestCreateAttemptCount() {
|
||||
s.AssertTTLWithin(ttl, 1*time.Second, apiKeyRateLimitDuration)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "delete_removes_key",
|
||||
fn: func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache) {
|
||||
userID := int64(1)
|
||||
|
||||
require.NoError(s.T(), cache.IncrementCreateAttemptCount(ctx, userID))
|
||||
require.NoError(s.T(), cache.DeleteCreateAttemptCount(ctx, userID), "DeleteCreateAttemptCount")
|
||||
|
||||
count, err := cache.GetCreateAttemptCount(ctx, userID)
|
||||
require.NoError(s.T(), err, "expected nil error after delete")
|
||||
require.Equal(s.T(), 0, count, "expected zero count after delete")
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAPIKeyCacheIncrementCreateCountUsesFixedWindow(t *testing.T) {
|
||||
server := miniredis.RunT(t)
|
||||
client := redis.NewClient(&redis.Options{Addr: server.Addr()})
|
||||
defer func() { _ = client.Close() }()
|
||||
cache := NewAPIKeyCache(client)
|
||||
ctx := context.Background()
|
||||
key := apiKeyCreateCountKey(7)
|
||||
|
||||
count, err := cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
require.Equal(t, time.Hour, server.TTL(key))
|
||||
|
||||
// 后续创建不得延长窗口,否则持续创建会让计数永不过期。
|
||||
server.FastForward(40 * time.Minute)
|
||||
count, err = cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(2), count)
|
||||
require.Equal(t, 20*time.Minute, server.TTL(key))
|
||||
|
||||
// 其他用户独立计数。
|
||||
count, err = cache.IncrementCreateCount(ctx, 8, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
|
||||
// 窗口到期后重新计数。
|
||||
server.FastForward(21 * time.Minute)
|
||||
count, err = cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
}
|
||||
@@ -652,6 +652,17 @@ func apiKeyListOrder(params pagination.PaginationParams) []func(*entsql.Selector
|
||||
sortBy := strings.ToLower(strings.TrimSpace(params.SortBy))
|
||||
sortOrder := params.NormalizedSortOrder(pagination.SortOrderDesc)
|
||||
|
||||
if sortBy == "group" {
|
||||
// Sort before pagination, keeping ungrouped keys last in either direction.
|
||||
opts := []entsql.OrderTermOption{entsql.OrderNullsLast()}
|
||||
tieOrder := dbent.Asc(apikey.FieldID)
|
||||
if sortOrder == pagination.SortOrderDesc {
|
||||
opts = append(opts, entsql.OrderDesc())
|
||||
tieOrder = dbent.Desc(apikey.FieldID)
|
||||
}
|
||||
return []func(*entsql.Selector){apikey.ByGroupField(group.FieldName, opts...), tieOrder}
|
||||
}
|
||||
|
||||
var field string
|
||||
switch sortBy {
|
||||
case "name":
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAPIKeyRepositoryListByUserIDSortByGroup(t *testing.T) {
|
||||
repo, client := newAPIKeyRepoSQLite(t)
|
||||
ctx := context.Background()
|
||||
user := mustCreateAPIKeyRepoUser(t, ctx, client, "group-sort@test.com")
|
||||
otherUser := mustCreateAPIKeyRepoUser(t, ctx, client, "other-group-sort@test.com")
|
||||
createGroup := func(name string) *dbent.Group {
|
||||
g, err := client.Group.Create().SetName(name).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
return g
|
||||
}
|
||||
// Create groups in reverse name order so sorting by group ID cannot pass.
|
||||
zulu := createGroup("Zulu")
|
||||
alpha := createGroup("Alpha")
|
||||
createKey := func(userID int64, name string, groupID *int64, status string) int64 {
|
||||
key := &service.APIKey{
|
||||
UserID: userID, Key: "sk-" + name, Name: name, GroupID: groupID, Status: status,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, key))
|
||||
return key.ID
|
||||
}
|
||||
zuluFirst := createKey(user.ID, "match-zulu-first", &zulu.ID, service.StatusActive)
|
||||
ungroupedFirst := createKey(user.ID, "match-ungrouped-first", nil, service.StatusActive)
|
||||
alphaFirst := createKey(user.ID, "match-alpha-first", &alpha.ID, service.StatusActive)
|
||||
zuluSecond := createKey(user.ID, "match-zulu-second", &zulu.ID, service.StatusActive)
|
||||
alphaSecond := createKey(user.ID, "match-alpha-second", &alpha.ID, service.StatusActive)
|
||||
ungroupedSecond := createKey(user.ID, "match-ungrouped-second", nil, service.StatusActive)
|
||||
createKey(otherUser.ID, "match-other-user", &alpha.ID, service.StatusActive)
|
||||
createKey(user.ID, "excluded-search", &alpha.ID, service.StatusActive)
|
||||
createKey(user.ID, "match-inactive", &alpha.ID, service.StatusDisabled)
|
||||
deleted := createKey(user.ID, "match-deleted", &alpha.ID, service.StatusActive)
|
||||
require.NoError(t, repo.Delete(ctx, deleted))
|
||||
|
||||
ungroupedID := int64(0)
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
order string
|
||||
groupID *int64
|
||||
want []int64
|
||||
}{
|
||||
{"ascending", "asc", nil, []int64{alphaFirst, alphaSecond, zuluFirst, zuluSecond, ungroupedFirst, ungroupedSecond}},
|
||||
{"descending", "desc", nil, []int64{zuluSecond, zuluFirst, alphaSecond, alphaFirst, ungroupedSecond, ungroupedFirst}},
|
||||
{"group filter", "asc", &alpha.ID, []int64{alphaFirst, alphaSecond}},
|
||||
{"ungrouped filter", "desc", &ungroupedID, []int64{ungroupedSecond, ungroupedFirst}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var got []int64
|
||||
const pageSize = 3
|
||||
pages := (len(tc.want) + pageSize - 1) / pageSize
|
||||
for page := 1; page <= pages; page++ {
|
||||
keys, result, err := repo.ListByUserID(ctx, user.ID, pagination.PaginationParams{
|
||||
Page: page, PageSize: pageSize, SortBy: "group", SortOrder: tc.order,
|
||||
}, service.APIKeyListFilters{Search: "match", Status: service.StatusActive, GroupID: tc.groupID})
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, len(tc.want), result.Total)
|
||||
require.Equal(t, pages, result.Pages)
|
||||
for _, key := range keys {
|
||||
got = append(got, key.ID)
|
||||
if key.GroupID != nil {
|
||||
require.NotNil(t, key.Group, "sorting must preserve group preloading")
|
||||
}
|
||||
}
|
||||
}
|
||||
require.Equal(t, tc.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// 余额在途预留:
|
||||
// - billing:inflight:{uid} ZSET member=requestID score=过期时间(ms)
|
||||
// - billing:inflight_amt:{uid} HASH field=requestID value=预留金额(USD)
|
||||
//
|
||||
// 两个 key 使用同一 hash tag,保证 Redis Cluster 下 Lua 可同时访问。
|
||||
const (
|
||||
billingInflightKeyPrefix = "billing:inflight:"
|
||||
billingInflightAmtKeyPrefix = "billing:inflight_amt:"
|
||||
)
|
||||
|
||||
func billingInflightKeys(userID int64) (string, string) {
|
||||
return fmt.Sprintf("%s{%d}", billingInflightKeyPrefix, userID),
|
||||
fmt.Sprintf("%s{%d}", billingInflightAmtKeyPrefix, userID)
|
||||
}
|
||||
|
||||
var (
|
||||
// KEYS: [1]=zset [2]=hash
|
||||
// ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=balance [6]=key_ttl_ms
|
||||
// 返回 {allowed, inflight_sum(string), inflight_count}
|
||||
// 规则:先清理已过期成员;若无在途预留则直接放行(与旧行为一致,外层已校验余额 > 阈值);
|
||||
// 否则要求 balance - sum(在途) >= amount。放行时登记预留。
|
||||
reserveInflightBalanceScript = 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
|
||||
local count = redis.call('ZCARD', KEYS[1])
|
||||
local amount = tonumber(ARGV[4])
|
||||
local balance = tonumber(ARGV[5])
|
||||
if count > 0 and (balance - sum) < amount then
|
||||
return {0, tostring(sum), count}
|
||||
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[6])
|
||||
redis.call('PEXPIRE', KEYS[2], ARGV[6])
|
||||
return {1, tostring(sum), count}
|
||||
`)
|
||||
|
||||
// KEYS: [1]=zset [2]=hash ARGV: [1]=member [2]=expire_at_ms [3]=key_ttl_ms [4]=now_ms
|
||||
// 仅当成员仍存在且尚未过期(score > now)时续期(XX),并把两个 key 的 TTL 至少延长到 key_ttl_ms。
|
||||
// 已过期但尚未被惰性清理的成员不得被复活:顺手清掉并返回 0。
|
||||
renewInflightBalanceScript = redis.NewScript(`
|
||||
local score = redis.call('ZSCORE', KEYS[1], ARGV[1])
|
||||
if not score then
|
||||
return 0
|
||||
end
|
||||
if tonumber(score) <= tonumber(ARGV[4]) then
|
||||
redis.call('ZREM', KEYS[1], ARGV[1])
|
||||
redis.call('HDEL', KEYS[2], ARGV[1])
|
||||
return 0
|
||||
end
|
||||
redis.call('ZADD', KEYS[1], 'XX', tonumber(ARGV[2]), ARGV[1])
|
||||
local ttl = tonumber(ARGV[3])
|
||||
if redis.call('PTTL', KEYS[1]) < ttl then
|
||||
redis.call('PEXPIRE', KEYS[1], ttl)
|
||||
end
|
||||
if redis.call('PTTL', KEYS[2]) < ttl then
|
||||
redis.call('PEXPIRE', KEYS[2], ttl)
|
||||
end
|
||||
return 1
|
||||
`)
|
||||
|
||||
releaseInflightBalanceScript = redis.NewScript(`
|
||||
redis.call('ZREM', KEYS[1], ARGV[1])
|
||||
redis.call('HDEL', KEYS[2], ARGV[1])
|
||||
return 1
|
||||
`)
|
||||
)
|
||||
|
||||
// ReserveInflightBalance 实现 service.InflightBalanceReservationCache。
|
||||
func (c *billingCache) ReserveInflightBalance(ctx context.Context, userID int64, requestID string, amount, balance float64, ttl time.Duration) (bool, float64, error) {
|
||||
zkey, hkey := billingInflightKeys(userID)
|
||||
now := time.Now().UnixMilli()
|
||||
ttlMs := ttl.Milliseconds()
|
||||
if ttlMs <= 0 {
|
||||
ttlMs = 1
|
||||
}
|
||||
res, err := reserveInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey},
|
||||
now,
|
||||
now+ttlMs,
|
||||
requestID,
|
||||
strconv.FormatFloat(amount, 'f', -1, 64),
|
||||
strconv.FormatFloat(balance, 'f', -1, 64),
|
||||
ttlMs,
|
||||
).Slice()
|
||||
if err != nil {
|
||||
return false, 0, err
|
||||
}
|
||||
if len(res) < 2 {
|
||||
return false, 0, fmt.Errorf("unexpected inflight reservation reply: %v", res)
|
||||
}
|
||||
allowed, _ := res[0].(int64)
|
||||
var sum float64
|
||||
if s, ok := res[1].(string); ok {
|
||||
sum, _ = strconv.ParseFloat(s, 64)
|
||||
}
|
||||
return allowed == 1, sum, nil
|
||||
}
|
||||
|
||||
// ReleaseInflightBalance 实现 service.InflightBalanceReservationCache。
|
||||
func (c *billingCache) ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error {
|
||||
zkey, hkey := billingInflightKeys(userID)
|
||||
return releaseInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID).Err()
|
||||
}
|
||||
|
||||
// RenewInflightBalance 实现 service.InflightBalanceReservationRenewer。
|
||||
func (c *billingCache) RenewInflightBalance(ctx context.Context, userID int64, requestID string, ttl time.Duration) (bool, error) {
|
||||
zkey, hkey := billingInflightKeys(userID)
|
||||
ttlMs := ttl.Milliseconds()
|
||||
if ttlMs <= 0 {
|
||||
ttlMs = 1
|
||||
}
|
||||
now := time.Now().UnixMilli()
|
||||
n, err := renewInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID, now+ttlMs, ttlMs, now).Int64()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return n == 1, nil
|
||||
}
|
||||
|
||||
var (
|
||||
_ service.InflightBalanceReservationCache = (*billingCache)(nil)
|
||||
_ service.InflightBalanceReservationRenewer = (*billingCache)(nil)
|
||||
)
|
||||
@@ -0,0 +1,287 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newInflightTestEnv(t *testing.T, enabled bool, ttlSeconds int) (*miniredis.Miniredis, *billingCache, *service.BillingCacheService) {
|
||||
t.Helper()
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
cache := &billingCache{rdb: rdb}
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{
|
||||
Enabled: enabled,
|
||||
TTLSeconds: ttlSeconds,
|
||||
DefaultMaxTokens: 8192,
|
||||
}
|
||||
svc := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(svc.Stop)
|
||||
return mr, cache, svc
|
||||
}
|
||||
|
||||
func inflightCount(t *testing.T, cache *billingCache, userID int64) int64 {
|
||||
t.Helper()
|
||||
zkey, _ := billingInflightKeys(userID)
|
||||
n, err := cache.rdb.ZCard(context.Background(), zkey).Result()
|
||||
require.NoError(t, err)
|
||||
return n
|
||||
}
|
||||
|
||||
func TestInflightReservation_ConcurrentAdmitsOnlyWhatBalanceCovers(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 42}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0))
|
||||
|
||||
const n = 20
|
||||
var admitted, rejected atomic.Int32
|
||||
releases := make(chan func(), n)
|
||||
var wg sync.WaitGroup
|
||||
start := make(chan struct{})
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.3)
|
||||
if err != nil {
|
||||
if !errors.Is(err, service.ErrInsufficientBalance) {
|
||||
t.Errorf("unexpected error: %v", err)
|
||||
}
|
||||
rejected.Add(1)
|
||||
return
|
||||
}
|
||||
admitted.Add(1)
|
||||
releases <- release
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(releases)
|
||||
|
||||
// balance 1.0, estimate 0.3: first always admitted, then while 1.0 - inflight >= 0.3 → 3 total.
|
||||
require.Equal(t, int32(3), admitted.Load())
|
||||
require.Equal(t, int32(n-3), rejected.Load())
|
||||
require.Equal(t, int64(3), inflightCount(t, cache, user.ID))
|
||||
|
||||
for release := range releases {
|
||||
release()
|
||||
release() // idempotent
|
||||
}
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
|
||||
// After release, capacity is available again.
|
||||
release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.3)
|
||||
require.NoError(t, err)
|
||||
release()
|
||||
}
|
||||
|
||||
func TestInflightReservation_FirstRequestAlwaysAdmitted(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 7}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 0.01))
|
||||
|
||||
release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 5)
|
||||
require.NoError(t, err, "a lone request keeps legacy behavior even if estimate exceeds balance")
|
||||
_, err = svc.ReserveInflightBalance(ctx, user, nil, nil, 0.001)
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance)
|
||||
release()
|
||||
release2, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.001)
|
||||
require.NoError(t, err)
|
||||
release2()
|
||||
}
|
||||
|
||||
func TestInflightReservation_TTLExpiryFreesLeakedReservation(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 1)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 9}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0))
|
||||
|
||||
_, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.8) // leaked: never released
|
||||
require.NoError(t, err)
|
||||
_, err = svc.ReserveInflightBalance(ctx, user, nil, nil, 0.8)
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance)
|
||||
|
||||
time.Sleep(1100 * time.Millisecond)
|
||||
release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.8)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), inflightCount(t, cache, user.ID))
|
||||
release()
|
||||
}
|
||||
|
||||
func TestInflightReservation_DisabledKeepsLegacyBehavior(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, false, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 11}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 0.01))
|
||||
require.False(t, svc.InflightReservationEnabled())
|
||||
for i := 0; i < 5; i++ {
|
||||
_, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 10)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
func TestInflightReservation_SubscriptionUnaffected(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 12}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 0))
|
||||
group := &service.Group{ID: 1, SubscriptionType: service.SubscriptionTypeSubscription}
|
||||
sub := &service.UserSubscription{ID: 1}
|
||||
for i := 0; i < 5; i++ {
|
||||
_, err := svc.ReserveInflightBalance(ctx, user, group, sub, 10)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
func TestInflightReservation_ZeroEstimateSkips(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 13}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1))
|
||||
_, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
// downReserveCache 余额读取正常,但预留走一个已关闭的 Redis。
|
||||
type downReserveCache struct {
|
||||
service.BillingCache
|
||||
down *billingCache
|
||||
}
|
||||
|
||||
func (d *downReserveCache) ReserveInflightBalance(ctx context.Context, userID int64, requestID string, amount, balance float64, ttl time.Duration) (bool, float64, error) {
|
||||
return d.down.ReserveInflightBalance(ctx, userID, requestID, amount, balance, ttl)
|
||||
}
|
||||
|
||||
func (d *downReserveCache) ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error {
|
||||
return d.down.ReleaseInflightBalance(ctx, userID, requestID)
|
||||
}
|
||||
|
||||
func TestInflightReservation_RedisDownFailsOpen(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
up := &billingCache{rdb: rdb}
|
||||
require.NoError(t, up.SetUserBalance(context.Background(), 5, 0.01))
|
||||
|
||||
downMR := miniredis.RunT(t)
|
||||
downRDB := redis.NewClient(&redis.Options{Addr: downMR.Addr(), MaxRetries: -1, DialTimeout: 100 * time.Millisecond})
|
||||
t.Cleanup(func() { _ = downRDB.Close() })
|
||||
downMR.Close()
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
svc := service.NewBillingCacheService(&downReserveCache{BillingCache: up, down: &billingCache{rdb: downRDB}}, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(svc.Stop)
|
||||
|
||||
user := &service.User{ID: 5}
|
||||
for i := 0; i < 3; i++ {
|
||||
release, err := svc.ReserveInflightBalance(context.Background(), user, nil, nil, 100)
|
||||
require.NoError(t, err)
|
||||
release()
|
||||
}
|
||||
}
|
||||
|
||||
func TestInflightReservation_ScriptErrorSurfaces(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr(), MaxRetries: -1})
|
||||
cache := &billingCache{rdb: rdb}
|
||||
mr.Close()
|
||||
_, _, err := cache.ReserveInflightBalance(context.Background(), 1, "r", 1, 1, time.Second)
|
||||
require.Error(t, err)
|
||||
require.False(t, errors.Is(err, context.Canceled))
|
||||
_ = rdb.Close()
|
||||
}
|
||||
|
||||
// 回归:原实现在 handler 返回时即释放预留,而计费(RecordUsage → 余额缓存扣减)是异步的,
|
||||
// 窗口内的顺序请求看到「在途=0 且余额未扣」全部放行(余额 $1、5×$0.9 → -4.40)。
|
||||
// 现在预留由计费任务在余额缓存扣减之后才释放。
|
||||
func TestInflightReservation_SequentialInBillingWindowAdmitsOnlyWhatBalanceCovers(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 77}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0))
|
||||
|
||||
const cost = 0.9
|
||||
var pendingBilling []func()
|
||||
admitted := 0
|
||||
for i := 0; i < 5; i++ {
|
||||
res, err := svc.ReserveInflight(ctx, user, nil, nil, cost)
|
||||
if err != nil {
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance)
|
||||
continue
|
||||
}
|
||||
admitted++
|
||||
taskDone := res.Acquire() // handler submits the async billing task
|
||||
res.HandlerDone() // handler returns before billing lands
|
||||
pendingBilling = append(pendingBilling, func() {
|
||||
require.NoError(t, cache.DeductUserBalance(ctx, user.ID, cost))
|
||||
taskDone()
|
||||
})
|
||||
}
|
||||
require.Equal(t, 1, admitted, "requests arriving before billing lands must not be admitted")
|
||||
|
||||
for _, bill := range pendingBilling {
|
||||
bill()
|
||||
}
|
||||
bal, err := cache.GetUserBalance(ctx, user.ID)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 0.1, bal, 1e-9)
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
func TestInflightReservation_RenewalKeepsStreamingReservationAlive(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 1)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 78}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0))
|
||||
|
||||
res, err := svc.ReserveInflight(ctx, user, nil, nil, 0.8)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(2500 * time.Millisecond) // streaming well past the 1s TTL
|
||||
_, err = svc.ReserveInflight(ctx, user, nil, nil, 0.8)
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance, "renewed reservation must still count")
|
||||
require.Equal(t, int64(1), inflightCount(t, cache, user.ID))
|
||||
|
||||
res.HandlerDone()
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
|
||||
ok, err := cache.RenewInflightBalance(ctx, user.ID, "missing", time.Second)
|
||||
require.NoError(t, err)
|
||||
require.False(t, ok, "renew must not resurrect a released reservation")
|
||||
}
|
||||
|
||||
func TestInflightReservation_RenewDoesNotResurrectExpiredMember(t *testing.T) {
|
||||
_, cache, _ := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
userID := int64(79)
|
||||
zkey, hkey := billingInflightKeys(userID)
|
||||
// 已过期但尚未被惰性清理的成员(score 在过去)。
|
||||
require.NoError(t, cache.rdb.ZAdd(ctx, zkey, redis.Z{Score: float64(time.Now().UnixMilli() - 1000), Member: "stale"}).Err())
|
||||
require.NoError(t, cache.rdb.HSet(ctx, hkey, "stale", "0.5").Err())
|
||||
|
||||
ok, err := cache.RenewInflightBalance(ctx, userID, "stale", time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.False(t, ok, "renew must not resurrect an expired member")
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, userID))
|
||||
exists, err := cache.rdb.HExists(ctx, hkey, "stale").Result()
|
||||
require.NoError(t, err)
|
||||
require.False(t, exists)
|
||||
}
|
||||
@@ -197,6 +197,8 @@ func (r *contentModerationRepository) CountFlaggedByUserSince(ctx context.Contex
|
||||
return 0, nil
|
||||
}
|
||||
// SQL 中的 'cyber_policy' 字面量须与 service.ContentModerationActionCyberPolicy 保持一致。
|
||||
// 'cyber_log_only' matches service.ContentModerationModeCyberLogOnly; these
|
||||
// events remain evidence but never become penalties after allowlist removal.
|
||||
var count int
|
||||
err := r.db.QueryRowContext(ctx, `
|
||||
WITH last_auto_ban AS (
|
||||
@@ -209,6 +211,8 @@ FROM content_moderation_logs
|
||||
WHERE user_id = $1
|
||||
AND flagged = TRUE
|
||||
AND action <> 'hash_block'
|
||||
AND mode <> 'cyber_log_only'
|
||||
AND mode <> 'risk_control_log_only'
|
||||
AND ($3::bool IS FALSE OR action <> 'cyber_policy')
|
||||
AND created_at >= $2
|
||||
AND created_at > COALESCE((SELECT at FROM last_auto_ban), '-infinity'::timestamptz)
|
||||
|
||||
@@ -102,3 +102,18 @@ func TestContentModerationRepositoryCountFlaggedByUserSince_ExcludesCyberPolicyW
|
||||
require.Equal(t, 3, count)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestContentModerationRepositoryCountFlaggedByUserSince_ExcludesCyberLogOnly(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
repo := NewContentModerationRepository(db)
|
||||
since := time.Now().Add(-time.Hour)
|
||||
mock.ExpectQuery(regexp.QuoteMeta("AND mode <> 'cyber_log_only'\n AND mode <> 'risk_control_log_only'")).
|
||||
WithArgs(int64(12), since, false).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
count, err := repo.CountFlaggedByUserSince(context.Background(), 12, since, false)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, count)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
@@ -17,8 +17,42 @@ const (
|
||||
passwordResetKeyPrefix = "password_reset:"
|
||||
passwordResetSentAtKeyPrefix = "password_reset_sent:"
|
||||
notifyCodeUserRateKeyPrefix = "notify_code_user_rate:"
|
||||
|
||||
// attemptsKeySuffix stores the failed-attempt counter next to a verification code.
|
||||
// Kept in a separate key so it can be incremented atomically with INCR.
|
||||
attemptsKeySuffix = ":attempts"
|
||||
)
|
||||
|
||||
// incrAttemptsScript atomically increments the attempt counter for an existing
|
||||
// verification code and aligns the counter TTL with the code TTL.
|
||||
// KEYS[1] = code key, KEYS[2] = attempts key. Returns -1 when the code is missing.
|
||||
var incrAttemptsScript = redis.NewScript(`
|
||||
if redis.call('EXISTS', KEYS[1]) == 0 then
|
||||
return -1
|
||||
end
|
||||
local n = redis.call('INCR', KEYS[2])
|
||||
local ttl = redis.call('PTTL', KEYS[1])
|
||||
if ttl > 0 then
|
||||
redis.call('PEXPIRE', KEYS[2], ttl)
|
||||
end
|
||||
return n
|
||||
`)
|
||||
|
||||
// consumeResetTokenScript atomically compares the stored token hash and deletes it.
|
||||
// KEYS[1] = reset key, ARGV[1] = expected token hash. Returns 1 on success, 0 otherwise.
|
||||
var consumeResetTokenScript = redis.NewScript(`
|
||||
local v = redis.call('GET', KEYS[1])
|
||||
if not v then
|
||||
return 0
|
||||
end
|
||||
local ok, d = pcall(cjson.decode, v)
|
||||
if not ok or type(d) ~= 'table' or d['Token'] ~= ARGV[1] then
|
||||
return 0
|
||||
end
|
||||
redis.call('DEL', KEYS[1])
|
||||
return 1
|
||||
`)
|
||||
|
||||
// verifyCodeKey generates the Redis key for email verification code.
|
||||
// Email is lowercased for case-insensitive consistency.
|
||||
func verifyCodeKey(email string) string {
|
||||
@@ -50,8 +84,7 @@ func NewEmailCache(rdb *redis.Client) service.EmailCache {
|
||||
return &emailCache{rdb: rdb}
|
||||
}
|
||||
|
||||
func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
|
||||
key := verifyCodeKey(email)
|
||||
func (c *emailCache) getCode(ctx context.Context, key string) (*service.VerificationCodeData, error) {
|
||||
val, err := c.rdb.Get(ctx, key).Result()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -60,21 +93,56 @@ func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*se
|
||||
if err := json.Unmarshal([]byte(val), &data); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if n, err := c.rdb.Get(ctx, key+attemptsKeySuffix).Int(); err == nil && n > data.Attempts {
|
||||
data.Attempts = n
|
||||
}
|
||||
return &data, nil
|
||||
}
|
||||
|
||||
func (c *emailCache) SetVerificationCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
|
||||
key := verifyCodeKey(email)
|
||||
func (c *emailCache) setCode(ctx context.Context, key string, data *service.VerificationCodeData, ttl time.Duration) error {
|
||||
val, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.rdb.Set(ctx, key, val, ttl).Err()
|
||||
pipe := c.rdb.TxPipeline()
|
||||
pipe.Set(ctx, key, val, ttl)
|
||||
pipe.Del(ctx, key+attemptsKeySuffix)
|
||||
if data.Attempts > 0 {
|
||||
pipe.Set(ctx, key+attemptsKeySuffix, data.Attempts, ttl)
|
||||
}
|
||||
_, err = pipe.Exec(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *emailCache) incrCodeAttempts(ctx context.Context, key string) (int, error) {
|
||||
n, err := incrAttemptsScript.Run(ctx, c.rdb, []string{key, key + attemptsKeySuffix}).Int()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if n < 0 {
|
||||
return 0, redis.Nil
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (c *emailCache) deleteCode(ctx context.Context, key string) error {
|
||||
return c.rdb.Del(ctx, key, key+attemptsKeySuffix).Err()
|
||||
}
|
||||
|
||||
func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
|
||||
return c.getCode(ctx, verifyCodeKey(email))
|
||||
}
|
||||
|
||||
func (c *emailCache) SetVerificationCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
|
||||
return c.setCode(ctx, verifyCodeKey(email), data, ttl)
|
||||
}
|
||||
|
||||
func (c *emailCache) IncrVerificationCodeAttempts(ctx context.Context, email string) (int, error) {
|
||||
return c.incrCodeAttempts(ctx, verifyCodeKey(email))
|
||||
}
|
||||
|
||||
func (c *emailCache) DeleteVerificationCode(ctx context.Context, email string) error {
|
||||
key := verifyCodeKey(email)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
return c.deleteCode(ctx, verifyCodeKey(email))
|
||||
}
|
||||
|
||||
// Password reset token methods
|
||||
@@ -101,6 +169,16 @@ func (c *emailCache) SetPasswordResetToken(ctx context.Context, email string, da
|
||||
return c.rdb.Set(ctx, key, val, ttl).Err()
|
||||
}
|
||||
|
||||
// ConsumePasswordResetToken atomically deletes the stored reset token when its
|
||||
// stored hash equals tokenHash. Returns true only for the single winning caller.
|
||||
func (c *emailCache) ConsumePasswordResetToken(ctx context.Context, email, tokenHash string) (bool, error) {
|
||||
n, err := consumeResetTokenScript.Run(ctx, c.rdb, []string{passwordResetKey(email)}, tokenHash).Int()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return n == 1, nil
|
||||
}
|
||||
|
||||
func (c *emailCache) DeletePasswordResetToken(ctx context.Context, email string) error {
|
||||
key := passwordResetKey(email)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
@@ -122,30 +200,19 @@ func (c *emailCache) SetPasswordResetEmailCooldown(ctx context.Context, email st
|
||||
// Notify email verification code methods
|
||||
|
||||
func (c *emailCache) GetNotifyVerifyCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
|
||||
key := notifyVerifyKey(email)
|
||||
val, err := c.rdb.Get(ctx, key).Result()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var data service.VerificationCodeData
|
||||
if err := json.Unmarshal([]byte(val), &data); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &data, nil
|
||||
return c.getCode(ctx, notifyVerifyKey(email))
|
||||
}
|
||||
|
||||
func (c *emailCache) SetNotifyVerifyCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
|
||||
key := notifyVerifyKey(email)
|
||||
val, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.rdb.Set(ctx, key, val, ttl).Err()
|
||||
return c.setCode(ctx, notifyVerifyKey(email), data, ttl)
|
||||
}
|
||||
|
||||
func (c *emailCache) IncrNotifyVerifyCodeAttempts(ctx context.Context, email string) (int, error) {
|
||||
return c.incrCodeAttempts(ctx, notifyVerifyKey(email))
|
||||
}
|
||||
|
||||
func (c *emailCache) DeleteNotifyVerifyCode(ctx context.Context, email string) error {
|
||||
key := notifyVerifyKey(email)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
return c.deleteCode(ctx, notifyVerifyKey(email))
|
||||
}
|
||||
|
||||
// User-level rate limiting for notify email verification codes
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newMiniredisEmailCache(t *testing.T) (service.EmailCache, *miniredis.Miniredis, *redis.Client) {
|
||||
t.Helper()
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
return NewEmailCache(rdb), mr, rdb
|
||||
}
|
||||
|
||||
func TestEmailCache_ConcurrentWrongCodesCannotExceedAttemptCap(t *testing.T) {
|
||||
cache, _, _ := newMiniredisEmailCache(t)
|
||||
ctx := context.Background()
|
||||
email := "user@example.com"
|
||||
|
||||
svc := service.NewEmailService(nil, cache)
|
||||
require.NoError(t, cache.SetVerificationCode(ctx, email, &service.VerificationCodeData{
|
||||
Code: "123456",
|
||||
CreatedAt: time.Now(),
|
||||
ExpiresAt: time.Now().Add(15 * time.Minute),
|
||||
}, 15*time.Minute))
|
||||
|
||||
const workers = 50
|
||||
var invalid, maxed atomic.Int32
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
err := svc.VerifyCode(ctx, email, "000000")
|
||||
switch {
|
||||
case errors.Is(err, service.ErrInvalidVerifyCode):
|
||||
invalid.Add(1)
|
||||
case errors.Is(err, service.ErrVerifyCodeMaxAttempts):
|
||||
maxed.Add(1)
|
||||
default:
|
||||
t.Errorf("unexpected result: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// Only attempts 1..4 may return "invalid"; every other guess is rejected by the cap.
|
||||
require.LessOrEqual(t, int(invalid.Load()), 4)
|
||||
require.Equal(t, workers, int(invalid.Load()+maxed.Load()))
|
||||
|
||||
// Even the correct code is now rejected.
|
||||
require.ErrorIs(t, svc.VerifyCode(ctx, email, "123456"), service.ErrVerifyCodeMaxAttempts)
|
||||
|
||||
data, err := cache.GetVerificationCode(ctx, email)
|
||||
require.NoError(t, err)
|
||||
require.GreaterOrEqual(t, data.Attempts, 5)
|
||||
}
|
||||
|
||||
func TestEmailCache_AttemptsResetOnNewCodeAndTTLFollowsCode(t *testing.T) {
|
||||
cache, mr, _ := newMiniredisEmailCache(t)
|
||||
ctx := context.Background()
|
||||
email := "User@Example.com"
|
||||
|
||||
require.NoError(t, cache.SetVerificationCode(ctx, email, &service.VerificationCodeData{Code: "1"}, time.Minute))
|
||||
n, err := cache.IncrVerificationCodeAttempts(ctx, email)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, n)
|
||||
require.Greater(t, mr.TTL(verifyCodeKey(email)+attemptsKeySuffix), time.Duration(0))
|
||||
|
||||
require.NoError(t, cache.SetVerificationCode(ctx, email, &service.VerificationCodeData{Code: "2"}, time.Minute))
|
||||
data, err := cache.GetVerificationCode(ctx, email)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, data.Attempts)
|
||||
|
||||
require.NoError(t, cache.DeleteVerificationCode(ctx, email))
|
||||
_, err = cache.IncrVerificationCodeAttempts(ctx, email)
|
||||
require.Error(t, err)
|
||||
require.False(t, mr.Exists(verifyCodeKey(email)+attemptsKeySuffix))
|
||||
}
|
||||
|
||||
func TestEmailCache_PasswordResetTokenHashedAndSingleUse(t *testing.T) {
|
||||
cache, mr, _ := newMiniredisEmailCache(t)
|
||||
ctx := context.Background()
|
||||
email := "reset@example.com"
|
||||
|
||||
svc := service.NewEmailService(nil, cache)
|
||||
|
||||
// Seed the token the same way SendPasswordResetEmail does (hash only).
|
||||
token, err := svc.GeneratePasswordResetToken()
|
||||
require.NoError(t, err)
|
||||
sum := sha256.Sum256([]byte(token))
|
||||
require.NoError(t, cache.SetPasswordResetToken(ctx, email, &service.PasswordResetTokenData{
|
||||
Token: hex.EncodeToString(sum[:]), CreatedAt: time.Now(),
|
||||
}, 30*time.Minute))
|
||||
|
||||
raw, err := mr.Get(passwordResetKey(email))
|
||||
require.NoError(t, err)
|
||||
require.False(t, strings.Contains(raw, token), "plaintext token must not be stored")
|
||||
|
||||
require.NoError(t, svc.VerifyPasswordResetToken(ctx, email, token))
|
||||
require.ErrorIs(t, svc.ConsumePasswordResetToken(ctx, email, "wrong"), service.ErrInvalidResetToken)
|
||||
|
||||
const workers = 30
|
||||
var ok atomic.Int32
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if svc.ConsumePasswordResetToken(ctx, email, token) == nil {
|
||||
ok.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
require.Equal(t, int32(1), ok.Load())
|
||||
require.False(t, mr.Exists(passwordResetKey(email)))
|
||||
}
|
||||
|
||||
func TestEmailCache_ConsumePasswordResetTokenMismatchKeepsToken(t *testing.T) {
|
||||
cache, mr, _ := newMiniredisEmailCache(t)
|
||||
ctx := context.Background()
|
||||
email := "keep@example.com"
|
||||
require.NoError(t, cache.SetPasswordResetToken(ctx, email, &service.PasswordResetTokenData{Token: "abc"}, time.Minute))
|
||||
|
||||
ok, err := cache.ConsumePasswordResetToken(ctx, email, "xyz")
|
||||
require.NoError(t, err)
|
||||
require.False(t, ok)
|
||||
require.True(t, mr.Exists(passwordResetKey(email)))
|
||||
|
||||
ok, err = cache.ConsumePasswordResetToken(ctx, email, "abc")
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
ok, err = cache.ConsumePasswordResetToken(ctx, email, "abc")
|
||||
require.NoError(t, err)
|
||||
require.False(t, ok)
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user