mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f2669c8cf6 | ||
|
|
85e3110953 | ||
|
|
64caa9af8b | ||
|
|
dab3b87ea2 | ||
|
|
d0392edf06 | ||
|
|
d85cebecc4 | ||
|
|
5f93bb3b3f | ||
|
|
c61f6ebaa8 | ||
|
|
387ee20ed1 | ||
|
|
7ca08d521f | ||
|
|
9264414612 | ||
|
|
874871a25a | ||
|
|
e2939cd8fa | ||
|
|
7ffa38a8ae | ||
|
|
c556aee8fd | ||
|
|
eefa7dd9dc | ||
|
|
3b89cc08dd | ||
|
|
105c955c88 | ||
|
|
8ac2936a3d | ||
|
|
3f8eea0807 | ||
|
|
c063a4016e | ||
|
|
5fc0e486c3 | ||
|
|
3e69643392 | ||
|
|
3a012ac0ce | ||
|
|
80dd7288a6 | ||
|
|
5a6e826f83 | ||
|
|
d068b866f3 | ||
|
|
a618338987 | ||
|
|
b0c2c4d623 | ||
|
|
9663b5ccf0 | ||
|
|
6cfa0f4ef8 | ||
|
|
3149d62851 | ||
|
|
0b51cbe15d | ||
|
|
3dc2a9c371 | ||
|
|
2c13fb595e | ||
|
|
068239305e | ||
|
|
9e46f2d77e | ||
|
|
2eb7ccb6d2 | ||
|
|
745ff1d3ca | ||
|
|
8e2c4fad32 | ||
|
|
0ac774e378 | ||
|
|
ba3af61251 | ||
|
|
48c647b5d2 | ||
|
|
b94974467b | ||
|
|
4747827e35 | ||
|
|
f3bcad3c33 | ||
|
|
2c6bd63de9 | ||
|
|
02aad478f3 | ||
|
|
486b226b69 | ||
|
|
d1202bea5f | ||
|
|
1865f28469 | ||
|
|
c2f0edc137 | ||
|
|
cdbf1d6f4d | ||
|
|
b126b94525 | ||
|
|
7fc5332b45 | ||
|
|
4bcad0657b | ||
|
|
7ff5cff1bb | ||
|
|
bc63bd2874 | ||
|
|
458865ed1b | ||
|
|
3a7fd4ae78 | ||
|
|
4b8c557705 | ||
|
|
bb0d4a53c9 | ||
|
|
3c274768c1 | ||
|
|
bbf1d37783 | ||
|
|
90afaf120e | ||
|
|
6cacccc767 | ||
|
|
2afe7d4577 | ||
|
|
129369204e | ||
|
|
f19cd4fdea | ||
|
|
e473fba99c | ||
|
|
cc4a1e3f61 | ||
|
|
9d49d275f3 | ||
|
|
8d50466809 | ||
|
|
f9bd97cc15 | ||
|
|
f087fdb205 | ||
|
|
ef16f8a3d1 | ||
|
|
6db4171cfb | ||
|
|
db29b3a773 | ||
|
|
7cd3dc85c3 | ||
|
|
18182f1127 | ||
|
|
1817a8428a | ||
|
|
38dd23f834 | ||
|
|
6f6a27f9c6 | ||
|
|
61e1afdd49 | ||
|
|
8811c125ec | ||
|
|
3f2f66f328 | ||
|
|
f022a04573 | ||
|
|
795a7e2783 | ||
|
|
7b1aa93148 | ||
|
|
d5f490e106 | ||
|
|
3f1a2ea0a7 | ||
|
|
0363b8cdba | ||
|
|
0f9d460dc6 | ||
|
|
6e315616ed | ||
|
|
d97ccc952f | ||
|
|
33d9e72f90 | ||
|
|
df89228c2d | ||
|
|
480f668ba9 | ||
|
|
aa3b049000 | ||
|
|
0dccd7b6f5 | ||
|
|
d1aac6b989 | ||
|
|
54810f8c82 | ||
|
|
5ef517d0fc | ||
|
|
0173861a5a | ||
|
|
412718edf5 | ||
|
|
9f84204c21 | ||
|
|
679a7186fb | ||
|
|
bbba01dae9 | ||
|
|
da4419e626 | ||
|
|
b07b1fb5d9 | ||
|
|
81bd85ec55 | ||
|
|
b472042513 | ||
|
|
4551fb0d86 | ||
|
|
8f636513f7 | ||
|
|
f385bba928 | ||
|
|
856d7ea905 | ||
|
|
abbe42ac44 | ||
|
|
852f0f824f | ||
|
|
5800abf1cb | ||
|
|
abdfbb704c | ||
|
|
1c996fe301 | ||
|
|
4b6ed15a8b | ||
|
|
b8dece9000 | ||
|
|
3040209f20 | ||
|
|
f5aca74623 | ||
|
|
432a6a461f | ||
|
|
c2d5bbd93e | ||
|
|
ae501cc220 | ||
|
|
99efaa477f | ||
|
|
d9d4333795 | ||
|
|
18380398d5 | ||
|
|
8b9788eeb4 | ||
|
|
66a9e2e6bb | ||
|
|
dd7ee21703 | ||
|
|
c0e5a4670e | ||
|
|
458b92abd4 | ||
|
|
4407849142 | ||
|
|
1888de135d | ||
|
|
b6a157c19c | ||
|
|
775c1a3c1d | ||
|
|
15f469c86b | ||
|
|
7a767d5ce8 | ||
|
|
9a8ab4f88e | ||
|
|
f3d752693f | ||
|
|
71e97d1353 | ||
|
|
73381fc9c7 | ||
|
|
27f4a1780f | ||
|
|
3613174cf6 | ||
|
|
21d90fb5e5 | ||
|
|
75420462a6 | ||
|
|
c641453123 | ||
|
|
e0019da463 | ||
|
|
00a922c429 | ||
|
|
2e2b8f6d59 | ||
|
|
8167d85b32 | ||
|
|
8a495b0a54 | ||
|
|
8bc375b592 | ||
|
|
cee76030be | ||
|
|
3a1a37b7b9 | ||
|
|
7d0bf9c87f | ||
|
|
636faa695b | ||
|
|
e180a939ac | ||
|
|
5ace0699b2 | ||
|
|
a4ea8059b7 |
@@ -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),并致谢报告者(如你希望匿名则不署名)。
|
||||
|
||||
在公告发布前请勿公开漏洞细节。本项目由志愿者维护,响应时间尽力而为。
|
||||
@@ -3,16 +3,16 @@ exceptions:
|
||||
- package: xlsx
|
||||
advisory: "GHSA-4r6h-8v6p-xvw6"
|
||||
severity: high
|
||||
reason: "Admin export only; switched to dynamic import to reduce exposure (CVE-2023-30533)"
|
||||
mitigation: "Load only on export; restrict export permissions and data scope"
|
||||
expires_on: "2026-10-06"
|
||||
reason: "Write-only: admin UsageView lazily imports xlsx solely for XLSX.write export; the advisory requires parsing a crafted file and no XLSX.read path exists (CVE-2023-30533)"
|
||||
mitigation: "Never parse untrusted spreadsheets with xlsx; no patched version on npm, migrate to SheetJS CDN build or another writer before expiry"
|
||||
expires_on: "2027-01-06"
|
||||
owner: "security@your-domain"
|
||||
- package: xlsx
|
||||
advisory: "GHSA-5pgg-2g8v-p4x9"
|
||||
severity: high
|
||||
reason: "Admin export only; switched to dynamic import to reduce exposure (CVE-2024-22363)"
|
||||
mitigation: "Load only on export; restrict export permissions and data scope"
|
||||
expires_on: "2026-10-06"
|
||||
reason: "Write-only: admin UsageView lazily imports xlsx solely for XLSX.write export; the advisory requires parsing a crafted file and no XLSX.read path exists (CVE-2024-22363)"
|
||||
mitigation: "Never parse untrusted spreadsheets with xlsx; no patched version on npm, migrate to SheetJS CDN build or another writer before expiry"
|
||||
expires_on: "2027-01-06"
|
||||
owner: "security@your-domain"
|
||||
- package: lodash
|
||||
advisory: "GHSA-r5fr-rjxr-66jc"
|
||||
|
||||
@@ -35,7 +35,7 @@ jobs:
|
||||
cache-dependency-path: backend/go.sum
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.27.0'
|
||||
go version | grep -q 'go1.27.2'
|
||||
- name: Unit tests
|
||||
working-directory: backend
|
||||
run: make test-unit
|
||||
@@ -75,11 +75,11 @@ jobs:
|
||||
cache-dependency-path: backend/go.sum
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.27.0'
|
||||
go version | grep -q 'go1.27.2'
|
||||
- name: golangci-lint
|
||||
uses: golangci/golangci-lint-action@v9
|
||||
with:
|
||||
version: v2.13
|
||||
version: v2.14.0
|
||||
args: --timeout=30m
|
||||
working-directory: backend
|
||||
|
||||
|
||||
@@ -137,7 +137,7 @@ jobs:
|
||||
check-latest: false
|
||||
cache: false
|
||||
- name: Verify Go version
|
||||
run: go version | grep -q 'go1.27.0'
|
||||
run: go version | grep -q 'go1.27.2'
|
||||
- name: Restore target-specific Go cache
|
||||
id: go-cache
|
||||
uses: actions/cache/restore@v6
|
||||
|
||||
@@ -23,7 +23,7 @@ jobs:
|
||||
cache-dependency-path: backend/go.sum
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.27.0'
|
||||
go version | grep -q 'go1.27.2'
|
||||
- name: Run govulncheck
|
||||
working-directory: backend
|
||||
run: |
|
||||
|
||||
+1
-1
@@ -8,7 +8,7 @@
|
||||
# =============================================================================
|
||||
|
||||
ARG NODE_IMAGE=node:24-alpine
|
||||
ARG GOLANG_IMAGE=golang:1.27.0-alpine
|
||||
ARG GOLANG_IMAGE=golang:1.27.2-alpine
|
||||
ARG ALPINE_IMAGE=alpine:3.21
|
||||
ARG POSTGRES_IMAGE=postgres:18-alpine
|
||||
ARG GOPROXY=https://goproxy.cn,direct
|
||||
|
||||
@@ -8,6 +8,12 @@ FRONTEND_CRITICAL_VITEST := \
|
||||
src/components/account/__tests__/OpenAIReferralCell.spec.ts \
|
||||
src/components/account/__tests__/OpenAIReferralCell.transport.spec.ts \
|
||||
src/components/account/__tests__/OpenAIQuotaResetCell.spark_shadow.spec.ts \
|
||||
src/constants/__tests__/platforms.spec.ts \
|
||||
src/components/account/__tests__/credentialsBuilder.platformCatalog.spec.ts \
|
||||
src/components/account/__tests__/CreateAccountModal.spec.ts \
|
||||
src/components/account/__tests__/EditAccountModal.spec.ts \
|
||||
src/components/account/__tests__/credentialsBuilder.spec.ts \
|
||||
src/components/account/__tests__/OpenCodeGoProtocolRulesEditor.spec.ts \
|
||||
src/components/keys/__tests__/BulkEditKeysModal.spec.ts \
|
||||
src/components/admin/user/__tests__/UserPlatformQuotaModal.spec.ts \
|
||||
src/views/user/__tests__/KeysView.spec.ts \
|
||||
@@ -18,10 +24,12 @@ FRONTEND_CRITICAL_VITEST := \
|
||||
src/views/user/__tests__/PaymentResultView.spec.ts \
|
||||
src/views/user/__tests__/ChannelStatusView.mode.spec.ts \
|
||||
src/components/user/profile/__tests__/ProfileInfoCard.spec.ts \
|
||||
src/components/user/profile/__tests__/ProfileIdentityBindingsSection.spec.ts \
|
||||
src/views/admin/__tests__/SettingsView.spec.ts \
|
||||
src/features/channel-monitor-v2/__tests__/designSystem.structure.spec.ts \
|
||||
src/features/channel-monitor-v2/__tests__/monitorFormat.spec.ts \
|
||||
src/features/channel-monitor-v2/__tests__/monitorZoom.spec.ts
|
||||
src/features/channel-monitor-v2/__tests__/monitorZoom.spec.ts \
|
||||
src/components/admin/channel/__tests__/PricingEntryCard.modelDefaultPrice.spec.ts
|
||||
|
||||
# 一键编译前后端
|
||||
build: build-backend build-frontend
|
||||
|
||||
@@ -354,8 +354,10 @@ JWT_SECRET=your_jwt_secret_here
|
||||
TOTP_ENCRYPTION_KEY=your_totp_key_here
|
||||
|
||||
# Optional: Admin account
|
||||
ADMIN_EMAIL=admin@example.com
|
||||
ADMIN_PASSWORD=your_admin_password
|
||||
# Leave empty to auto-generate a random email (login username) and password, shown in logs on first startup.
|
||||
# Avoid guessable values such as admin@example.com: they are brute-force targets.
|
||||
ADMIN_EMAIL=
|
||||
ADMIN_PASSWORD=
|
||||
|
||||
# Optional: Custom port
|
||||
SERVER_PORT=8080
|
||||
@@ -404,9 +406,9 @@ docker compose -f docker-compose.local.yml logs -f sub2api
|
||||
|
||||
Open `http://YOUR_SERVER_IP:8080` in your browser.
|
||||
|
||||
If admin password was auto-generated, find it in logs:
|
||||
If the admin email (login username) or password was auto-generated, find them in logs:
|
||||
```bash
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "Generated admin"
|
||||
```
|
||||
|
||||
#### Upgrade
|
||||
|
||||
+6
-4
@@ -380,8 +380,10 @@ JWT_SECRET=your_jwt_secret_here
|
||||
TOTP_ENCRYPTION_KEY=your_totp_key_here
|
||||
|
||||
# 可选:管理员账号
|
||||
ADMIN_EMAIL=admin@example.com
|
||||
ADMIN_PASSWORD=your_admin_password
|
||||
# 留空则首次启动时自动生成随机邮箱(登录用户名)和密码,并输出到日志。
|
||||
# 请勿使用 admin@example.com 这类可被猜中的值,它们是暴力破解的目标。
|
||||
ADMIN_EMAIL=
|
||||
ADMIN_PASSWORD=
|
||||
|
||||
# 可选:自定义端口
|
||||
SERVER_PORT=8080
|
||||
@@ -442,9 +444,9 @@ docker compose -f docker-compose.local.yml logs -f sub2api
|
||||
|
||||
在浏览器中打开 `http://你的服务器IP:8080`
|
||||
|
||||
如果管理员密码是自动生成的,在日志中查找:
|
||||
如果管理员邮箱(登录用户名)或密码是自动生成的,在日志中查找:
|
||||
```bash
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "Generated admin"
|
||||
```
|
||||
|
||||
#### 升级
|
||||
|
||||
+6
-4
@@ -355,8 +355,10 @@ JWT_SECRET=your_jwt_secret_here
|
||||
TOTP_ENCRYPTION_KEY=your_totp_key_here
|
||||
|
||||
# オプション: 管理者アカウント
|
||||
ADMIN_EMAIL=admin@example.com
|
||||
ADMIN_PASSWORD=your_admin_password
|
||||
# 空欄の場合、初回起動時にランダムなメールアドレス(ログインユーザー名)とパスワードが自動生成され、ログに出力されます。
|
||||
# admin@example.com のような推測されやすい値は総当たり攻撃の標的になるため避けてください。
|
||||
ADMIN_EMAIL=
|
||||
ADMIN_PASSWORD=
|
||||
|
||||
# オプション: カスタムポート
|
||||
SERVER_PORT=8080
|
||||
@@ -405,9 +407,9 @@ docker compose -f docker-compose.local.yml logs -f sub2api
|
||||
|
||||
ブラウザで `http://YOUR_SERVER_IP:8080` を開いてください。
|
||||
|
||||
管理者パスワードが自動生成された場合は、ログで確認できます:
|
||||
管理者メールアドレス(ログインユーザー名)またはパスワードが自動生成された場合は、ログで確認できます:
|
||||
```bash
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "Generated admin"
|
||||
```
|
||||
|
||||
#### アップグレード
|
||||
|
||||
@@ -123,6 +123,15 @@ linters:
|
||||
|
||||
exclusions:
|
||||
rules:
|
||||
# x/net v0.60 deprecated the HTTP/2 compatibility adapters. Keep their
|
||||
# existing per-protocol idle timeout and PING configuration in this CI fix.
|
||||
- path: 'internal/(repository/http_upstream(_http2_keepalive_test)?|server/http)\.go$'
|
||||
linters: [ staticcheck ]
|
||||
text: 'SA1019: (\(golang\.org/x/net/http2\.(Server|Transport)\)\.|golang\.org/x/net/http2\.(ConfigureServer|ConfigureTransports|Server|Transport)\b)'
|
||||
# Upstream libraries can still return the legacy typed GOAWAY error.
|
||||
- path: 'internal/service/openai_codex_models_service(_test)?\.go$'
|
||||
linters: [ staticcheck ]
|
||||
text: 'SA1019: golang\.org/x/net/http2\.GoAwayError\b'
|
||||
# G703/G704 污点分析在测试文件中只会命中 httptest mock 与测试者自设的
|
||||
# 环境变量路径,不构成攻击面;且该分析跨环境结果不稳定(本地/CI 报告的
|
||||
# 位置集合不同),逐点 nolint 无法收敛,故按路径豁免。生产代码不豁免,
|
||||
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
FROM golang:1.27.0-alpine
|
||||
FROM golang:1.27.2-alpine
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
0.2.11
|
||||
0.2.14
|
||||
|
||||
@@ -906,7 +906,21 @@ func init() {
|
||||
// compositemodelroute.DefaultTargetPlatform holds the default value on creation for the target_platform field.
|
||||
compositemodelroute.DefaultTargetPlatform = compositemodelrouteDescTargetPlatform.Default.(string)
|
||||
// compositemodelroute.TargetPlatformValidator is a validator for the "target_platform" field. It is called by the builders before save.
|
||||
compositemodelroute.TargetPlatformValidator = compositemodelrouteDescTargetPlatform.Validators[0].(func(string) error)
|
||||
compositemodelroute.TargetPlatformValidator = func() func(string) error {
|
||||
validators := compositemodelrouteDescTargetPlatform.Validators
|
||||
fns := [...]func(string) error{
|
||||
validators[0].(func(string) error),
|
||||
validators[1].(func(string) error),
|
||||
}
|
||||
return func(target_platform string) error {
|
||||
for _, fn := range fns {
|
||||
if err := fn(target_platform); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}()
|
||||
// compositemodelrouteDescUpstreamModel is the schema descriptor for upstream_model field.
|
||||
compositemodelrouteDescUpstreamModel := compositemodelrouteFields[4].Descriptor()
|
||||
// compositemodelroute.DefaultUpstreamModel holds the default value on creation for the upstream_model field.
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package schema
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/ent/schema/mixins"
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
|
||||
@@ -45,6 +47,13 @@ func (CompositeModelRoute) Fields() []ent.Field {
|
||||
field.String("target_platform").
|
||||
MaxLen(50).
|
||||
Default(domain.PlatformOpenAI).
|
||||
Validate(func(s string) error {
|
||||
// 目标平台须为平台清单中的具体平台;数据库不再维护 CHECK 约束。
|
||||
if !domain.IsConcretePlatform(s) {
|
||||
return fmt.Errorf("target_platform %q is not a concrete platform", s)
|
||||
}
|
||||
return nil
|
||||
}).
|
||||
Comment("Concrete provider platform."),
|
||||
field.String("upstream_model").
|
||||
MaxLen(200).
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"entgo.io/ent/schema/field"
|
||||
"entgo.io/ent/schema/index"
|
||||
"github.com/Wei-Shaw/sub2api/ent/schema/mixins"
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
)
|
||||
|
||||
// UserPlatformQuota holds the schema definition for per-user per-platform quota.
|
||||
@@ -38,15 +39,12 @@ func (UserPlatformQuota) Fields() []ent.Field {
|
||||
MaxLen(32).
|
||||
NotEmpty().
|
||||
Validate(func(s string) error {
|
||||
// 注意:平台列表的单一权威源为 service.AllowedQuotaPlatforms;
|
||||
// 此处为 ent 构建期约束,需与 service.AllowedQuotaPlatforms 保持同步。
|
||||
switch s {
|
||||
case "anthropic", "openai", "gemini", "antigravity", "grok",
|
||||
"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe":
|
||||
return nil
|
||||
default:
|
||||
// 平台白名单为平台清单(domain/platforms.go)中的全部具体平台,
|
||||
// 与 service.AllowedQuotaPlatforms 同源;数据库不再维护 CHECK 约束。
|
||||
if !domain.IsConcretePlatform(s) {
|
||||
return fmt.Errorf("platform %q is not allowed", s)
|
||||
}
|
||||
return nil
|
||||
}),
|
||||
|
||||
// 日 / 周 / 月 USD 上限:
|
||||
|
||||
+10
-10
@@ -1,6 +1,6 @@
|
||||
module github.com/Wei-Shaw/sub2api
|
||||
|
||||
go 1.27.0
|
||||
go 1.27.2
|
||||
|
||||
require (
|
||||
entgo.io/ent v0.14.5
|
||||
@@ -19,6 +19,7 @@ require (
|
||||
github.com/coder/websocket v1.8.14
|
||||
github.com/dgraph-io/ristretto v0.2.0
|
||||
github.com/gin-gonic/gin v1.9.1
|
||||
github.com/go-playground/validator/v10 v10.14.0
|
||||
github.com/go-webauthn/webauthn v0.17.4
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1
|
||||
github.com/google/uuid v1.6.0
|
||||
@@ -50,12 +51,12 @@ require (
|
||||
github.com/wechatpay-apiv3/wechatpay-go v0.2.21
|
||||
github.com/zeromicro/go-zero v1.9.4
|
||||
go.uber.org/zap v1.24.0
|
||||
golang.org/x/crypto v0.55.0
|
||||
golang.org/x/crypto v0.57.0
|
||||
golang.org/x/image v0.41.0
|
||||
golang.org/x/mod v0.38.0
|
||||
golang.org/x/net v0.58.0
|
||||
golang.org/x/sync v0.22.0
|
||||
golang.org/x/term v0.45.0
|
||||
golang.org/x/mod v0.41.0
|
||||
golang.org/x/net v0.60.0
|
||||
golang.org/x/sync v0.23.0
|
||||
golang.org/x/term v0.46.0
|
||||
google.golang.org/grpc v1.83.2
|
||||
google.golang.org/protobuf v1.36.11
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1
|
||||
@@ -120,7 +121,6 @@ require (
|
||||
github.com/go-openapi/inflect v0.19.0 // indirect
|
||||
github.com/go-playground/locales v0.14.1 // indirect
|
||||
github.com/go-playground/universal-translator v0.18.1 // indirect
|
||||
github.com/go-playground/validator/v10 v10.14.0 // indirect
|
||||
github.com/go-viper/mapstructure/v2 v2.5.0 // indirect
|
||||
github.com/go-webauthn/x v0.2.6 // indirect
|
||||
github.com/goccy/go-json v0.10.2 // indirect
|
||||
@@ -201,10 +201,10 @@ require (
|
||||
go.uber.org/multierr v1.9.0 // indirect
|
||||
golang.org/x/arch v0.3.0 // indirect
|
||||
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect
|
||||
golang.org/x/sys v0.47.0 // indirect
|
||||
golang.org/x/text v0.41.0 // indirect
|
||||
golang.org/x/sys v0.48.0 // indirect
|
||||
golang.org/x/text v0.42.0 // indirect
|
||||
golang.org/x/time v0.12.0 // indirect
|
||||
golang.org/x/tools v0.48.0 // indirect
|
||||
golang.org/x/tools v0.50.0 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect
|
||||
gopkg.in/ini.v1 v1.67.0 // indirect
|
||||
modernc.org/libc v1.67.6 // indirect
|
||||
|
||||
@@ -564,6 +564,8 @@ golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v
|
||||
golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
|
||||
golang.org/x/crypto v0.55.0 h1:+KWHjbgOaAQ66dh/YlkZKHlz9ZUlq61AFirAR9ntP8M=
|
||||
golang.org/x/crypto v0.55.0/go.mod h1:uq0V9dE/fzQuJtbnL+2EhWOE63vo164FY8xqEnV9xis=
|
||||
golang.org/x/crypto v0.57.0 h1:3ZVCjf8Ggz7zneR/EHRVx68Ctf+2pmIMP2UFhh9cC6M=
|
||||
golang.org/x/crypto v0.57.0/go.mod h1:Fdz0i5U6CoizGwLda9DttjSk6qlZo25zYNtR+ycvuZA=
|
||||
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY=
|
||||
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70=
|
||||
@@ -580,6 +582,8 @@ golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/mod v0.38.0 h1:MECBjubtXD7yj4HrhIUcywNaGeNVUdfVnxmPajOk4yk=
|
||||
golang.org/x/mod v0.38.0/go.mod h1:V6Xz0pq8TQ3dGqVQ1FVHuelZpAL0uNhSkk9ogYP3c40=
|
||||
golang.org/x/mod v0.41.0 h1:qJmnOUb4YB+FsEuM3HcWucdZASCPGhsX6uljO6pog0c=
|
||||
golang.org/x/mod v0.41.0/go.mod h1:Ek9pY8RKWXwsWvd3rQiHYtMqkjSUV+s1Rj7j4H5Ur6o=
|
||||
golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
|
||||
golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
|
||||
golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
|
||||
@@ -603,6 +607,8 @@ golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
|
||||
golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
|
||||
golang.org/x/net v0.58.0 h1:ynWG7rqYi4ccpTEuPZ2QGWHktVEM9DMCj9yzDE0Q7To=
|
||||
golang.org/x/net v0.58.0/go.mod h1:YwCddHnFlT7eLQqVprV19OnhLGtc5xOKgE0RyqgfWAU=
|
||||
golang.org/x/net v0.60.0 h1:79p50tfZlm0J9YfoDsSi639qSXNGVwEzOPLCxM2FsYU=
|
||||
golang.org/x/net v0.60.0/go.mod h1:2DA/G1UfVbCpQPeWTmMPGY7Cs2PkBkwu743bVX5PIVg=
|
||||
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
|
||||
golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
@@ -616,6 +622,8 @@ golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
|
||||
golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk=
|
||||
golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0=
|
||||
golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
@@ -649,6 +657,8 @@ golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo=
|
||||
golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og=
|
||||
golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
@@ -663,6 +673,8 @@ golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
|
||||
golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0=
|
||||
golang.org/x/term v0.45.0 h1:NwWyBmoJCbfTHpxrWoZ9C6/VxOf7ic219I8xZZFdrf0=
|
||||
golang.org/x/term v0.45.0/go.mod h1:9aqxs0blBcrm/n0L9QW0aRVD+ktan8ssZromtqJC43w=
|
||||
golang.org/x/term v0.46.0 h1:3+OXuTbaKDgwk8jTi3aSLHRlmWqHEUDUtxnbFigO4YE=
|
||||
golang.org/x/term v0.46.0/go.mod h1:+K02xbkittuwc0Am4abfA3Fc+XRGXkvBXNO88NCXPoc=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
|
||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
@@ -675,6 +687,8 @@ golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
|
||||
golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8=
|
||||
golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M=
|
||||
golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI=
|
||||
golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E=
|
||||
golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE=
|
||||
golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
@@ -691,6 +705,10 @@ golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58
|
||||
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
|
||||
golang.org/x/tools v0.48.0 h1:3+hClM1aLL5mjMKm5ovokw9epgRXPuu2tILgismM6RE=
|
||||
golang.org/x/tools v0.48.0/go.mod h1:08xX0orndb/F7jJxGDicx061tyd5pcMto75YMAXr6lk=
|
||||
golang.org/x/tools v0.49.0 h1:3NI7VXzL9+1WZD52Dx2ttoPwD5DWrFGpl9mFZDlmisI=
|
||||
golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo=
|
||||
golang.org/x/tools v0.50.0 h1:c2ifzfcuY7L90lZ2aKd8S4K2NpASF08SZx9ZuJkHmSU=
|
||||
golang.org/x/tools v0.50.0/go.mod h1:7ulVMw3831Mwi5EZD6RomGyffr4VFjuNYXf2BbCEAV0=
|
||||
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
|
||||
@@ -30,9 +30,15 @@ const (
|
||||
PlatformMiniMax = "minimax" // MiniMax (M 系列)
|
||||
PlatformTypeSafe = "typesafe" // TypeSafe AI System One (Jev)
|
||||
// PlatformOpenCodeGo 是 OpenCode 平台(账号类型 Zen 按量 / Go 订阅)。
|
||||
// 值保持 opencode_go 以兼容已落库的分组、配额与 Composite 路由 CHECK。
|
||||
// 值保持 opencode_go 以兼容已落库的账号、分组、配额与 Composite 路由。
|
||||
PlatformOpenCodeGo = "opencode_go"
|
||||
PlatformComposite = "composite"
|
||||
// PlatformCommandCode 是 Command Code Provider API(多模型聚合,积分计费,
|
||||
// 订阅套餐另有 5 小时 / 每周滚动窗口)。
|
||||
PlatformCommandCode = "command_code"
|
||||
// PlatformCline 是 Cline API(多模型聚合,OpenAI 兼容 Chat Completions;按量积分计费,
|
||||
// 另有 ClinePass 订阅)。
|
||||
PlatformCline = "cline"
|
||||
PlatformComposite = "composite"
|
||||
)
|
||||
|
||||
// Account mode constants 区分国产供应商的「按量付费(余额)」与「Coding Plan」两种接入方式。
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
package domain
|
||||
|
||||
// 平台清单:平台的单一数据来源。"所有平台"类列表(分组平台、用户额度、组合
|
||||
// 路由目标、调度快照、定价匹配等)与"某网关族"类列表都从这里派生;新增平台
|
||||
// 登记一条即可在这些位置生效,无需逐处修改或新增数据库约束迁移。
|
||||
//
|
||||
// 依赖具体实现的能力(OAuth、渠道监控、余额 / Coding Plan 探测、用量窗口、
|
||||
// Codex 模型目录等)仍由各自的代码显式列出,不在此处声明。
|
||||
|
||||
// PlatformGateway 标识平台请求由哪个网关服务处理。
|
||||
type PlatformGateway string
|
||||
|
||||
const (
|
||||
// PlatformGatewayAnthropic 由 GatewayService 处理(Claude Messages 协议,以及 TypeSafe 的 /v1/systemone)。
|
||||
PlatformGatewayAnthropic PlatformGateway = "anthropic"
|
||||
// PlatformGatewayOpenAI 由 OpenAIGatewayService 处理(OpenAI、Grok 与多协议 API Key 供应商)。
|
||||
PlatformGatewayOpenAI PlatformGateway = "openai"
|
||||
// PlatformGatewayGemini 由 Gemini 兼容网关处理。
|
||||
PlatformGatewayGemini PlatformGateway = "gemini"
|
||||
// PlatformGatewayAntigravity 由 Antigravity 网关处理。
|
||||
PlatformGatewayAntigravity PlatformGateway = "antigravity"
|
||||
)
|
||||
|
||||
// PlatformSpec 描述一个具体(非 composite)平台。
|
||||
type PlatformSpec struct {
|
||||
ID string `json:"id"`
|
||||
DisplayName string `json:"display_name"`
|
||||
Gateway PlatformGateway `json:"gateway"`
|
||||
// CNProvider 标识国产单一厂商(Kimi / 智谱 / DeepSeek / MiniMax),共享国产
|
||||
// 供应商的余额 / Coding Plan 等能力。
|
||||
CNProvider bool `json:"cn_provider"`
|
||||
// LiteLLMProvider 是定价目录中的 provider 名;空表示不支持从定价目录同步模型。
|
||||
LiteLLMProvider string `json:"litellm_provider,omitempty"`
|
||||
}
|
||||
|
||||
// platformList 的顺序即平台的展示顺序。
|
||||
var platformList = []PlatformSpec{
|
||||
{ID: PlatformAnthropic, DisplayName: "Anthropic", Gateway: PlatformGatewayAnthropic, LiteLLMProvider: "anthropic"},
|
||||
{ID: PlatformOpenAI, DisplayName: "OpenAI", Gateway: PlatformGatewayOpenAI, LiteLLMProvider: "openai"},
|
||||
{ID: PlatformGemini, DisplayName: "Gemini", Gateway: PlatformGatewayGemini, LiteLLMProvider: "gemini"},
|
||||
{ID: PlatformAntigravity, DisplayName: "Antigravity", Gateway: PlatformGatewayAntigravity, LiteLLMProvider: "anthropic"},
|
||||
{ID: PlatformGrok, DisplayName: "Grok", Gateway: PlatformGatewayOpenAI, LiteLLMProvider: "xai"},
|
||||
{ID: PlatformKimi, DisplayName: "Kimi", Gateway: PlatformGatewayOpenAI, CNProvider: true, LiteLLMProvider: "moonshot"},
|
||||
{ID: PlatformZhipu, DisplayName: "Zhipu GLM", Gateway: PlatformGatewayOpenAI, CNProvider: true, LiteLLMProvider: "zhipu"},
|
||||
{ID: PlatformDeepseek, DisplayName: "DeepSeek", Gateway: PlatformGatewayOpenAI, CNProvider: true, LiteLLMProvider: "deepseek"},
|
||||
{ID: PlatformMiniMax, DisplayName: "MiniMax", Gateway: PlatformGatewayOpenAI, CNProvider: true, LiteLLMProvider: "minimax"},
|
||||
{ID: PlatformOpenCodeGo, DisplayName: "OpenCode", Gateway: PlatformGatewayOpenAI, LiteLLMProvider: "opencode-go"},
|
||||
{ID: PlatformTypeSafe, DisplayName: "TypeSafe / Jev", Gateway: PlatformGatewayAnthropic, LiteLLMProvider: "typesafe"},
|
||||
{ID: PlatformCommandCode, DisplayName: "Command Code", Gateway: PlatformGatewayOpenAI},
|
||||
{ID: PlatformCline, DisplayName: "Cline", Gateway: PlatformGatewayOpenAI},
|
||||
}
|
||||
|
||||
var platformIndex = func() map[string]int {
|
||||
index := make(map[string]int, len(platformList))
|
||||
for i, spec := range platformList {
|
||||
index[spec.ID] = i
|
||||
}
|
||||
return index
|
||||
}()
|
||||
|
||||
// Platforms 返回全部具体平台(副本),按展示顺序。
|
||||
func Platforms() []PlatformSpec {
|
||||
out := make([]PlatformSpec, len(platformList))
|
||||
copy(out, platformList)
|
||||
return out
|
||||
}
|
||||
|
||||
// LookupPlatform 返回具体平台的登记信息;composite 与未知值返回 false。
|
||||
func LookupPlatform(id string) (PlatformSpec, bool) {
|
||||
i, ok := platformIndex[id]
|
||||
if !ok {
|
||||
return PlatformSpec{}, false
|
||||
}
|
||||
return platformList[i], true
|
||||
}
|
||||
|
||||
// IsConcretePlatform 报告 id 是否为已登记的具体平台(不含 composite)。
|
||||
func IsConcretePlatform(id string) bool {
|
||||
_, ok := platformIndex[id]
|
||||
return ok
|
||||
}
|
||||
|
||||
// IsGroupPlatform 报告 id 是否可作为分组平台(具体平台或 composite)。
|
||||
func IsGroupPlatform(id string) bool {
|
||||
return id == PlatformComposite || IsConcretePlatform(id)
|
||||
}
|
||||
|
||||
// ConcretePlatformIDs 返回全部具体平台标识,按展示顺序。
|
||||
func ConcretePlatformIDs() []string {
|
||||
return PlatformIDsWhere(func(PlatformSpec) bool { return true })
|
||||
}
|
||||
|
||||
// PlatformIDsWhere 返回满足条件的具体平台标识,按展示顺序。
|
||||
func PlatformIDsWhere(match func(PlatformSpec) bool) []string {
|
||||
out := make([]string, 0, len(platformList))
|
||||
for _, spec := range platformList {
|
||||
if match(spec) {
|
||||
out = append(out, spec.ID)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// compositePrecedenceHead 是组合分组查找顺序的固定前缀(历史顺序:Gemini 先于 OpenAI)。
|
||||
var compositePrecedenceHead = []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI}
|
||||
|
||||
// CompositePrecedencePlatformIDs 返回按平台依次回退时的顺序:组合分组在请求目标
|
||||
// 解析前逐平台查找定价 / 模型映射(先命中者生效)、拼接默认模型列表,以及调度快照
|
||||
// 遍历。沿用历史顺序 anthropic、gemini、openai,其余按展示顺序,新登记的平台追加在末尾。
|
||||
func CompositePrecedencePlatformIDs() []string {
|
||||
out := make([]string, 0, len(platformList))
|
||||
seen := make(map[string]struct{}, len(platformList))
|
||||
for _, id := range compositePrecedenceHead {
|
||||
if IsConcretePlatform(id) {
|
||||
out = append(out, id)
|
||||
seen[id] = struct{}{}
|
||||
}
|
||||
}
|
||||
for _, spec := range platformList {
|
||||
if _, ok := seen[spec.ID]; !ok {
|
||||
out = append(out, spec.ID)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// UsesOpenAIGateway 报告平台是否由 OpenAI 网关处理。
|
||||
func UsesOpenAIGateway(id string) bool {
|
||||
spec, ok := LookupPlatform(id)
|
||||
return ok && spec.Gateway == PlatformGatewayOpenAI
|
||||
}
|
||||
|
||||
// IsCNProviderPlatform 报告平台是否为国产单一厂商。
|
||||
func IsCNProviderPlatform(id string) bool {
|
||||
spec, ok := LookupPlatform(id)
|
||||
return ok && spec.CNProvider
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package domain
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 重构前各处手写的平台列表,作为平台清单派生结果的等价基准;重构后新登记的平台
|
||||
// (Command Code、Cline)按同类平台(OpenCode)的位置补入。
|
||||
var (
|
||||
legacyDisplayOrder = []string{
|
||||
PlatformAnthropic, PlatformOpenAI, PlatformGemini, PlatformAntigravity, PlatformGrok,
|
||||
PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe,
|
||||
PlatformCommandCode,
|
||||
PlatformCline,
|
||||
}
|
||||
legacyCompositePrecedence = []string{
|
||||
PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok,
|
||||
PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe,
|
||||
PlatformCommandCode,
|
||||
PlatformCline,
|
||||
}
|
||||
legacyOpenAIGateway = []string{
|
||||
PlatformOpenAI, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo,
|
||||
PlatformCommandCode,
|
||||
PlatformCline,
|
||||
}
|
||||
legacyCNProviders = []string{PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax}
|
||||
legacyLiteLLMByPlatf = map[string]string{
|
||||
PlatformAnthropic: "anthropic", PlatformOpenAI: "openai", PlatformGemini: "gemini",
|
||||
PlatformAntigravity: "anthropic", PlatformGrok: "xai", PlatformKimi: "moonshot",
|
||||
PlatformZhipu: "zhipu", PlatformDeepseek: "deepseek", PlatformMiniMax: "minimax",
|
||||
PlatformOpenCodeGo: "opencode-go", PlatformTypeSafe: "typesafe", PlatformCommandCode: "", PlatformCline: "",
|
||||
}
|
||||
)
|
||||
|
||||
func TestPlatformListMatchesLegacyLists(t *testing.T) {
|
||||
require.Equal(t, legacyDisplayOrder, ConcretePlatformIDs())
|
||||
require.Equal(t, legacyCompositePrecedence, CompositePrecedencePlatformIDs())
|
||||
require.Equal(t, legacyOpenAIGateway, PlatformIDsWhere(func(spec PlatformSpec) bool {
|
||||
return spec.Gateway == PlatformGatewayOpenAI
|
||||
}))
|
||||
require.Equal(t, legacyCNProviders, PlatformIDsWhere(func(spec PlatformSpec) bool { return spec.CNProvider }))
|
||||
for platform, provider := range legacyLiteLLMByPlatf {
|
||||
spec, ok := LookupPlatform(platform)
|
||||
require.True(t, ok, platform)
|
||||
require.Equal(t, provider, spec.LiteLLMProvider, platform)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlatformListPredicates(t *testing.T) {
|
||||
for _, platform := range legacyDisplayOrder {
|
||||
require.True(t, IsConcretePlatform(platform), platform)
|
||||
require.True(t, IsGroupPlatform(platform), platform)
|
||||
spec, ok := LookupPlatform(platform)
|
||||
require.True(t, ok, platform)
|
||||
require.NotEmpty(t, spec.DisplayName, platform)
|
||||
require.NotEmpty(t, spec.Gateway, platform)
|
||||
}
|
||||
require.False(t, IsConcretePlatform(PlatformComposite))
|
||||
require.True(t, IsGroupPlatform(PlatformComposite))
|
||||
for _, invalid := range []string{"", "moonshot", "Kimi", "openai ", "glm", "bogus"} {
|
||||
require.False(t, IsConcretePlatform(invalid), invalid)
|
||||
require.False(t, IsGroupPlatform(invalid), invalid)
|
||||
require.False(t, UsesOpenAIGateway(invalid), invalid)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlatformsReturnsCopy(t *testing.T) {
|
||||
platforms := Platforms()
|
||||
platforms[0].ID = "mutated"
|
||||
require.Equal(t, PlatformAnthropic, Platforms()[0].ID)
|
||||
}
|
||||
@@ -1226,7 +1226,7 @@ func (h *AccountHandler) Update(c *gin.Context) {
|
||||
// 网关会按"现状即证据"默认走 Responses。
|
||||
func (h *AccountHandler) scheduleOpenAIResponsesProbe(account *service.Account) {
|
||||
if account == nil || account.Type != service.AccountTypeAPIKey ||
|
||||
(account.Platform != service.PlatformOpenAI && !service.IsCNProvider(account.Platform)) {
|
||||
(account.Platform != service.PlatformOpenAI && !account.RoutesProtocolByInbound()) {
|
||||
return
|
||||
}
|
||||
if h.accountTestService == nil {
|
||||
@@ -2894,10 +2894,9 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// Handle Antigravity accounts: return Claude + Gemini models
|
||||
// Explicit account mappings expose their request-side names to connectivity tests.
|
||||
if account.Platform == service.PlatformAntigravity {
|
||||
// 直接复用 antigravity.DefaultModels(),与 /v1/models 端点保持同步
|
||||
response.Success(c, antigravity.DefaultModels())
|
||||
response.Success(c, antigravityAccountTestModels(account.Credentials["model_mapping"]))
|
||||
return
|
||||
}
|
||||
|
||||
@@ -2998,6 +2997,48 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
|
||||
response.Success(c, models)
|
||||
}
|
||||
|
||||
// antigravityAccountTestModels uses the stored mapping rather than GetModelMapping,
|
||||
// which supplies defaults and compatibility aliases that the administrator did not configure.
|
||||
func antigravityAccountTestModels(rawMapping any) []antigravity.ClaudeModel {
|
||||
var mappedIDs []string
|
||||
switch mapping := rawMapping.(type) {
|
||||
case map[string]any:
|
||||
for id := range mapping {
|
||||
if strings.TrimSpace(id) != "" {
|
||||
mappedIDs = append(mappedIDs, id)
|
||||
}
|
||||
}
|
||||
case map[string]string:
|
||||
for id := range mapping {
|
||||
if strings.TrimSpace(id) != "" {
|
||||
mappedIDs = append(mappedIDs, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(mappedIDs) == 0 {
|
||||
return antigravity.DefaultModels()
|
||||
}
|
||||
|
||||
sort.Strings(mappedIDs)
|
||||
defaultByID := make(map[string]antigravity.ClaudeModel)
|
||||
for _, model := range antigravity.DefaultModels() {
|
||||
defaultByID[model.ID] = model
|
||||
}
|
||||
models := make([]antigravity.ClaudeModel, 0, len(mappedIDs))
|
||||
for _, id := range mappedIDs {
|
||||
if model, ok := defaultByID[id]; ok {
|
||||
models = append(models, model)
|
||||
continue
|
||||
}
|
||||
models = append(models, antigravity.ClaudeModel{
|
||||
ID: id,
|
||||
Type: "model",
|
||||
DisplayName: id,
|
||||
})
|
||||
}
|
||||
return models
|
||||
}
|
||||
|
||||
// SyncUpstreamModels handles syncing live supported models from an account's upstream.
|
||||
// POST /api/v1/admin/accounts/:id/models/sync-upstream
|
||||
func (h *AccountHandler) SyncUpstreamModels(c *gin.Context) {
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -330,6 +331,49 @@ func TestAccountHandlerGetAvailableModels_GeminiGoogleOneUsesConservativeCatalog
|
||||
require.NotContains(t, ids, "gemini-2.5-flash-image")
|
||||
}
|
||||
|
||||
func TestAccountHandlerGetAvailableModels_Antigravity(t *testing.T) {
|
||||
defaults := antigravity.DefaultModels()
|
||||
known := defaults[0]
|
||||
custom := antigravity.ClaudeModel{ID: "claude-custom-alias", Type: "model", DisplayName: "claude-custom-alias"}
|
||||
cases := []struct {
|
||||
name string
|
||||
credentials map[string]any
|
||||
want []antigravity.ClaudeModel
|
||||
}{
|
||||
{name: "nil credentials", want: defaults},
|
||||
{name: "missing mapping", credentials: map[string]any{}, want: defaults},
|
||||
{name: "empty JSON mapping", credentials: map[string]any{"model_mapping": map[string]any{}}, want: defaults},
|
||||
{name: "empty string mapping", credentials: map[string]any{"model_mapping": map[string]string{}}, want: defaults},
|
||||
{name: "blank key", credentials: map[string]any{"model_mapping": map[string]any{" ": "upstream-model"}}, want: defaults},
|
||||
{name: "JSON mapping keys and metadata", credentials: map[string]any{"model_mapping": map[string]any{
|
||||
known.ID: known.ID, custom.ID: "upstream-model",
|
||||
}}, want: []antigravity.ClaudeModel{custom, known}},
|
||||
{name: "string mapping keys and metadata", credentials: map[string]any{"model_mapping": map[string]string{
|
||||
known.ID: known.ID, custom.ID: "upstream-model",
|
||||
}}, want: []antigravity.ClaudeModel{custom, known}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
for _, accountType := range []string{service.AccountTypeOAuth, service.AccountTypeAPIKey} {
|
||||
t.Run(accountType, func(t *testing.T) {
|
||||
svc := &availableModelsAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
account: service.Account{ID: 51, Platform: service.PlatformAntigravity, Type: accountType, Credentials: tc.credentials},
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
setupAvailableModelsRouter(svc).ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/51/models", nil))
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var resp struct {
|
||||
Data []antigravity.ClaudeModel `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, tc.want, resp.Data)
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccountHandlerSyncUpstreamModels_ConfigErrorReturnsBadRequest(t *testing.T) {
|
||||
svc := &availableModelsAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
@@ -631,20 +632,14 @@ func (h *ChannelHandler) GetModelDefaultPricing(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
// platformToLiteLLMProvider maps a channel platform name to the corresponding
|
||||
// LiteLLM provider string used as the key in the pricing catalog.
|
||||
var platformToLiteLLMProvider = map[string]string{
|
||||
service.PlatformAnthropic: "anthropic",
|
||||
service.PlatformOpenAI: "openai",
|
||||
service.PlatformGemini: "gemini",
|
||||
service.PlatformAntigravity: "anthropic",
|
||||
service.PlatformGrok: "xai",
|
||||
service.PlatformKimi: "moonshot",
|
||||
service.PlatformZhipu: "zhipu",
|
||||
service.PlatformDeepseek: "deepseek",
|
||||
service.PlatformMiniMax: "minimax",
|
||||
service.PlatformOpenCodeGo: "opencode-go",
|
||||
service.PlatformTypeSafe: "typesafe",
|
||||
// platformLiteLLMProvider maps a channel platform name to the corresponding
|
||||
// LiteLLM provider string used as the key in the pricing catalog (platform list).
|
||||
func platformLiteLLMProvider(platform string) (string, bool) {
|
||||
spec, ok := domain.LookupPlatform(platform)
|
||||
if !ok || spec.LiteLLMProvider == "" {
|
||||
return "", false
|
||||
}
|
||||
return spec.LiteLLMProvider, true
|
||||
}
|
||||
|
||||
// SyncPricingModels 返回 LiteLLM 定价目录中指定平台的最新模型列表
|
||||
@@ -657,7 +652,7 @@ func (h *ChannelHandler) SyncPricingModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
provider, ok := platformToLiteLLMProvider[platform]
|
||||
provider, ok := platformLiteLLMProvider(platform)
|
||||
if !ok {
|
||||
response.ErrorFrom(c, infraerrors.BadRequest("UNSUPPORTED_PLATFORM",
|
||||
fmt.Sprintf("unsupported platform: %s", platform)).
|
||||
|
||||
@@ -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 typesafe composite"`
|
||||
Platform string `json:"platform" binding:"omitempty,group_platform"`
|
||||
RateMultiplier float64 `json:"rate_multiplier"`
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
|
||||
@@ -258,7 +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 typesafe composite"`
|
||||
Platform string `json:"platform" binding:"omitempty,group_platform"`
|
||||
RateMultiplier *float64 `json:"rate_multiplier"`
|
||||
IsExclusive *bool `json:"is_exclusive"`
|
||||
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
|
||||
@@ -332,7 +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 typesafe"`
|
||||
TargetPlatform string `json:"target_platform" binding:"required,concrete_platform"`
|
||||
UpstreamModel string `json:"upstream_model"`
|
||||
Endpoint string `json:"endpoint" binding:"omitempty,oneof=any messages count_tokens responses chat_completions embeddings images gemini"`
|
||||
Priority int `json:"priority"`
|
||||
|
||||
@@ -29,6 +29,7 @@ func TestGroupPlatformBinding_AllowedPlatforms(t *testing.T) {
|
||||
"anthropic", "openai", "gemini", "antigravity", "grok",
|
||||
"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "composite",
|
||||
"typesafe",
|
||||
"command_code", "cline",
|
||||
}
|
||||
for _, platform := range allowed {
|
||||
t.Run("create_"+platform, func(t *testing.T) {
|
||||
@@ -73,7 +74,7 @@ func TestGroupPlatformBinding_RejectsInvalidPlatforms(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestCompositeRouteTargetPlatform_AllowsConcreteProviders(t *testing.T) {
|
||||
for _, platform := range []string{"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe"} {
|
||||
for _, platform := range []string{"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe", "command_code", "cline"} {
|
||||
var req CompositeRouteRequest
|
||||
body := fmt.Sprintf(`{"public_model":"m","target_platform":%q}`, platform)
|
||||
require.NoError(t, bindGroupPlatformJSON(t, &req, body))
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/gin-gonic/gin/binding"
|
||||
"github.com/go-playground/validator/v10"
|
||||
)
|
||||
|
||||
// 平台相关的请求绑定校验标签,由平台清单(domain/platforms.go)驱动,替代在
|
||||
// binding 标签里手写 oneof 白名单(新增平台时易漏加,导致分组无法创建):
|
||||
// - group_platform:已登记的具体平台或 composite;
|
||||
// - concrete_platform:已登记的具体平台。
|
||||
func init() {
|
||||
engine, ok := binding.Validator.Engine().(*validator.Validate)
|
||||
if !ok {
|
||||
panic("platform validation requires a go-playground validator engine")
|
||||
}
|
||||
if err := engine.RegisterValidation("group_platform", func(fl validator.FieldLevel) bool {
|
||||
return domain.IsGroupPlatform(fl.Field().String())
|
||||
}); err != nil {
|
||||
panic(fmt.Errorf("register group_platform validation: %w", err))
|
||||
}
|
||||
if err := engine.RegisterValidation("concrete_platform", func(fl validator.FieldLevel) bool {
|
||||
return domain.IsConcretePlatform(fl.Field().String())
|
||||
}); err != nil {
|
||||
panic(fmt.Errorf("register concrete_platform validation: %w", err))
|
||||
}
|
||||
}
|
||||
@@ -1166,7 +1166,7 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
}
|
||||
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID, true)
|
||||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID, "", true)
|
||||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
if len(source) == 0 {
|
||||
@@ -1265,19 +1265,19 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
platform = group.Platform
|
||||
}
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(ctx, groupID, false)
|
||||
availableModels := h.compositeAvailableModels(ctx, groupID, service.CompositeRouteEndpointResponses, false)
|
||||
fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite)
|
||||
models := availableModels
|
||||
if len(models) == 0 {
|
||||
models = fallbackModels
|
||||
}
|
||||
if group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
if len(source) == 0 {
|
||||
source = fallbackModels
|
||||
}
|
||||
return group.ModelAllowlist.FilterForListing(source)
|
||||
models = group.ModelAllowlist.FilterForListing(models)
|
||||
}
|
||||
if len(availableModels) > 0 {
|
||||
return availableModels
|
||||
if filtered, err := h.gatewayService.FilterCompositeCodexModels(ctx, group.ID, models); err == nil {
|
||||
return filtered
|
||||
}
|
||||
return fallbackModels
|
||||
return models
|
||||
}
|
||||
|
||||
availableModels := h.gatewayService.GetAvailableModels(ctx, groupID, platform)
|
||||
@@ -1294,14 +1294,14 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
// compositeAvailableModels lists the models the composite group can serve.
|
||||
// includeSystemOne adds TypeSafe models, which only work through /v1/systemone;
|
||||
// LLM client catalogs (Codex) must exclude them.
|
||||
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64, includeSystemOne bool) []string {
|
||||
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64, endpoint string, includeSystemOne bool) []string {
|
||||
if h == nil || h.gatewayService == nil {
|
||||
return nil
|
||||
}
|
||||
seen := make(map[string]struct{})
|
||||
models := make([]string, 0)
|
||||
schedulablePlatforms := h.gatewayService.GetSchedulablePlatforms(ctx, groupID)
|
||||
for _, platform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax, service.PlatformOpenCodeGo, service.PlatformTypeSafe} {
|
||||
for _, platform := range domain.CompositePrecedencePlatformIDs() {
|
||||
if platform == service.PlatformTypeSafe && !includeSystemOne {
|
||||
continue
|
||||
}
|
||||
@@ -1325,6 +1325,16 @@ func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *
|
||||
models = append(models, model)
|
||||
}
|
||||
}
|
||||
// A route can expose a public ID that no account model mapping contains.
|
||||
// On lookup failure, retain the existing account-derived catalog only.
|
||||
if routeModels, err := h.gatewayService.GetCompositeRouteModels(ctx, groupID, endpoint, includeSystemOne); err == nil {
|
||||
for _, model := range routeModels {
|
||||
if _, ok := seen[model]; !ok {
|
||||
seen[model] = struct{}{}
|
||||
models = append(models, model)
|
||||
}
|
||||
}
|
||||
}
|
||||
return models
|
||||
}
|
||||
|
||||
@@ -1493,10 +1503,13 @@ func defaultModelIDsForPlatform(platform string) []string {
|
||||
case service.PlatformComposite:
|
||||
ids := make([]string, 0)
|
||||
seen := make(map[string]struct{})
|
||||
// TypeSafe is deliberately absent: jev-latest only works through
|
||||
// /v1/systemone, so the static fallback never advertises it to LLM
|
||||
// clients. compositeAvailableModels lists it when the group can serve it.
|
||||
for _, concretePlatform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax, service.PlatformOpenCodeGo} {
|
||||
for _, concretePlatform := range domain.CompositePrecedencePlatformIDs() {
|
||||
// TypeSafe is deliberately skipped: jev-latest only works through
|
||||
// /v1/systemone, so the static fallback never advertises it to LLM
|
||||
// clients. compositeAvailableModels lists it when the group can serve it.
|
||||
if concretePlatform == service.PlatformTypeSafe {
|
||||
continue
|
||||
}
|
||||
for _, id := range defaultModelIDsForPlatform(concretePlatform) {
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
|
||||
@@ -3,6 +3,7 @@ package handler
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -22,6 +23,16 @@ type gatewayModelsAccountRepoStub struct {
|
||||
byGroup map[int64][]service.Account
|
||||
}
|
||||
|
||||
type gatewayModelsRouteRepoStub struct {
|
||||
service.CompositeModelRouteRepository
|
||||
routes []service.CompositeModelRoute
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *gatewayModelsRouteRepoStub) ListByGroup(_ context.Context, _ int64, _ bool) ([]service.CompositeModelRoute, error) {
|
||||
return s.routes, s.err
|
||||
}
|
||||
|
||||
type gatewayModelsResponseForTest struct {
|
||||
Object string `json:"object"`
|
||||
Data []gatewayModelItemForTest `json:"data"`
|
||||
@@ -81,16 +92,80 @@ func (s *gatewayModelsAccountRepoStub) ListModelAvailabilityCandidates(ctx conte
|
||||
return s.ListSchedulableByGroupID(ctx, *groupID)
|
||||
}
|
||||
|
||||
func newGatewayModelsHandlerForTest(repo service.AccountRepository) *GatewayHandler {
|
||||
func newGatewayModelsHandlerForTest(repo service.AccountRepository, routeRepo ...service.CompositeModelRouteRepository) *GatewayHandler {
|
||||
var resolver *service.CompositeRouteResolver
|
||||
if len(routeRepo) > 0 {
|
||||
resolver = service.NewCompositeRouteResolver(routeRepo[0])
|
||||
}
|
||||
return &GatewayHandler{
|
||||
gatewayService: service.NewGatewayService(
|
||||
repo,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, resolver, nil, nil,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
func TestGatewayCompositeRouteModelsAppearInBothCatalogs(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 7784
|
||||
accounts := &gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{
|
||||
groupID: {{ID: 1, Platform: service.PlatformOpenAI, Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{"mapped-model": "upstream"},
|
||||
}}},
|
||||
}}
|
||||
routes := &gatewayModelsRouteRepoStub{routes: []service.CompositeModelRoute{
|
||||
{PublicModel: "route-only", MatchType: service.CompositeRouteMatchExact, Enabled: true},
|
||||
{PublicModel: "route-only", MatchType: service.CompositeRouteMatchExact, Enabled: true},
|
||||
{PublicModel: "mapped-model", MatchType: service.CompositeRouteMatchExact, Enabled: true},
|
||||
{PublicModel: "image-only", MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointImages, Enabled: true},
|
||||
{PublicModel: "prefix-", MatchType: service.CompositeRouteMatchPrefix, Enabled: true},
|
||||
{PublicModel: "disabled", MatchType: service.CompositeRouteMatchExact, Enabled: false},
|
||||
}}
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
repo service.CompositeModelRouteRepository
|
||||
allowlist []string
|
||||
want []string
|
||||
}{
|
||||
{"routes", routes, nil, []string{"mapped-model", "route-only", "image-only"}},
|
||||
{"allowlist", routes, []string{"route-only", "disabled", "unknown"}, []string{"route-only"}},
|
||||
{"lookup failure", &gatewayModelsRouteRepoStub{err: errors.New("unavailable")}, nil, []string{"mapped-model"}},
|
||||
{"allowlisted lookup failure", &gatewayModelsRouteRepoStub{err: errors.New("unavailable")}, []string{"mapped-model", "unknown"}, []string{"mapped-model"}},
|
||||
{"no route repo", nil, nil, []string{"mapped-model"}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
h := newGatewayModelsHandlerForTest(accounts, tc.repo)
|
||||
group := &service.Group{ID: groupID, Platform: service.PlatformComposite}
|
||||
if tc.allowlist != nil {
|
||||
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: tc.allowlist}
|
||||
}
|
||||
for _, endpoint := range []string{"models", "codex"} {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group})
|
||||
if endpoint == "models" {
|
||||
h.Models(c)
|
||||
var got gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Equal(t, tc.want, modelIDsForTest(got.Data))
|
||||
} else {
|
||||
h.CodexModels(c)
|
||||
var got codexModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
want := tc.want
|
||||
if tc.name == "routes" {
|
||||
want = []string{"mapped-model", "route-only"}
|
||||
}
|
||||
require.ElementsMatch(t, want, codexModelSlugsForTest(got.Models))
|
||||
}
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultModelIDsForCompositeIncludesAntigravityDefaults(t *testing.T) {
|
||||
antigravityIDs := defaultModelIDsForPlatform(service.PlatformAntigravity)
|
||||
require.NotEmpty(t, antigravityIDs)
|
||||
@@ -1570,3 +1645,188 @@ func TestDefaultModelIDsForPlatform_CompositeFallbackExcludesTypeSafe(t *testing
|
||||
require.NotContains(t, defaultModelIDsForPlatform(service.PlatformComposite), "jev-latest")
|
||||
require.NotContains(t, defaultCodexModelIDsForPlatform(service.PlatformComposite), "jev-latest")
|
||||
}
|
||||
|
||||
func TestGatewayModels_CompositeExactRoutesRespectTypeSafeIsolation(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 7784
|
||||
routes := &gatewayModelsRouteRepoStub{routes: []service.CompositeModelRoute{
|
||||
{PublicModel: "systemone-any-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointAny, Enabled: true},
|
||||
{PublicModel: "systemone-responses-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
|
||||
{PublicModel: "systemone-legacy-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Enabled: true},
|
||||
{PublicModel: "jev-latest", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Enabled: true},
|
||||
{PublicModel: "responses-alias", TargetPlatform: service.PlatformOpenAI, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
|
||||
{PublicModel: "image-alias", TargetPlatform: service.PlatformOpenAI, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointImages, Enabled: true},
|
||||
}}
|
||||
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{}, routes)
|
||||
wantModels := make([]string, 0, len(routes.routes))
|
||||
for _, route := range routes.routes {
|
||||
wantModels = append(wantModels, route.PublicModel)
|
||||
}
|
||||
for _, allowlisted := range []bool{false, true} {
|
||||
name := "unrestricted"
|
||||
if allowlisted {
|
||||
name = "allowlisted"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
group := &service.Group{ID: groupID, Platform: service.PlatformComposite}
|
||||
if allowlisted {
|
||||
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{
|
||||
"systemone-any-alias", "systemone-responses-alias", "systemone-legacy-alias", "jev-latest", "responses-alias", "image-alias",
|
||||
}}
|
||||
}
|
||||
for _, codex := range []bool{false, true} {
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
ctx.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group})
|
||||
if codex {
|
||||
h.CodexModels(ctx)
|
||||
var manifest codexModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &manifest))
|
||||
require.Equal(t, []string{"responses-alias"}, codexModelSlugsForTest(manifest.Models))
|
||||
} else {
|
||||
h.Models(ctx)
|
||||
var models gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &models))
|
||||
require.ElementsMatch(t, wantModels, modelIDsForTest(models.Data))
|
||||
}
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGatewayModels_CompositeCodexFiltersEffectiveResponsesRoutes(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 7793
|
||||
defaultModel := defaultModelIDsForPlatform(service.PlatformOpenAI)[0]
|
||||
for _, scenario := range []struct {
|
||||
name string
|
||||
model string
|
||||
accounts []service.Account
|
||||
routes []service.CompositeModelRoute
|
||||
}{
|
||||
{
|
||||
name: "colliding exact routes",
|
||||
model: "shared-alias",
|
||||
routes: []service.CompositeModelRoute{
|
||||
{PublicModel: "shared-alias", TargetPlatform: service.PlatformOpenAI, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointAny, Priority: 1, Enabled: true},
|
||||
{PublicModel: "shared-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Priority: 100, Enabled: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "account alias exact override",
|
||||
model: "account-alias",
|
||||
accounts: []service.Account{{ID: 1, Platform: service.PlatformOpenAI, Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{"account-alias": "gpt-5"},
|
||||
}}},
|
||||
routes: []service.CompositeModelRoute{
|
||||
{PublicModel: "account-alias", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "account alias prefix override",
|
||||
model: "account-alias",
|
||||
accounts: []service.Account{{ID: 1, Platform: service.PlatformOpenAI, Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{"account-alias": "gpt-5"},
|
||||
}}},
|
||||
routes: []service.CompositeModelRoute{
|
||||
{PublicModel: "account-", TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchPrefix, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "account default override",
|
||||
model: defaultModel,
|
||||
accounts: []service.Account{{ID: 1, Platform: service.PlatformOpenAI}},
|
||||
routes: []service.CompositeModelRoute{
|
||||
{PublicModel: defaultModel, TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "fallback default override",
|
||||
model: defaultModel,
|
||||
routes: []service.CompositeModelRoute{
|
||||
{PublicModel: defaultModel, TargetPlatform: service.PlatformTypeSafe, MatchType: service.CompositeRouteMatchPrefix, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true},
|
||||
},
|
||||
},
|
||||
} {
|
||||
t.Run(scenario.name, func(t *testing.T) {
|
||||
for _, allowlisted := range []bool{false, true} {
|
||||
name := "unrestricted"
|
||||
if allowlisted {
|
||||
name = "allowlisted"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
routes := append([]service.CompositeModelRoute(nil), scenario.routes...)
|
||||
if scenario.name != "fallback default override" {
|
||||
routes = append(routes, service.CompositeModelRoute{
|
||||
PublicModel: "responses-control", TargetPlatform: service.PlatformOpenAI, MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointResponses, Enabled: true,
|
||||
})
|
||||
}
|
||||
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{groupID: scenario.accounts},
|
||||
}, &gatewayModelsRouteRepoStub{routes: routes})
|
||||
group := &service.Group{ID: groupID, Platform: service.PlatformComposite}
|
||||
if allowlisted {
|
||||
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{scenario.model, "responses-control"}}
|
||||
}
|
||||
for _, codex := range []bool{false, true} {
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
ctx.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group})
|
||||
if codex {
|
||||
h.CodexModels(ctx)
|
||||
var manifest codexModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &manifest))
|
||||
slugs := codexModelSlugsForTest(manifest.Models)
|
||||
require.NotContains(t, slugs, scenario.model)
|
||||
if scenario.name != "fallback default override" {
|
||||
require.Contains(t, slugs, "responses-control")
|
||||
} else if allowlisted {
|
||||
require.Empty(t, slugs)
|
||||
}
|
||||
} else {
|
||||
h.Models(ctx)
|
||||
var models gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &models))
|
||||
require.Contains(t, modelIDsForTest(models.Data), scenario.model)
|
||||
}
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGatewayModels_CompositeCodexRouteLookupErrorKeepsFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 7793
|
||||
model := defaultCodexModelIDsForPlatform(service.PlatformComposite)[0]
|
||||
for _, allowlisted := range []bool{false, true} {
|
||||
name := "unrestricted"
|
||||
if allowlisted {
|
||||
name = "allowlisted"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{}, &gatewayModelsRouteRepoStub{err: errors.New("unavailable")})
|
||||
group := &service.Group{ID: groupID, Platform: service.PlatformComposite}
|
||||
if allowlisted {
|
||||
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{model}}
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
ctx.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group})
|
||||
h.CodexModels(ctx)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var manifest codexModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &manifest))
|
||||
require.Contains(t, codexModelSlugsForTest(manifest.Models), model)
|
||||
if allowlisted {
|
||||
require.Len(t, manifest.Models, 1)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1001,3 +1001,32 @@ func TestCodexModelsPinnedAccountsETagMatchReturns304(t *testing.T) {
|
||||
require.Equal(t, http.StatusNotModified, second.Code, second.Body.String())
|
||||
require.Empty(t, second.Body.Bytes())
|
||||
}
|
||||
|
||||
func TestCodexModelsDiscoveryProjectsWildcardMappingWithoutAllowlist(t *testing.T) {
|
||||
for _, mode := range []string{"ordinary", "pinned", "fallback"} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
account := newPinnedCodexAccount(1, service.StatusActive, true, false)
|
||||
account.Credentials["model_mapping"] = map[string]any{"gpt-*": "gpt-5.4"}
|
||||
upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
|
||||
1: `{"models":[{"slug":"gpt-5.4"},{"slug":"blocked"},{"slug":"codex-auto-review"},{"slug":"gpt-image-2"}]}`,
|
||||
}}
|
||||
handler := newPinnedCodexTestHandler([]service.Account{account}, upstream, 3)
|
||||
group := &service.Group{ID: 84, Platform: service.PlatformOpenAI,
|
||||
ModelAllowlist: service.GroupModelAllowlist{Models: []string{"blocked"}},
|
||||
}
|
||||
if mode != "ordinary" {
|
||||
group.CodexModelsManifestConfig = service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{1}}
|
||||
}
|
||||
if mode == "fallback" {
|
||||
group.CodexModelsManifestConfig.AccountIDs = []int64{999}
|
||||
group.CodexModelsManifestConfig.FallbackToScheduler = true
|
||||
}
|
||||
first := performCodexModelsRequestForGroup(t, handler, group, "")
|
||||
require.Equal(t, http.StatusOK, first.Code, first.Body.String())
|
||||
require.Equal(t, []string{"gpt-5.4"}, codexHandlerManifestSlugs(t, first))
|
||||
require.NotEmpty(t, first.Header().Get("ETag"))
|
||||
second := performCodexModelsRequestForGroup(t, handler, group, first.Header().Get("ETag"))
|
||||
require.Equal(t, http.StatusNotModified, second.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
@@ -330,9 +331,9 @@ func allowOpenAICompatibleMessagesDispatch(c *gin.Context, apiKey *service.APIKe
|
||||
|
||||
func openAICompatibleTextTargetAllowed(c *gin.Context, apiKey *service.APIKey, model string) bool {
|
||||
return compositeTargetPlatformAllowed(c, apiKey, model,
|
||||
service.PlatformOpenAI, service.PlatformGrok,
|
||||
service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek,
|
||||
service.PlatformMiniMax, service.PlatformOpenCodeGo)
|
||||
domain.PlatformIDsWhere(func(spec domain.PlatformSpec) bool {
|
||||
return spec.Gateway == domain.PlatformGatewayOpenAI
|
||||
})...)
|
||||
}
|
||||
|
||||
// isResponsesWebSocketCompositePlatform 限定 composite 分组在 Responses WebSocket
|
||||
@@ -2125,6 +2126,101 @@ func (p *openAIWSTurnPricing) currentOr(fallback time.Time) time.Time {
|
||||
return fallback
|
||||
}
|
||||
|
||||
// openAIWSTurnAPIKeyLookup 是后续 turn 重取 API Key 认证快照的入口
|
||||
// (APIKeyService.GetByKey:经 L1/L2 认证缓存,未命中才回源)。
|
||||
type openAIWSTurnAPIKeyLookup interface {
|
||||
GetByKey(ctx context.Context, key string) (*service.APIKey, error)
|
||||
}
|
||||
|
||||
// openAIWSTurnBillingAPIKeys 按 turn 号保存每个 turn 计费用的 API Key 快照。
|
||||
//
|
||||
// API Key 认证快照只在建连时取一次。若连接内所有 turn 共用它,分组定价(倍率、
|
||||
// 高峰、图片定价、利润门参数)会停留在建连时刻:管理员调价后,已打开的长连接
|
||||
// 在客户端重连前一直按旧价计费、按旧售价过利润门,而 HTTP 请求每次都经认证缓存
|
||||
// 取快照,缓存失效后即生效。
|
||||
//
|
||||
// 因此后续 turn 在 BeforeTurn(与 HTTP 请求进入认证中间件同位)经同一认证缓存
|
||||
// 重取快照,只采用其中的分组,作为该 turn 利润门准入与用量计费共同的分组;用户、
|
||||
// Key 限额、订阅仍沿用建连快照。只在同一把 Key、同一分组且平台与订阅类型未变时
|
||||
// 采用:Key 换组或分组改平台/订阅类型不是调价,该连接按建连分组调度,继续按建连
|
||||
// 快照计费;重取失败同样保留建连快照,不断连。
|
||||
//
|
||||
// 首轮沿用刚经认证中间件取得的建连快照;没有经过 BeforeTurn 的 turn 回退建连
|
||||
// 快照。只保留当前与上一个 turn:下一 turn 的 BeforeTurn 先于上一 turn 的
|
||||
// AfterTurn 执行时,上一 turn 仍取到自己的快照。
|
||||
type openAIWSTurnBillingAPIKeys struct {
|
||||
mu sync.Mutex
|
||||
keys map[int]*service.APIKey
|
||||
}
|
||||
|
||||
// begin 在 BeforeTurn 重装利润门前调用:后续 turn 重取计费分组并按 turn 记下,
|
||||
// 返回换入该分组的上下文供本 turn 的利润门使用。
|
||||
func (k *openAIWSTurnBillingAPIKeys) begin(ctx context.Context, apiKeyService *service.APIKeyService, turn int, conn *service.APIKey) context.Context {
|
||||
turnKey := conn
|
||||
if turn > 1 && apiKeyService != nil {
|
||||
turnKey = refreshOpenAIWSTurnBillingAPIKey(ctx, apiKeyService, conn)
|
||||
}
|
||||
k.set(turn, turnKey)
|
||||
return withOpenAIWSTurnBillingGroup(ctx, turnKey)
|
||||
}
|
||||
|
||||
func (k *openAIWSTurnBillingAPIKeys) set(turn int, key *service.APIKey) {
|
||||
k.mu.Lock()
|
||||
defer k.mu.Unlock()
|
||||
if k.keys == nil {
|
||||
k.keys = make(map[int]*service.APIKey, 2)
|
||||
}
|
||||
for t := range k.keys {
|
||||
if t < turn-1 {
|
||||
delete(k.keys, t)
|
||||
}
|
||||
}
|
||||
k.keys[turn] = key
|
||||
}
|
||||
|
||||
func (k *openAIWSTurnBillingAPIKeys) forTurn(turn int, conn *service.APIKey) *service.APIKey {
|
||||
k.mu.Lock()
|
||||
defer k.mu.Unlock()
|
||||
if key := k.keys[turn]; key != nil {
|
||||
return key
|
||||
}
|
||||
return conn
|
||||
}
|
||||
|
||||
// refreshOpenAIWSTurnBillingAPIKey 返回本 turn 计费用的 API Key:分组取当前认证
|
||||
// 快照,其余字段与建连快照共享。不满足采用条件时原样返回建连快照。
|
||||
func refreshOpenAIWSTurnBillingAPIKey(ctx context.Context, lookup openAIWSTurnAPIKeyLookup, conn *service.APIKey) *service.APIKey {
|
||||
if lookup == nil || conn == nil || conn.Key == "" || conn.GroupID == nil || conn.Group == nil {
|
||||
return conn
|
||||
}
|
||||
latest, err := lookup.GetByKey(ctx, conn.Key)
|
||||
if err != nil || latest == nil || latest.ID != conn.ID || latest.GroupID == nil || *latest.GroupID != *conn.GroupID {
|
||||
return conn
|
||||
}
|
||||
group := latest.Group
|
||||
if group == nil || group.ID != conn.Group.ID ||
|
||||
group.Platform != conn.Group.Platform ||
|
||||
group.SubscriptionType != conn.Group.SubscriptionType {
|
||||
return conn
|
||||
}
|
||||
turnKey := *conn
|
||||
turnKey.Group = group
|
||||
return &turnKey
|
||||
}
|
||||
|
||||
// withOpenAIWSTurnBillingGroup 把本 turn 的计费分组换进认证分组上下文,使利润门
|
||||
// 的售价与本 turn 计费同源。上下文里没有同 ID 的认证分组时不改动。
|
||||
func withOpenAIWSTurnBillingGroup(ctx context.Context, turnKey *service.APIKey) context.Context {
|
||||
if turnKey == nil || turnKey.Group == nil {
|
||||
return ctx
|
||||
}
|
||||
current, ok := ctx.Value(ctxkey.Group).(*service.Group)
|
||||
if !ok || current == nil || current == turnKey.Group || current.ID != turnKey.Group.ID {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, ctxkey.Group, turnKey.Group)
|
||||
}
|
||||
|
||||
// recordOpenAIProfitVeto 记录 OpenAI 侧选号循环的一次利润门终检否决:把账号
|
||||
// 加入本请求排除集并递增否决计数。返回 false 表示否决次数已达
|
||||
// maxProfitVetoAttempts,调用方必须停止重选并按「无可用账号」终止。
|
||||
@@ -2856,6 +2952,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// turn 级定价:首轮回退到 TurnStarted 的所属 turn 时刻;后续 turn 由
|
||||
// BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号。
|
||||
var turnPricing openAIWSTurnPricing
|
||||
// 后续 turn 的计费分组同样在 BeforeTurn 经认证缓存重取,使分组调价对
|
||||
// 已打开的连接生效。
|
||||
var turnBillingAPIKeys openAIWSTurnBillingAPIKeys
|
||||
// Passthrough ingress does not invoke BeforeTurn for the first frame.
|
||||
if err := checkSimpleModeTurnBilling(); err != nil {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "billing check failed")
|
||||
@@ -2943,7 +3042,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// 长连接跨峰谷/倍率刷新防护:每个 turn 按当前时刻重装门并复核
|
||||
// 当前账号,越线即要求客户端重连重选(连接绑定单一上游账号,
|
||||
// 无法中途换号)。本 turn 的准入与计费共用同一 pricingAt。
|
||||
turnCtx, turnAt := h.gatewayService.WithOpenAITurnPricingContext(ctx, apiKey.GroupID)
|
||||
turnBillingCtx := turnBillingAPIKeys.begin(ctx, h.apiKeyService, turn, apiKey)
|
||||
turnCtx, turnAt := h.gatewayService.WithOpenAITurnPricingContext(turnBillingCtx, apiKey.GroupID)
|
||||
if _, vetoed, reason := h.gatewayService.ProfitControlVetoLatest(turnCtx, account); vetoed {
|
||||
reqLog.Info("openai.websocket_turn_profit_vetoed",
|
||||
zap.Int("turn", turn),
|
||||
@@ -3059,11 +3159,12 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
turnRecordPricingAt := turnPricing.currentOr(turnStart)
|
||||
turnBillingAPIKey := turnBillingAPIKeys.forTurn(turn, apiKey)
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{
|
||||
Result: result,
|
||||
APIKey: apiKey,
|
||||
APIKey: turnBillingAPIKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
@@ -1948,6 +1949,12 @@ type openAIResponsesWSUsageLogCase struct {
|
||||
afterFirstUpstreamRequest func(channelSvc *service.ChannelService) error
|
||||
// group 覆盖 apiKey.Group(分组级模型白名单测试用);nil 保持原有无分组行为。
|
||||
group *service.Group
|
||||
// apiKeyService 非 nil 时模拟 API Key 认证中间件:连接认证快照经它按
|
||||
// apiKeyCredential 取得,并把其分组放入请求 ctx;handler 也使用它。
|
||||
apiKeyService *service.APIKeyService
|
||||
apiKeyCredential string
|
||||
// accountRateMultiplier 覆盖账号倍率(利润门测试用)。
|
||||
accountRateMultiplier *float64
|
||||
// firstFrameCloseExpected:首帧即被拒(连接被 1008 关闭),不期待任何响应帧。
|
||||
firstFrameCloseExpected bool
|
||||
// secondTurnCloseExpected:第二个 turn 被拒(连接被 1008 关闭)。
|
||||
@@ -2962,6 +2969,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
if tc.accountPlatform != "" {
|
||||
account.Platform = tc.accountPlatform
|
||||
}
|
||||
account.RateMultiplier = tc.accountRateMultiplier
|
||||
if strings.TrimSpace(tc.ingressMode) != "" {
|
||||
account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
|
||||
}
|
||||
@@ -3050,6 +3058,12 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 1701, Status: service.StatusActive},
|
||||
}
|
||||
if tc.apiKeyService != nil {
|
||||
h.apiKeyService = tc.apiKeyService
|
||||
authKey, err := tc.apiKeyService.GetByKey(context.Background(), tc.apiKeyCredential)
|
||||
require.NoError(t, err)
|
||||
apiKey = authKey
|
||||
}
|
||||
if tc.simpleModeRejectAtRead > 0 {
|
||||
apiKey.RateLimit5h = 1
|
||||
}
|
||||
@@ -3060,6 +3074,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
|
||||
if tc.apiKeyService != nil && apiKey.Group != nil {
|
||||
c.Request = c.Request.WithContext(context.WithValue(c.Request.Context(), ctxkey.Group, apiKey.Group))
|
||||
}
|
||||
c.Next()
|
||||
})
|
||||
router.GET("/openai/v1/responses", h.ResponsesWebSocket)
|
||||
|
||||
@@ -0,0 +1,293 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 分组调价对已打开的 WebSocket 连接生效:同一条连接的两个 turn 之间由管理端
|
||||
// 改分组倍率(并按分组失效认证缓存),第二个 turn 必须按新倍率计费、按新售价
|
||||
// 过利润门;HTTP 请求每次经认证缓存取快照,本来就是这个行为。
|
||||
|
||||
const (
|
||||
wsGroupPricingGroupID = int64(4201) // 与 runOpenAIResponsesWebSocketUsageLogCase 的调度分组一致
|
||||
wsGroupPricingKey = "sk-ws-group-pricing"
|
||||
wsGroupPricingFirst = `{"type":"response.create","model":"gpt-5.6-sol","stream":false}`
|
||||
wsGroupPricingSecond = `{"type":"response.create","model":"gpt-5.6-sol","stream":false}`
|
||||
)
|
||||
|
||||
// wsGroupPricingAPIKeyRepoStub 只实现认证查询与按分组列 Key(认证缓存失效用)。
|
||||
type wsGroupPricingAPIKeyRepoStub struct {
|
||||
service.APIKeyRepository
|
||||
mu sync.Mutex
|
||||
apiKey service.APIKey
|
||||
group service.Group
|
||||
err error
|
||||
lookups int
|
||||
}
|
||||
|
||||
func newWSGroupPricingAPIKeyRepoStub(rate float64) *wsGroupPricingAPIKeyRepoStub {
|
||||
groupID := wsGroupPricingGroupID
|
||||
return &wsGroupPricingAPIKeyRepoStub{
|
||||
apiKey: service.APIKey{
|
||||
ID: 1801,
|
||||
UserID: 1701,
|
||||
Key: wsGroupPricingKey,
|
||||
Status: service.StatusActive,
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 1701, Status: service.StatusActive, Concurrency: 1},
|
||||
},
|
||||
group: service.Group{
|
||||
ID: groupID,
|
||||
Name: "ws-group-pricing",
|
||||
Platform: service.PlatformOpenAI,
|
||||
Status: service.StatusActive,
|
||||
RateMultiplier: rate,
|
||||
Hydrated: true,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *wsGroupPricingAPIKeyRepoStub) GetByKeyForAuth(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.lookups++
|
||||
if s.err != nil {
|
||||
return nil, s.err
|
||||
}
|
||||
if key != s.apiKey.Key {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
apiKey := s.apiKey
|
||||
user := *s.apiKey.User
|
||||
apiKey.User = &user
|
||||
group := s.group
|
||||
apiKey.Group = &group
|
||||
return &apiKey, nil
|
||||
}
|
||||
|
||||
func (s *wsGroupPricingAPIKeyRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
|
||||
return []string{s.apiKey.Key}, nil
|
||||
}
|
||||
|
||||
// adminUpdate 模拟管理端改 Key/分组后按分组失效认证缓存(与 admin 分组更新同一入口)。
|
||||
func (s *wsGroupPricingAPIKeyRepoStub) adminUpdate(apiKeyService *service.APIKeyService, fn func(apiKey *service.APIKey, group *service.Group)) error {
|
||||
s.mu.Lock()
|
||||
fn(&s.apiKey, &s.group)
|
||||
s.mu.Unlock()
|
||||
apiKeyService.InvalidateAuthCacheByGroupID(context.Background(), wsGroupPricingGroupID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *wsGroupPricingAPIKeyRepoStub) lookupCount() int {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.lookups
|
||||
}
|
||||
|
||||
func newWSGroupPricingAPIKeyService(repo *wsGroupPricingAPIKeyRepoStub) *service.APIKeyService {
|
||||
return service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, &config.Config{})
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_GroupRateChangeAppliesToNextTurn(t *testing.T) {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeCtxPool} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
|
||||
apiKeyService := newWSGroupPricingAPIKeyService(repo)
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: wsGroupPricingFirst,
|
||||
secondPayload: wsGroupPricingSecond,
|
||||
ingressMode: mode,
|
||||
apiKeyService: apiKeyService,
|
||||
apiKeyCredential: wsGroupPricingKey,
|
||||
afterFirstUpstreamRequest: func(*service.ChannelService) error {
|
||||
return repo.adminUpdate(apiKeyService, func(_ *service.APIKey, group *service.Group) { group.RateMultiplier = 0.3 })
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.logs, 2)
|
||||
require.InDelta(t, 3.0, got.logs[0].RateMultiplier, 1e-12, "turn 1 bills at the rate in force at connect")
|
||||
require.InDelta(t, 0.3, got.logs[1].RateMultiplier, 1e-12, "turn 2 must bill at the group rate in force when it starts")
|
||||
require.Greater(t, got.logs[0].ActualCost, 0.0)
|
||||
require.InDelta(t, got.logs[0].ActualCost/10, got.logs[1].ActualCost, 1e-15)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_GroupRateChangeReachesProfitGateOnNextTurn(t *testing.T) {
|
||||
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
|
||||
repo.group.ProfitControlEnabled = true
|
||||
repo.group.ProfitMinMargin = 0.1
|
||||
apiKeyService := newWSGroupPricingAPIKeyService(repo)
|
||||
accountRate := 1.0
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: wsGroupPricingFirst,
|
||||
secondPayload: wsGroupPricingSecond,
|
||||
apiKeyService: apiKeyService,
|
||||
apiKeyCredential: wsGroupPricingKey,
|
||||
accountRateMultiplier: &accountRate,
|
||||
// 3.0 → 0.3:账号倍率 1.0 在新售价下已越线,第二个 turn 必须在转发上游前被拒。
|
||||
afterFirstUpstreamRequest: func(*service.ChannelService) error {
|
||||
return repo.adminUpdate(apiKeyService, func(_ *service.APIKey, group *service.Group) { group.RateMultiplier = 0.3 })
|
||||
},
|
||||
secondTurnCloseExpected: true,
|
||||
closeStatus: coderws.StatusTryAgainLater,
|
||||
closeReason: "no longer eligible for this connection",
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_KeyMovedToAnotherGroupKeepsConnectionGroup(t *testing.T) {
|
||||
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
|
||||
apiKeyService := newWSGroupPricingAPIKeyService(repo)
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: wsGroupPricingFirst,
|
||||
secondPayload: wsGroupPricingSecond,
|
||||
apiKeyService: apiKeyService,
|
||||
apiKeyCredential: wsGroupPricingKey,
|
||||
afterFirstUpstreamRequest: func(*service.ChannelService) error {
|
||||
return repo.adminUpdate(apiKeyService, func(apiKey *service.APIKey, group *service.Group) {
|
||||
other := wsGroupPricingGroupID + 1
|
||||
apiKey.GroupID = &other
|
||||
group.ID = other
|
||||
group.RateMultiplier = 0.3
|
||||
})
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.logs, 2)
|
||||
require.InDelta(t, 3.0, got.logs[1].RateMultiplier, 1e-12,
|
||||
"a key moved to another group keeps the connection's group: the connection was scheduled from it")
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_FailedKeyRefreshKeepsConnectionGroup(t *testing.T) {
|
||||
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
|
||||
apiKeyService := newWSGroupPricingAPIKeyService(repo)
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: wsGroupPricingFirst,
|
||||
secondPayload: wsGroupPricingSecond,
|
||||
apiKeyService: apiKeyService,
|
||||
apiKeyCredential: wsGroupPricingKey,
|
||||
afterFirstUpstreamRequest: func(*service.ChannelService) error {
|
||||
return repo.adminUpdate(apiKeyService, func(_ *service.APIKey, group *service.Group) {
|
||||
group.RateMultiplier = 0.3
|
||||
repo.err = errors.New("auth store unavailable")
|
||||
})
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.logs, 2)
|
||||
require.InDelta(t, 3.0, got.logs[1].RateMultiplier, 1e-12, "a failed refresh keeps the connection snapshot")
|
||||
require.GreaterOrEqual(t, repo.lookupCount(), 2, "the second turn must have attempted a refresh")
|
||||
}
|
||||
|
||||
type wsTurnAPIKeyLookupFunc func(ctx context.Context, key string) (*service.APIKey, error)
|
||||
|
||||
func (f wsTurnAPIKeyLookupFunc) GetByKey(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return f(ctx, key)
|
||||
}
|
||||
|
||||
func TestRefreshOpenAIWSTurnBillingAPIKey(t *testing.T) {
|
||||
groupID := int64(77)
|
||||
conn := &service.APIKey{
|
||||
ID: 5,
|
||||
Key: "sk-conn",
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 9},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, SubscriptionType: service.SubscriptionTypeStandard, RateMultiplier: 3},
|
||||
}
|
||||
latestWith := func(mutate func(k *service.APIKey)) *service.APIKey {
|
||||
gid := groupID
|
||||
k := &service.APIKey{
|
||||
ID: 5,
|
||||
Key: "sk-conn",
|
||||
GroupID: &gid,
|
||||
User: &service.User{ID: 9, Balance: 123},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, SubscriptionType: service.SubscriptionTypeStandard, RateMultiplier: 0.3},
|
||||
}
|
||||
if mutate != nil {
|
||||
mutate(k)
|
||||
}
|
||||
return k
|
||||
}
|
||||
|
||||
t.Run("same group adopts the latest group and nothing else", func(t *testing.T) {
|
||||
got := refreshOpenAIWSTurnBillingAPIKey(context.Background(), wsTurnAPIKeyLookupFunc(func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
require.Equal(t, "sk-conn", key)
|
||||
return latestWith(nil), nil
|
||||
}), conn)
|
||||
require.NotSame(t, conn, got)
|
||||
require.InDelta(t, 0.3, got.Group.RateMultiplier, 1e-12)
|
||||
require.Same(t, conn.User, got.User, "only the group snapshot is refreshed")
|
||||
require.Equal(t, conn.ID, got.ID)
|
||||
require.InDelta(t, 3.0, conn.Group.RateMultiplier, 1e-12, "the connection snapshot is never mutated")
|
||||
})
|
||||
|
||||
keep := map[string]wsTurnAPIKeyLookupFunc{
|
||||
"lookup error": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return nil, errors.New("boom")
|
||||
},
|
||||
"not found": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
},
|
||||
"nil key": func(ctx context.Context, key string) (*service.APIKey, error) { return nil, nil },
|
||||
"different key id": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { k.ID = 6 }), nil
|
||||
},
|
||||
"moved group": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { other := int64(78); k.GroupID = &other; k.Group.ID = other }), nil
|
||||
},
|
||||
"no group": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { k.Group = nil }), nil
|
||||
},
|
||||
"platform changed": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { k.Group.Platform = service.PlatformAnthropic }), nil
|
||||
},
|
||||
"billing type changed": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { k.Group.SubscriptionType = service.SubscriptionTypeSubscription }), nil
|
||||
},
|
||||
}
|
||||
for name, lookup := range keep {
|
||||
t.Run(name+" keeps the connection snapshot", func(t *testing.T) {
|
||||
require.Same(t, conn, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, conn))
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("no credential or no group skips the lookup", func(t *testing.T) {
|
||||
called := false
|
||||
lookup := wsTurnAPIKeyLookupFunc(func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
called = true
|
||||
return latestWith(nil), nil
|
||||
})
|
||||
noKey := *conn
|
||||
noKey.Key = ""
|
||||
require.Same(t, &noKey, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, &noKey))
|
||||
noGroup := *conn
|
||||
noGroup.Group = nil
|
||||
require.Same(t, &noGroup, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, &noGroup))
|
||||
require.Nil(t, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, nil))
|
||||
require.Same(t, conn, refreshOpenAIWSTurnBillingAPIKey(context.Background(), nil, conn))
|
||||
require.False(t, called)
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIWSTurnBillingAPIKeysKeepPreviousTurnUntilItIsRecorded(t *testing.T) {
|
||||
conn := &service.APIKey{ID: 1}
|
||||
turn2 := &service.APIKey{ID: 2}
|
||||
turn3 := &service.APIKey{ID: 3}
|
||||
var keys openAIWSTurnBillingAPIKeys
|
||||
require.Same(t, conn, keys.forTurn(1, conn), "a turn without BeforeTurn bills with the connection snapshot")
|
||||
keys.set(2, turn2)
|
||||
keys.set(3, turn3) // turn 3 starts before turn 2's usage is submitted
|
||||
require.Same(t, turn2, keys.forTurn(2, conn))
|
||||
require.Same(t, turn3, keys.forTurn(3, conn))
|
||||
keys.set(4, conn)
|
||||
require.Same(t, conn, keys.forTurn(2, conn), "only the current and previous turn are kept")
|
||||
require.Len(t, keys.keys, 2)
|
||||
}
|
||||
@@ -264,10 +264,11 @@ func loginTestUser(t *testing.T) string {
|
||||
t.Helper()
|
||||
|
||||
// 先尝试用管理员账户登录
|
||||
adminEmail := getEnv("ADMIN_EMAIL", "admin@sub2api.local")
|
||||
// 管理员邮箱不再有固定默认值(未设置时安装阶段随机生成),需同时提供邮箱与密码。
|
||||
adminEmail := getEnv("ADMIN_EMAIL", "")
|
||||
adminPassword := getEnv("ADMIN_PASSWORD", "")
|
||||
|
||||
if adminPassword == "" {
|
||||
if adminEmail == "" || adminPassword == "" {
|
||||
// 尝试用测试用户
|
||||
adminEmail = testUserEmail
|
||||
adminPassword = testUserPassword
|
||||
|
||||
@@ -49,21 +49,9 @@ const (
|
||||
PlatformTypeSafe = domain.PlatformTypeSafe
|
||||
)
|
||||
|
||||
// AllPlatforms 返回所有支持的平台列表
|
||||
// AllPlatforms 返回所有支持的平台列表(平台清单,按展示顺序)。
|
||||
func AllPlatforms() []string {
|
||||
return []string{
|
||||
PlatformAnthropic,
|
||||
PlatformOpenAI,
|
||||
PlatformGemini,
|
||||
PlatformAntigravity,
|
||||
PlatformGrok,
|
||||
PlatformKimi,
|
||||
PlatformZhipu,
|
||||
PlatformDeepseek,
|
||||
PlatformMiniMax,
|
||||
PlatformOpenCodeGo,
|
||||
PlatformTypeSafe,
|
||||
}
|
||||
return domain.ConcretePlatformIDs()
|
||||
}
|
||||
|
||||
// Validate 验证规则配置的有效性
|
||||
|
||||
@@ -19,5 +19,7 @@ func TestAllPlatformsIncludesEveryConcretePlatform(t *testing.T) {
|
||||
"minimax",
|
||||
"opencode_go",
|
||||
"typesafe",
|
||||
"command_code",
|
||||
"cline",
|
||||
}, AllPlatforms())
|
||||
}
|
||||
|
||||
@@ -376,6 +376,9 @@ func (e *EasyPay) VerifyNotification(_ context.Context, rawBody string, _ map[st
|
||||
// url.ParseQuery already decodes values — no additional decode needed.
|
||||
params := make(map[string]string)
|
||||
for k := range values {
|
||||
if !easyPayNotifyAllowedParams[k] {
|
||||
return nil, fmt.Errorf("unexpected notify param: %s", k)
|
||||
}
|
||||
params[k] = values.Get(k)
|
||||
}
|
||||
sign := params["sign"]
|
||||
@@ -601,3 +604,26 @@ func easyPaySign(params map[string]string, pkey string) string {
|
||||
func easyPayVerifySign(params map[string]string, pkey string, sign string) bool {
|
||||
return hmac.Equal([]byte(easyPaySign(params, pkey)), []byte(sign))
|
||||
}
|
||||
|
||||
// easyPayNotifyAllowedParams is the exact parameter set a genuine EasyPay
|
||||
// (彩虹易支付-compatible) async notification carries. Order-creation-only
|
||||
// fields (notify_url, return_url, cid, device, clientip, ...) must never
|
||||
// appear in a callback: because the sign base string concatenates values
|
||||
// unescaped, a signed order URL whose return_url embeds e.g.
|
||||
// "trade_status=TRADE_SUCCESS" could otherwise be replayed as a forged
|
||||
// payment-success notification (issue #7881). Rejecting unknown keys closes
|
||||
// the whole smuggling class; genuinely paid orders rejected by an exotic
|
||||
// upstream variant are still recovered by the upstream QueryOrder reconcile
|
||||
// path.
|
||||
var easyPayNotifyAllowedParams = map[string]bool{
|
||||
"pid": true,
|
||||
"trade_no": true,
|
||||
"out_trade_no": true,
|
||||
"type": true,
|
||||
"name": true,
|
||||
"money": true,
|
||||
"trade_status": true,
|
||||
"param": true,
|
||||
"sign": true,
|
||||
"sign_type": true,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
package provider
|
||||
|
||||
// Regression tests for the EasyPay forged-callback vulnerability
|
||||
// (Wei-Shaw/sub2api issue #7875).
|
||||
//
|
||||
// Attack recap: the order-creation signature (visible to the payer in the
|
||||
// submit.php popup URL) signs return_url among other fields, and the sign
|
||||
// base string concatenates values unescaped. A client-supplied return_url
|
||||
// ending in "&trade_status=TRADE_SUCCESS" therefore produced a signature
|
||||
// that was byte-identical to a payment-success notification's signature
|
||||
// once the callback decoded return_url partially and promoted the smuggled
|
||||
// pair to a top-level parameter. Merchant-key rotation could not help
|
||||
// because the attacker never needed the key.
|
||||
//
|
||||
// Fixes covered here:
|
||||
// 1. service.CanonicalizeReturnURL drops client-supplied query parameters
|
||||
// (see payment_resume_service_test.go).
|
||||
// 2. VerifyNotification rejects any parameter outside the genuine async
|
||||
// notify set, so smuggled order-creation fields (return_url et al.)
|
||||
// fail closed even if a signed value smuggles a fake pair.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/url"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/payment"
|
||||
)
|
||||
|
||||
// easyPayPoCProvider returns a provider with a fixed test credential set.
|
||||
func easyPayPoCProvider() *EasyPay {
|
||||
return &EasyPay{config: map[string]string{
|
||||
"pid": "1000",
|
||||
"pkey": "MERCHANT_SECRET_KEY",
|
||||
"apiBase": "https://pay.example.com",
|
||||
"notifyUrl": "https://site.example.com/api/v1/payment/webhook/easypay",
|
||||
"returnUrl": "https://site.example.com/payment/result",
|
||||
}}
|
||||
}
|
||||
|
||||
// TestEasyPayNotifyRejectsForgedSignReuseCallback is the exact PoC payload
|
||||
// that was accepted before the fix: the order's own submit.php signature
|
||||
// replayed with trade_status smuggled out of the return_url value.
|
||||
func TestEasyPayNotifyRejectsForgedSignReuseCallback(t *testing.T) {
|
||||
t.Parallel()
|
||||
e := easyPayPoCProvider()
|
||||
|
||||
// Order creation (popup mode): return_url carries the smuggled pair as
|
||||
// its last sorted inner query key, exactly as buildPaymentReturnURL
|
||||
// would have produced before CanonicalizeReturnURL stripped user query.
|
||||
returnURL := "https://site.example.com/payment/result?order_id=99&out_trade_no=ORDER123&status=success&trade_status=TRADE_SUCCESS"
|
||||
createParams := map[string]string{
|
||||
"pid": "1000",
|
||||
"type": "alipay",
|
||||
"out_trade_no": "ORDER123",
|
||||
"notify_url": e.config["notifyUrl"],
|
||||
"return_url": returnURL,
|
||||
"name": "balance recharge",
|
||||
"money": "650.00",
|
||||
}
|
||||
sign := easyPaySign(createParams, e.config["pkey"])
|
||||
|
||||
// Forged callback: return_url encoded only up to status=success, then a
|
||||
// raw &trade_status=TRADE_SUCCESS promotes it to a top-level parameter.
|
||||
prefix := "https://site.example.com/payment/result?order_id=99&out_trade_no=ORDER123&status=success"
|
||||
cb := url.Values{}
|
||||
cb.Set("pid", "1000")
|
||||
cb.Set("type", "alipay")
|
||||
cb.Set("out_trade_no", "ORDER123")
|
||||
cb.Set("notify_url", e.config["notifyUrl"])
|
||||
cb.Set("name", "balance recharge")
|
||||
cb.Set("money", "650.00")
|
||||
cb.Set("return_url", prefix)
|
||||
rawCallback := cb.Encode() + "&trade_status=TRADE_SUCCESS" + "&sign=" + sign + "&sign_type=MD5"
|
||||
|
||||
if _, err := e.VerifyNotification(context.Background(), rawCallback, nil); err == nil {
|
||||
t.Fatal("forged sign-reuse callback must be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestEasyPayNotifyRejectsOrderURLReplay covers the milder variant where the
|
||||
// attacker replays the complete signed pay URL unchanged: return_url itself
|
||||
// is not a legitimate notify parameter and must also be rejected.
|
||||
func TestEasyPayNotifyRejectsOrderURLReplay(t *testing.T) {
|
||||
t.Parallel()
|
||||
e := easyPayPoCProvider()
|
||||
|
||||
createParams := map[string]string{
|
||||
"pid": "1000",
|
||||
"type": "alipay",
|
||||
"out_trade_no": "ORDER123",
|
||||
"notify_url": e.config["notifyUrl"],
|
||||
"return_url": "https://site.example.com/payment/result",
|
||||
"name": "balance recharge",
|
||||
"money": "650.00",
|
||||
}
|
||||
createParams["sign"] = easyPaySign(createParams, e.config["pkey"])
|
||||
createParams["sign_type"] = signTypeMD5
|
||||
q := url.Values{}
|
||||
for k, v := range createParams {
|
||||
q.Set(k, v)
|
||||
}
|
||||
|
||||
if _, err := e.VerifyNotification(context.Background(), q.Encode(), nil); err == nil {
|
||||
t.Fatal("replayed order URL must be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestEasyPayNotifyAcceptsGenuineCallback locks in the legitimate notify
|
||||
// contract: the canonical parameter set with a valid signature succeeds.
|
||||
func TestEasyPayNotifyAcceptsGenuineCallback(t *testing.T) {
|
||||
t.Parallel()
|
||||
e := easyPayPoCProvider()
|
||||
|
||||
params := map[string]string{
|
||||
"pid": "1000",
|
||||
"trade_no": "2026100622001400000001",
|
||||
"out_trade_no": "ORDER123",
|
||||
"type": "alipay",
|
||||
"name": "balance recharge",
|
||||
"money": "650.00",
|
||||
"trade_status": tradeStatusSuccess,
|
||||
}
|
||||
sign := easyPaySign(params, e.config["pkey"])
|
||||
params["sign"] = sign
|
||||
params["sign_type"] = signTypeMD5
|
||||
q := url.Values{}
|
||||
for k, v := range params {
|
||||
q.Set(k, v)
|
||||
}
|
||||
|
||||
n, err := e.VerifyNotification(context.Background(), q.Encode(), nil)
|
||||
if err != nil {
|
||||
t.Fatalf("genuine callback rejected: %v", err)
|
||||
}
|
||||
if n.Status != payment.ProviderStatusSuccess {
|
||||
t.Fatalf("status = %v, want success", n.Status)
|
||||
}
|
||||
if n.OrderID != "ORDER123" || n.TradeNo != "2026100622001400000001" || n.Amount != 650.00 {
|
||||
t.Fatalf("unexpected notification: %+v", n)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEasyPayNotifyRejectsUnknownParam ensures any parameter outside the
|
||||
// canonical notify set fails closed, including empty-valued ones that the
|
||||
// signer itself would skip.
|
||||
func TestEasyPayNotifyRejectsUnknownParam(t *testing.T) {
|
||||
t.Parallel()
|
||||
e := easyPayPoCProvider()
|
||||
|
||||
params := map[string]string{
|
||||
"pid": "1000",
|
||||
"trade_no": "T1",
|
||||
"out_trade_no": "ORDER123",
|
||||
"type": "alipay",
|
||||
"name": "balance recharge",
|
||||
"money": "650.00",
|
||||
"trade_status": tradeStatusSuccess,
|
||||
}
|
||||
sign := easyPaySign(params, e.config["pkey"])
|
||||
q := url.Values{}
|
||||
for k, v := range params {
|
||||
q.Set(k, v)
|
||||
}
|
||||
q.Set("sign", sign)
|
||||
q.Set("sign_type", signTypeMD5)
|
||||
q.Set("device", "") // empty value: invisible to the signer, still rejected
|
||||
|
||||
if _, err := e.VerifyNotification(context.Background(), q.Encode(), nil); err == nil {
|
||||
t.Fatal("callback with unknown param must be rejected")
|
||||
}
|
||||
}
|
||||
@@ -122,6 +122,20 @@ type ClaudeContentItem struct {
|
||||
Input any `json:"input,omitempty"`
|
||||
}
|
||||
|
||||
// MarshalJSON 保证 thinking 块始终带 thinking/signature 字段(即使为空),
|
||||
// 与 Anthropic 官方一致;Grok Build 等严格客户端缺少 signature 会解析失败。
|
||||
func (c ClaudeContentItem) MarshalJSON() ([]byte, error) {
|
||||
type claudeContentItem ClaudeContentItem
|
||||
if c.Type != "thinking" {
|
||||
return json.Marshal(claudeContentItem(c))
|
||||
}
|
||||
return json.Marshal(struct {
|
||||
Thinking string `json:"thinking"`
|
||||
Signature string `json:"signature"`
|
||||
claudeContentItem
|
||||
}{Thinking: c.Thinking, Signature: c.Signature, claudeContentItem: claudeContentItem(c)})
|
||||
}
|
||||
|
||||
// ClaudeUsage Claude 用量统计
|
||||
type ClaudeUsage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
package antigravity
|
||||
|
||||
import "testing"
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDefaultModels_ContainsNewAndLegacyImageModels(t *testing.T) {
|
||||
t.Parallel()
|
||||
@@ -44,3 +47,23 @@ func TestDefaultModels_ContainsNewAndLegacyImageModels(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaudeContentItem_ThinkingAlwaysHasSignature(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
data, err := json.Marshal(ClaudeContentItem{Type: "thinking", Thinking: "plan"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := string(data), `{"thinking":"plan","signature":"","type":"thinking"}`; got != want {
|
||||
t.Fatalf("thinking block = %s, want %s", got, want)
|
||||
}
|
||||
|
||||
data, err = json.Marshal(ClaudeContentItem{Type: "text", Text: "hi"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := string(data), `{"type":"text","text":"hi"}`; got != want {
|
||||
t.Fatalf("text block = %s, want %s", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -296,8 +296,9 @@ func (p *StreamingProcessor) processThinking(text, signature string) []byte {
|
||||
// 开始或继续 thinking 块
|
||||
if p.blockType != BlockTypeThinking {
|
||||
_, _ = result.Write(p.startBlock(BlockTypeThinking, map[string]any{
|
||||
"type": "thinking",
|
||||
"thinking": "",
|
||||
"type": "thinking",
|
||||
"thinking": "",
|
||||
"signature": "",
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -475,8 +476,9 @@ func (p *StreamingProcessor) emitEmptyThinkingWithSignature(signature string) []
|
||||
var result bytes.Buffer
|
||||
|
||||
_, _ = result.Write(p.startBlock(BlockTypeThinking, map[string]any{
|
||||
"type": "thinking",
|
||||
"thinking": "",
|
||||
"type": "thinking",
|
||||
"thinking": "",
|
||||
"signature": "",
|
||||
}))
|
||||
_, _ = result.Write(p.emitDelta("thinking_delta", map[string]any{
|
||||
"thinking": "",
|
||||
|
||||
@@ -844,7 +844,7 @@ func TestStreamingReasoning(t *testing.T) {
|
||||
|
||||
sse, err := ResponsesAnthropicEventToSSE(events[0])
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, sse, `"content_block":{"thinking":"","type":"thinking"}`)
|
||||
assert.Contains(t, sse, `"content_block":{"thinking":"","signature":"","type":"thinking"}`)
|
||||
|
||||
// reasoning text delta
|
||||
events = ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
|
||||
|
||||
@@ -88,10 +88,11 @@ func ChatCompletionsToResponses(req *ChatCompletionsRequest) (*ResponsesRequest,
|
||||
out.Tools = convertChatToolsToResponses(req.Tools, req.Functions)
|
||||
}
|
||||
|
||||
// tool_choice: already compatible format — pass through directly.
|
||||
// tool_choice: strings and Responses-shaped objects pass through; a named
|
||||
// Chat choice nests the name under "function" and must be flattened.
|
||||
// Legacy function_call needs mapping.
|
||||
if len(req.ToolChoice) > 0 {
|
||||
out.ToolChoice = req.ToolChoice
|
||||
out.ToolChoice = convertChatToolChoiceToResponses(req.ToolChoice)
|
||||
} else if len(req.FunctionCall) > 0 {
|
||||
tc, err := convertChatFunctionCallToToolChoice(req.FunctionCall)
|
||||
if err != nil {
|
||||
@@ -107,7 +108,18 @@ func ChatCompletionsToResponses(req *ChatCompletionsRequest) (*ResponsesRequest,
|
||||
// array into a Responses API input items array.
|
||||
func convertChatMessagesToResponsesInput(msgs []ChatMessage) ([]ResponsesInputItem, error) {
|
||||
var out []ResponsesInputItem
|
||||
legacyIDs := legacyFunctionCallIDs{msgs: msgs}
|
||||
for _, m := range msgs {
|
||||
switch {
|
||||
case m.Role == "assistant" && m.FunctionCall != nil && len(m.ToolCalls) == 0:
|
||||
m.ToolCalls = []ChatToolCall{{
|
||||
ID: legacyIDs.assign(m.FunctionCall.Name),
|
||||
Type: "function",
|
||||
Function: *m.FunctionCall,
|
||||
}}
|
||||
case m.Role == "function" && m.ToolCallID == "":
|
||||
m.ToolCallID = legacyIDs.claim(m.Name)
|
||||
}
|
||||
items, err := chatMessageToResponsesItems(m)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -117,11 +129,63 @@ func convertChatMessagesToResponsesInput(msgs []ChatMessage) ([]ResponsesInputIt
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// legacyFunctionCallIDs pairs legacy function calls with their results. Legacy
|
||||
// function calling carries no call IDs: the assistant turn has function_call
|
||||
// and the role=function result refers to it only by name, while the Responses
|
||||
// API pairs function_call and function_call_output by call_id. IDs follow
|
||||
// message order, so a replayed history converts to the same input, and skip
|
||||
// IDs already used by tool_calls in the same conversation.
|
||||
type legacyFunctionCallIDs struct {
|
||||
msgs []ChatMessage
|
||||
used map[string]bool
|
||||
next int
|
||||
pending map[string][]string
|
||||
}
|
||||
|
||||
// assign returns a new call ID for a legacy function call and queues it for
|
||||
// the next result with the same function name.
|
||||
func (ids *legacyFunctionCallIDs) assign(name string) string {
|
||||
if ids.used == nil {
|
||||
ids.used = make(map[string]bool)
|
||||
ids.pending = make(map[string][]string)
|
||||
for _, m := range ids.msgs {
|
||||
for _, tc := range m.ToolCalls {
|
||||
ids.used[tc.ID] = true
|
||||
}
|
||||
if m.ToolCallID != "" {
|
||||
ids.used[m.ToolCallID] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
var id string
|
||||
for {
|
||||
ids.next++
|
||||
id = fmt.Sprintf("call_legacy_%d", ids.next)
|
||||
if !ids.used[id] {
|
||||
break
|
||||
}
|
||||
}
|
||||
ids.used[id] = true
|
||||
ids.pending[name] = append(ids.pending[name], id)
|
||||
return id
|
||||
}
|
||||
|
||||
// claim returns the oldest unanswered call ID for name, or "" when no legacy
|
||||
// call with that name precedes the result.
|
||||
func (ids *legacyFunctionCallIDs) claim(name string) string {
|
||||
queue := ids.pending[name]
|
||||
if len(queue) == 0 {
|
||||
return ""
|
||||
}
|
||||
ids.pending[name] = queue[1:]
|
||||
return queue[0]
|
||||
}
|
||||
|
||||
// chatMessageToResponsesItems converts a single ChatMessage into one or more
|
||||
// ResponsesInputItem values.
|
||||
func chatMessageToResponsesItems(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
switch m.Role {
|
||||
case "system":
|
||||
case "system", "developer":
|
||||
return chatSystemToResponses(m)
|
||||
case "user":
|
||||
return chatUserToResponses(m)
|
||||
@@ -136,7 +200,9 @@ func chatMessageToResponsesItems(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
}
|
||||
}
|
||||
|
||||
// chatSystemToResponses converts a system message.
|
||||
// chatSystemToResponses converts a system or developer message, keeping its
|
||||
// role: the Responses API accepts both, and developer must not be demoted to
|
||||
// a user turn.
|
||||
func chatSystemToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
parsed, err := parseChatMessageContent(m.Content)
|
||||
if err != nil {
|
||||
@@ -146,7 +212,7 @@ func chatSystemToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []ResponsesInputItem{{Type: "message", Role: "system", Content: content}}, nil
|
||||
return []ResponsesInputItem{{Type: "message", Role: m.Role, Content: content}}, nil
|
||||
}
|
||||
|
||||
// chatUserToResponses converts a user message, handling both plain strings and
|
||||
@@ -303,8 +369,9 @@ func chatToolToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
}
|
||||
|
||||
// chatFunctionToResponses converts a legacy function result message
|
||||
// (role=function) into a function_call_output item. The Name field is used as
|
||||
// call_id since legacy function calls do not carry a separate call_id.
|
||||
// (role=function) into a function_call_output item. ToolCallID holds the ID
|
||||
// assigned to the matching legacy call (see legacyFunctionCallIDs); without a
|
||||
// matching call the Name field is used as call_id.
|
||||
func chatFunctionToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
output, err := parseChatContent(m.Content)
|
||||
if err != nil {
|
||||
@@ -313,9 +380,13 @@ func chatFunctionToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
if output == "" {
|
||||
output = "(empty)"
|
||||
}
|
||||
callID := m.ToolCallID
|
||||
if callID == "" {
|
||||
callID = m.Name
|
||||
}
|
||||
return []ResponsesInputItem{{
|
||||
Type: "function_call_output",
|
||||
CallID: m.Name,
|
||||
CallID: callID,
|
||||
Output: output,
|
||||
}}, nil
|
||||
}
|
||||
@@ -490,6 +561,34 @@ func defaultStrictFalse(src *bool) *bool {
|
||||
return src
|
||||
}
|
||||
|
||||
// convertChatToolChoiceToResponses maps a Chat Completions tool_choice to the
|
||||
// Responses API shape.
|
||||
//
|
||||
// {"type":"function","function":{"name":"X"}} → {"type":"function","name":"X"}
|
||||
//
|
||||
// Strings ("auto", "none", "required") and objects that already use the
|
||||
// Responses shape are returned unchanged.
|
||||
func convertChatToolChoiceToResponses(raw json.RawMessage) json.RawMessage {
|
||||
var choice struct {
|
||||
Type string `json:"type"`
|
||||
Function *struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"function"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &choice); err != nil || choice.Type != "function" || choice.Function == nil {
|
||||
return raw
|
||||
}
|
||||
name := strings.TrimSpace(choice.Function.Name)
|
||||
if name == "" {
|
||||
return raw
|
||||
}
|
||||
flat, err := json.Marshal(map[string]string{"type": "function", "name": name})
|
||||
if err != nil {
|
||||
return raw
|
||||
}
|
||||
return flat
|
||||
}
|
||||
|
||||
// convertChatFunctionCallToToolChoice maps the legacy function_call field to a
|
||||
// Responses API tool_choice value.
|
||||
//
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// The Responses API accepts role "developer" in input messages. Converting it
|
||||
// to "user" silently demotes the caller's instructions to user turns.
|
||||
func TestChatCompletionsToResponses_DeveloperRolePreserved(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
wantContent string
|
||||
}{
|
||||
{name: "string content", content: `"Answer in French."`, wantContent: `"Answer in French."`},
|
||||
{name: "text parts", content: `[{"type":"text","text":"Answer in French."}]`, wantContent: `[{"type":"input_text","text":"Answer in French."}]`},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req := &ChatCompletionsRequest{
|
||||
Model: "gpt-4o",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "developer", Content: json.RawMessage(tt.content)},
|
||||
{Role: "user", Content: json.RawMessage(`"Hello"`)},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := ChatCompletionsToResponses(req)
|
||||
require.NoError(t, err)
|
||||
|
||||
var items []ResponsesInputItem
|
||||
require.NoError(t, json.Unmarshal(resp.Input, &items))
|
||||
require.Len(t, items, 2)
|
||||
require.Equal(t, "message", items[0].Type)
|
||||
require.Equal(t, "developer", items[0].Role)
|
||||
require.JSONEq(t, tt.wantContent, string(items[0].Content))
|
||||
require.Equal(t, "user", items[1].Role)
|
||||
})
|
||||
}
|
||||
}
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func convertLegacyFunctionHistory(t *testing.T, msgs []ChatMessage) []ResponsesInputItem {
|
||||
t.Helper()
|
||||
resp, err := ChatCompletionsToResponses(&ChatCompletionsRequest{Model: "gpt-4o", Messages: msgs})
|
||||
require.NoError(t, err)
|
||||
var items []ResponsesInputItem
|
||||
require.NoError(t, json.Unmarshal(resp.Input, &items))
|
||||
return items
|
||||
}
|
||||
|
||||
// Legacy function calling has no call IDs: the assistant turn carries
|
||||
// function_call and the result is a role=function message keyed by name. The
|
||||
// Responses API needs a function_call item and a function_call_output with the
|
||||
// same call_id, otherwise the output is rejected as having no matching call.
|
||||
func TestChatCompletionsToResponses_LegacyFunctionCallHistoryPaired(t *testing.T) {
|
||||
msgs := []ChatMessage{
|
||||
{Role: "user", Content: json.RawMessage(`"Weather in Paris and Rome?"`)},
|
||||
{Role: "assistant", Content: json.RawMessage(`"Checking Paris."`), FunctionCall: &ChatFunctionCall{Name: "get_weather", Arguments: `{"city":"Paris"}`}},
|
||||
{Role: "function", Name: "get_weather", Content: json.RawMessage(`"sunny"`)},
|
||||
{Role: "assistant", FunctionCall: &ChatFunctionCall{Name: "get_weather", Arguments: `{"city":"Rome"}`}},
|
||||
{Role: "function", Name: "get_weather", Content: json.RawMessage(`"rainy"`)},
|
||||
}
|
||||
|
||||
items := convertLegacyFunctionHistory(t, msgs)
|
||||
require.Len(t, items, 6)
|
||||
|
||||
require.Equal(t, "user", items[0].Role)
|
||||
require.Equal(t, "assistant", items[1].Role)
|
||||
|
||||
firstCall, firstOutput := items[2], items[3]
|
||||
require.Equal(t, "function_call", firstCall.Type)
|
||||
require.Equal(t, "get_weather", firstCall.Name)
|
||||
require.Equal(t, `{"city":"Paris"}`, firstCall.Arguments)
|
||||
require.Equal(t, "function_call_output", firstOutput.Type)
|
||||
require.Equal(t, "sunny", firstOutput.Output)
|
||||
require.NotEmpty(t, firstCall.CallID)
|
||||
require.Equal(t, firstCall.CallID, firstOutput.CallID)
|
||||
|
||||
secondCall, secondOutput := items[4], items[5]
|
||||
require.Equal(t, "function_call", secondCall.Type)
|
||||
require.Equal(t, `{"city":"Rome"}`, secondCall.Arguments)
|
||||
require.Equal(t, "function_call_output", secondOutput.Type)
|
||||
require.Equal(t, "rainy", secondOutput.Output)
|
||||
require.Equal(t, secondCall.CallID, secondOutput.CallID)
|
||||
require.NotEqual(t, firstCall.CallID, secondCall.CallID)
|
||||
|
||||
// Synthesized IDs must be stable across turns so the replayed history keeps
|
||||
// the same prompt prefix.
|
||||
require.Equal(t, items, convertLegacyFunctionHistory(t, msgs))
|
||||
}
|
||||
|
||||
func TestChatCompletionsToResponses_LegacyFunctionCallIDAvoidsToolCallIDs(t *testing.T) {
|
||||
items := convertLegacyFunctionHistory(t, []ChatMessage{
|
||||
{Role: "user", Content: json.RawMessage(`"hi"`)},
|
||||
{Role: "assistant", FunctionCall: &ChatFunctionCall{Name: "lookup", Arguments: `{}`}},
|
||||
{Role: "function", Name: "lookup", Content: json.RawMessage(`"legacy result"`)},
|
||||
{Role: "assistant", ToolCalls: []ChatToolCall{{ID: "call_legacy_1", Type: "function", Function: ChatFunctionCall{Name: "lookup", Arguments: `{}`}}}},
|
||||
{Role: "tool", ToolCallID: "call_legacy_1", Content: json.RawMessage(`"tool result"`)},
|
||||
})
|
||||
require.Len(t, items, 5)
|
||||
|
||||
legacyCall, legacyOutput := items[1], items[2]
|
||||
require.Equal(t, "function_call", legacyCall.Type)
|
||||
require.Equal(t, legacyCall.CallID, legacyOutput.CallID)
|
||||
require.NotEqual(t, "call_legacy_1", legacyCall.CallID)
|
||||
|
||||
require.Equal(t, "call_legacy_1", items[3].CallID)
|
||||
require.Equal(t, "call_legacy_1", items[4].CallID)
|
||||
}
|
||||
|
||||
func TestChatCompletionsToResponses_LegacyFunctionResultWithoutCallKeepsName(t *testing.T) {
|
||||
items := convertLegacyFunctionHistory(t, []ChatMessage{
|
||||
{Role: "user", Content: json.RawMessage(`"hi"`)},
|
||||
{Role: "function", Name: "lookup", Content: json.RawMessage(`"orphan"`)},
|
||||
})
|
||||
require.Len(t, items, 2)
|
||||
require.Equal(t, "function_call_output", items[1].Type)
|
||||
require.Equal(t, "lookup", items[1].CallID)
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// Chat Completions nests a named choice under "function"; the Responses API
|
||||
// expects the name at the top level and rejects the nested shape.
|
||||
func TestChatCompletionsToResponses_ToolChoiceShape(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
choice string
|
||||
want string
|
||||
}{
|
||||
{name: "named chat function is flattened", choice: `{"type":"function","function":{"name":"get_weather"}}`, want: `{"type":"function","name":"get_weather"}`},
|
||||
{name: "flat responses function is kept", choice: `{"type":"function","name":"get_weather"}`, want: `{"type":"function","name":"get_weather"}`},
|
||||
{name: "auto string is kept", choice: `"auto"`, want: `"auto"`},
|
||||
{name: "required string is kept", choice: `"required"`, want: `"required"`},
|
||||
{name: "none string is kept", choice: `"none"`, want: `"none"`},
|
||||
{name: "allowed_tools policy is kept", choice: `{"type":"allowed_tools","mode":"auto","tools":[{"type":"function","name":"get_weather"}]}`, want: `{"type":"allowed_tools","mode":"auto","tools":[{"type":"function","name":"get_weather"}]}`},
|
||||
{name: "function without name is kept", choice: `{"type":"function","function":{}}`, want: `{"type":"function","function":{}}`},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req := &ChatCompletionsRequest{
|
||||
Model: "gpt-4o",
|
||||
Messages: []ChatMessage{{Role: "user", Content: json.RawMessage(`"weather?"`)}},
|
||||
Tools: []ChatTool{{
|
||||
Type: "function",
|
||||
Function: &ChatFunction{Name: "get_weather", Parameters: json.RawMessage(`{"type":"object"}`)},
|
||||
}},
|
||||
ToolChoice: json.RawMessage(tt.choice),
|
||||
}
|
||||
|
||||
resp, err := ChatCompletionsToResponses(req)
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, tt.want, string(resp.ToolChoice))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -15,7 +15,8 @@ import (
|
||||
|
||||
// ResponsesToChatCompletions converts a Responses API response into a Chat
|
||||
// Completions response. Text output items are concatenated into
|
||||
// choices[0].message.content; function_call items become tool_calls.
|
||||
// choices[0].message.content, refusal parts into choices[0].message.refusal;
|
||||
// function_call items become tool_calls.
|
||||
func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatCompletionsResponse {
|
||||
id := resp.ID
|
||||
if id == "" {
|
||||
@@ -31,6 +32,7 @@ func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatComp
|
||||
}
|
||||
|
||||
var contentText string
|
||||
var refusalText string
|
||||
var reasoningText string
|
||||
var toolCalls []ChatToolCall
|
||||
|
||||
@@ -38,8 +40,11 @@ func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatComp
|
||||
switch item.Type {
|
||||
case "message":
|
||||
for _, part := range item.Content {
|
||||
if part.Type == "output_text" && part.Text != "" {
|
||||
switch {
|
||||
case part.Type == "output_text" && part.Text != "":
|
||||
contentText += part.Text
|
||||
case part.Type == "refusal" && part.Refusal != "":
|
||||
refusalText += part.Refusal
|
||||
}
|
||||
}
|
||||
case "function_call":
|
||||
@@ -62,7 +67,7 @@ func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatComp
|
||||
}
|
||||
}
|
||||
|
||||
msg := ChatMessage{Role: "assistant"}
|
||||
msg := ChatMessage{Role: "assistant", Refusal: refusalText}
|
||||
if len(toolCalls) > 0 {
|
||||
msg.ToolCalls = toolCalls
|
||||
}
|
||||
@@ -149,6 +154,8 @@ func ResponsesEventToChatChunks(evt *ResponsesStreamEvent, state *ResponsesEvent
|
||||
return resToChatHandleCreated(evt, state)
|
||||
case "response.output_text.delta":
|
||||
return resToChatHandleTextDelta(evt, state)
|
||||
case "response.refusal.delta":
|
||||
return resToChatHandleRefusalDelta(evt, state)
|
||||
case "response.output_item.added":
|
||||
return resToChatHandleOutputItemAdded(evt, state)
|
||||
case "response.function_call_arguments.delta",
|
||||
@@ -248,6 +255,14 @@ func resToChatHandleTextDelta(evt *ResponsesStreamEvent, state *ResponsesEventTo
|
||||
return []ChatCompletionsChunk{makeChatDeltaChunk(state, ChatDelta{Content: &content})}
|
||||
}
|
||||
|
||||
func resToChatHandleRefusalDelta(evt *ResponsesStreamEvent, state *ResponsesEventToChatState) []ChatCompletionsChunk {
|
||||
if evt.Delta == "" {
|
||||
return nil
|
||||
}
|
||||
refusal := evt.Delta
|
||||
return []ChatCompletionsChunk{makeChatDeltaChunk(state, ChatDelta{Refusal: &refusal})}
|
||||
}
|
||||
|
||||
func resToChatHandleOutputItemAdded(evt *ResponsesStreamEvent, state *ResponsesEventToChatState) []ChatCompletionsChunk {
|
||||
// function_call 与 custom_tool_call(custom/freeform 工具)均按工具调用注册,
|
||||
// 以便后续 *_input.delta / *_arguments.delta 能映射到正确的工具索引。
|
||||
@@ -496,6 +511,7 @@ type bufferedFuncCall struct {
|
||||
// (response.completed / response.done) carries an empty output array.
|
||||
type BufferedResponseAccumulator struct {
|
||||
text strings.Builder
|
||||
refusal strings.Builder
|
||||
reasoning strings.Builder
|
||||
funcCalls []bufferedFuncCall
|
||||
outputIndexToFuncIdx map[int]int
|
||||
@@ -517,6 +533,10 @@ func (a *BufferedResponseAccumulator) ProcessEvent(event *ResponsesStreamEvent)
|
||||
if event.Delta != "" {
|
||||
_, _ = a.text.WriteString(event.Delta)
|
||||
}
|
||||
case "response.refusal.delta":
|
||||
if event.Delta != "" {
|
||||
_, _ = a.refusal.WriteString(event.Delta)
|
||||
}
|
||||
case "response.output_item.added":
|
||||
if event.Item != nil && (event.Item.Type == "function_call" || event.Item.Type == "custom_tool_call") {
|
||||
idx := len(a.funcCalls)
|
||||
@@ -553,7 +573,7 @@ func (a *BufferedResponseAccumulator) ProcessEvent(event *ResponsesStreamEvent)
|
||||
|
||||
// HasContent reports whether any content has been accumulated.
|
||||
func (a *BufferedResponseAccumulator) HasContent() bool {
|
||||
return a.text.Len() > 0 || len(a.funcCalls) > 0 || a.reasoning.Len() > 0
|
||||
return a.text.Len() > 0 || a.refusal.Len() > 0 || len(a.funcCalls) > 0 || a.reasoning.Len() > 0
|
||||
}
|
||||
|
||||
// BuildOutput constructs a []ResponsesOutput from the accumulated delta
|
||||
@@ -572,14 +592,18 @@ func (a *BufferedResponseAccumulator) BuildOutput() []ResponsesOutput {
|
||||
})
|
||||
}
|
||||
|
||||
if a.text.Len() > 0 {
|
||||
if a.text.Len() > 0 || a.refusal.Len() > 0 {
|
||||
var content []ResponsesContentPart
|
||||
if a.text.Len() > 0 {
|
||||
content = append(content, ResponsesContentPart{Type: "output_text", Text: a.text.String()})
|
||||
}
|
||||
if a.refusal.Len() > 0 {
|
||||
content = append(content, ResponsesContentPart{Type: "refusal", Refusal: a.refusal.String()})
|
||||
}
|
||||
out = append(out, ResponsesOutput{
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []ResponsesContentPart{{
|
||||
Type: "output_text",
|
||||
Text: a.text.String(),
|
||||
}},
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: content,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const refusalText = "I can't help with that."
|
||||
|
||||
func decodeResponsesEvent(t *testing.T, raw string) ResponsesStreamEvent {
|
||||
t.Helper()
|
||||
var evt ResponsesStreamEvent
|
||||
require.NoError(t, json.Unmarshal([]byte(raw), &evt))
|
||||
return evt
|
||||
}
|
||||
|
||||
func chatMessageJSON(t *testing.T, resp *ChatCompletionsResponse) map[string]any {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(resp)
|
||||
require.NoError(t, err)
|
||||
var decoded struct {
|
||||
Choices []struct {
|
||||
Message map[string]any `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(raw, &decoded))
|
||||
require.Len(t, decoded.Choices, 1)
|
||||
return decoded.Choices[0].Message
|
||||
}
|
||||
|
||||
// A refusal arrives as a "refusal" content part, not output_text. Dropping it
|
||||
// leaves the Chat client with an empty assistant message.
|
||||
func TestResponsesToChatCompletions_RefusalPreserved(t *testing.T) {
|
||||
var resp ResponsesResponse
|
||||
require.NoError(t, json.Unmarshal([]byte(`{
|
||||
"id":"resp_refusal",
|
||||
"status":"completed",
|
||||
"output":[{"type":"message","role":"assistant","content":[{"type":"refusal","refusal":"`+refusalText+`"}]}]
|
||||
}`), &resp))
|
||||
|
||||
chat := ResponsesToChatCompletions(&resp, "gpt-4o")
|
||||
|
||||
msg := chatMessageJSON(t, chat)
|
||||
require.Equal(t, refusalText, msg["refusal"])
|
||||
require.NotContains(t, msg, "content")
|
||||
require.Equal(t, "stop", chat.Choices[0].FinishReason)
|
||||
}
|
||||
|
||||
func TestResponsesEventToChatChunks_RefusalDelta(t *testing.T) {
|
||||
state := NewResponsesEventToChatState()
|
||||
events := []string{
|
||||
`{"type":"response.created","response":{"id":"resp_refusal","model":"gpt-4o"}}`,
|
||||
`{"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant"}}`,
|
||||
`{"type":"response.refusal.delta","output_index":0,"content_index":0,"delta":"I can't "}`,
|
||||
`{"type":"response.refusal.delta","output_index":0,"content_index":0,"delta":"help with that."}`,
|
||||
`{"type":"response.refusal.done","output_index":0,"content_index":0,"refusal":"` + refusalText + `"}`,
|
||||
`{"type":"response.completed","response":{"id":"resp_refusal","status":"completed"}}`,
|
||||
}
|
||||
|
||||
var refusal strings.Builder
|
||||
var finishReason string
|
||||
for _, raw := range events {
|
||||
evt := decodeResponsesEvent(t, raw)
|
||||
for _, chunk := range ResponsesEventToChatChunks(&evt, state) {
|
||||
encoded, err := json.Marshal(chunk)
|
||||
require.NoError(t, err)
|
||||
var decoded struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Refusal *string `json:"refusal"`
|
||||
} `json:"delta"`
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(encoded, &decoded))
|
||||
for _, choice := range decoded.Choices {
|
||||
if choice.Delta.Refusal != nil {
|
||||
_, _ = refusal.WriteString(*choice.Delta.Refusal)
|
||||
}
|
||||
if choice.FinishReason != nil {
|
||||
finishReason = *choice.FinishReason
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
require.Equal(t, refusalText, refusal.String())
|
||||
require.Equal(t, "stop", finishReason)
|
||||
}
|
||||
|
||||
// Non-streaming Chat requests are served from a forced upstream stream; when
|
||||
// the terminal event carries an empty output array the message is rebuilt from
|
||||
// the accumulated deltas, which must include refusal deltas.
|
||||
func TestBufferedResponseAccumulator_RefusalRebuiltFromDeltas(t *testing.T) {
|
||||
acc := NewBufferedResponseAccumulator()
|
||||
for _, raw := range []string{
|
||||
`{"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant"}}`,
|
||||
`{"type":"response.refusal.delta","output_index":0,"content_index":0,"delta":"I can't "}`,
|
||||
`{"type":"response.refusal.delta","output_index":0,"content_index":0,"delta":"help with that."}`,
|
||||
} {
|
||||
evt := decodeResponsesEvent(t, raw)
|
||||
acc.ProcessEvent(&evt)
|
||||
}
|
||||
require.True(t, acc.HasContent())
|
||||
|
||||
final := &ResponsesResponse{ID: "resp_refusal", Status: "completed"}
|
||||
acc.SupplementResponseOutput(final)
|
||||
chat := ResponsesToChatCompletions(final, "gpt-4o")
|
||||
|
||||
require.Equal(t, refusalText, chatMessageJSON(t, chat)["refusal"])
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// Thinking blocks must always carry the `signature` key (Anthropic wire shape);
|
||||
// strict clients such as Grok Build fail with "missing field `signature`".
|
||||
func TestAnthropicContentBlock_ThinkingAlwaysHasSignature(t *testing.T) {
|
||||
data, err := json.Marshal(AnthropicContentBlock{Type: "thinking", Thinking: "plan"})
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{"type":"thinking","thinking":"plan","signature":""}`, string(data))
|
||||
|
||||
data, err = json.Marshal(AnthropicContentBlock{Type: "thinking", Signature: "sig"})
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{"type":"thinking","thinking":"","signature":"sig"}`, string(data))
|
||||
|
||||
data, err = json.Marshal(AnthropicContentBlock{Type: "text", Text: "hi"})
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{"type":"text","text":"hi"}`, string(data))
|
||||
}
|
||||
@@ -95,10 +95,14 @@ func (b AnthropicContentBlock) MarshalJSON() ([]byte, error) {
|
||||
anthropicContentBlock
|
||||
}{Text: b.Text, anthropicContentBlock: anthropicContentBlock(b)})
|
||||
case "thinking":
|
||||
// Anthropic always sends `signature` on thinking blocks (empty on
|
||||
// content_block_start); strict clients such as Grok Build reject the
|
||||
// block with "missing field `signature`" when the key is absent.
|
||||
return json.Marshal(struct {
|
||||
Thinking string `json:"thinking"`
|
||||
Thinking string `json:"thinking"`
|
||||
Signature string `json:"signature"`
|
||||
anthropicContentBlock
|
||||
}{Thinking: b.Thinking, anthropicContentBlock: anthropicContentBlock(b)})
|
||||
}{Thinking: b.Thinking, Signature: b.Signature, anthropicContentBlock: anthropicContentBlock(b)})
|
||||
default:
|
||||
return json.Marshal(base)
|
||||
}
|
||||
@@ -312,8 +316,9 @@ func (i *ResponsesInputItem) UnmarshalJSON(data []byte) error {
|
||||
// ResponsesContentPart is a typed content part in a Responses message.
|
||||
type ResponsesContentPart struct {
|
||||
PromptCacheBreakpoint json.RawMessage `json:"prompt_cache_breakpoint,omitempty"`
|
||||
Type string `json:"type"` // "input_text" | "output_text" | "input_image" | "input_file"
|
||||
Type string `json:"type"` // "input_text" | "output_text" | "refusal" | "input_image" | "input_file"
|
||||
Text string `json:"text,omitempty"`
|
||||
Refusal string `json:"refusal,omitempty"` // type=refusal
|
||||
ImageURL string `json:"image_url,omitempty"` // data URI for input_image
|
||||
|
||||
// input_file fields.
|
||||
@@ -690,6 +695,7 @@ type ChatStreamOptions struct {
|
||||
type ChatMessage struct {
|
||||
Role string `json:"role"` // "system" | "user" | "assistant" | "tool" | "function"
|
||||
Content json.RawMessage `json:"content,omitempty"`
|
||||
Refusal string `json:"refusal,omitempty"`
|
||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||
Reasoning string `json:"reasoning,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
@@ -830,6 +836,7 @@ type ChatChunkChoice struct {
|
||||
type ChatDelta struct {
|
||||
Role string `json:"role,omitempty"`
|
||||
Content *string `json:"content,omitempty"` // pointer: omit when not present, null vs "" matters
|
||||
Refusal *string `json:"refusal,omitempty"`
|
||||
ReasoningContent *string `json:"reasoning_content,omitempty"`
|
||||
Reasoning *string `json:"reasoning,omitempty"`
|
||||
ToolCalls []ChatToolCall `json:"tool_calls,omitempty"`
|
||||
|
||||
@@ -20,6 +20,9 @@ const (
|
||||
BetaFastMode = "fast-mode-2026-02-01"
|
||||
// Legacy structured output compatibility; forwarded only when explicitly requested.
|
||||
BetaStructuredOutputs = "structured-outputs-2025-11-13"
|
||||
// Tool-change compatibility tokens; never enabled without an explicit client request.
|
||||
BetaMidConversationToolChanges = "mid-conversation-tool-changes-2026-07-01"
|
||||
BetaInlineTools = "inline-tools-2026-09-15"
|
||||
|
||||
// 新增(对齐官方 CLI 2.1.9x 以来的流量)
|
||||
BetaPromptCachingScope = "prompt-caching-scope-2026-01-05"
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package xai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
@@ -14,6 +15,8 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/html"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -147,7 +150,7 @@ func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) {
|
||||
return nil, fmt.Errorf("open xAI device verification page: %w", SSOHTTPError{Status: status})
|
||||
}
|
||||
|
||||
status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOVerifyURL, url.Values{"user_code": {device.UserCode}})
|
||||
status, finalURL, body, err = f.do(ctx, http.MethodPost, SSOVerifyURL, url.Values{"user_code": {device.UserCode}})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -158,11 +161,19 @@ func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) {
|
||||
return nil, errors.New("xAI device verification did not reach consent page")
|
||||
}
|
||||
|
||||
status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOApproveURL, url.Values{
|
||||
approval := url.Values{
|
||||
"user_code": {device.UserCode},
|
||||
"action": {"allow"},
|
||||
"principal_type": {"User"},
|
||||
"principal_id": {""},
|
||||
}
|
||||
if consentToken := ssoConsentToken(body); consentToken != "" {
|
||||
approval.Set("consent_token", consentToken)
|
||||
}
|
||||
consentURL, _ := url.Parse(finalURL) // f.do only returns validated xAI URLs.
|
||||
status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOApproveURL, approval, http.Header{
|
||||
"Origin": {consentURL.Scheme + "://" + consentURL.Host},
|
||||
"Referer": {finalURL},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -177,6 +188,37 @@ func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) {
|
||||
return f.pollToken(ctx, device.DeviceCode, time.Duration(device.Interval)*time.Second, time.Duration(device.ExpiresIn)*time.Second)
|
||||
}
|
||||
|
||||
// xAI binds device approval to a hidden token on the consent page. Keep it in
|
||||
// this flow only; older consent pages without the field remain supported.
|
||||
func ssoConsentToken(body []byte) string {
|
||||
tokenizer := html.NewTokenizer(bytes.NewReader(body))
|
||||
for {
|
||||
switch tokenizer.Next() {
|
||||
case html.ErrorToken:
|
||||
return ""
|
||||
case html.StartTagToken, html.SelfClosingTagToken:
|
||||
token := tokenizer.Token()
|
||||
if token.Data != "input" {
|
||||
continue
|
||||
}
|
||||
var name, value, inputType string
|
||||
for _, attr := range token.Attr {
|
||||
switch attr.Key {
|
||||
case "name":
|
||||
name = attr.Val
|
||||
case "value":
|
||||
value = attr.Val
|
||||
case "type":
|
||||
inputType = attr.Val
|
||||
}
|
||||
}
|
||||
if name == "consent_token" && strings.EqualFold(inputType, "hidden") {
|
||||
return value
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (f *ssoDeviceFlow) pollToken(ctx context.Context, deviceCode string, interval, expiresIn time.Duration) (*TokenResponse, error) {
|
||||
if interval < time.Second {
|
||||
interval = time.Second
|
||||
@@ -241,7 +283,7 @@ func (f *ssoDeviceFlow) pollToken(ctx context.Context, deviceCode string, interv
|
||||
return nil, errors.New("xAI device flow token polling timed out")
|
||||
}
|
||||
|
||||
func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form url.Values) (int, string, []byte, error) {
|
||||
func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form url.Values, headers ...http.Header) (int, string, []byte, error) {
|
||||
if !safeXAIAuthURL(endpoint) {
|
||||
return 0, "", nil, errors.New("xAI OAuth URL is not trusted")
|
||||
}
|
||||
@@ -260,6 +302,15 @@ func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form ur
|
||||
request.Header.Set("Accept", "application/json, text/html;q=0.9, */*;q=0.8")
|
||||
request.Header.Set("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8")
|
||||
request.Header.Set("User-Agent", f.userAgent)
|
||||
if redirects == 0 {
|
||||
for _, header := range headers {
|
||||
for name, values := range header {
|
||||
for _, value := range values {
|
||||
request.Header.Add(name, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if cookie := f.cookieHeader(request.URL); cookie != "" {
|
||||
request.Header.Set("Cookie", cookie)
|
||||
}
|
||||
|
||||
@@ -23,6 +23,10 @@ type ssoDeviceFakeClient struct {
|
||||
|
||||
func (c *ssoDeviceFakeClient) Do(req *http.Request) (*http.Response, error) {
|
||||
c.cookieHeaders = append(c.cookieHeaders, req.Header.Get("Cookie"))
|
||||
if req.URL.String() != SSOApproveURL {
|
||||
require.Empty(c.t, req.Header.Get("Origin"))
|
||||
require.Empty(c.t, req.Header.Get("Referer"))
|
||||
}
|
||||
switch req.URL.String() {
|
||||
case SSOAccountsURL:
|
||||
require.Equal(c.t, http.MethodGet, req.Method)
|
||||
@@ -43,13 +47,16 @@ func (c *ssoDeviceFakeClient) Do(req *http.Request) (*http.Response, error) {
|
||||
return ssoDeviceResponse(http.StatusFound, http.Header{"Location": {"/oauth2/device/consent"}}, ``), nil
|
||||
case "https://auth.x.ai/oauth2/device/consent":
|
||||
require.Equal(c.t, http.MethodGet, req.Method)
|
||||
return ssoDeviceResponse(http.StatusOK, nil, `<html>consent</html>`), nil
|
||||
return ssoDeviceResponse(http.StatusOK, nil, `<html><form><input type="hidden" name="consent_token" value="consent&token"/></form></html>`), nil
|
||||
case SSOApproveURL:
|
||||
require.Equal(c.t, http.MethodPost, req.Method)
|
||||
require.Equal(c.t, "https://auth.x.ai", req.Header.Get("Origin"))
|
||||
require.Equal(c.t, "https://auth.x.ai/oauth2/device/consent", req.Header.Get("Referer"))
|
||||
values := readSSODeviceForm(c.t, req)
|
||||
require.Equal(c.t, "USER-1", values.Get("user_code"))
|
||||
require.Equal(c.t, "allow", values.Get("action"))
|
||||
require.Equal(c.t, "User", values.Get("principal_type"))
|
||||
require.Equal(c.t, "consent&token", values.Get("consent_token"))
|
||||
return ssoDeviceResponse(http.StatusSeeOther, http.Header{"Location": {"/oauth2/device/done"}}, ``), nil
|
||||
case "https://auth.x.ai/oauth2/device/done":
|
||||
require.Equal(c.t, http.MethodGet, req.Method)
|
||||
@@ -87,6 +94,24 @@ func TestConvertSSOToBuildCompletesDeviceFlow(t *testing.T) {
|
||||
require.Contains(t, client.cookieHeaders[0], "sso-rw=sso-token")
|
||||
require.Contains(t, client.cookieHeaders[len(client.cookieHeaders)-1], "session=web-session")
|
||||
require.Contains(t, client.cookieHeaders[len(client.cookieHeaders)-1], "csrf=csrf-token")
|
||||
for _, cookie := range client.cookieHeaders {
|
||||
require.NotContains(t, cookie, "consent&token")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSSOConsentToken(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, body, want string
|
||||
}{
|
||||
{"legacy page", `<html>consent</html>`, ""},
|
||||
{"encoded hidden value", `<input value='a&b+c' name='consent_token' type='hidden'>`, "a&b+c"},
|
||||
{"case insensitive HTML", `<INPUT TYPE="HIDDEN" NAME="consent_token" VALUE="token"/>`, "token"},
|
||||
{"ignore script and visible fields", `<script>"<input type='hidden' name='consent_token' value='fake'>"</script><input name='consent_token' value='visible'><input type='hidden' name='user_code' value='user'><input type='hidden' name='consent_token' value='real'>`, "real"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
require.Equal(t, tc.want, ssoConsentToken([]byte(tc.body)))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeSSOTokenAcceptsCookieHeader(t *testing.T) {
|
||||
|
||||
@@ -33,18 +33,63 @@ func (r *opsRepository) GetDashboardOverview(ctx context.Context, filter *servic
|
||||
mode = service.OpsQueryModeRaw
|
||||
}
|
||||
|
||||
var overview *service.OpsDashboardOverview
|
||||
var err error
|
||||
switch mode {
|
||||
case service.OpsQueryModePreagg:
|
||||
return r.getDashboardOverviewPreaggregated(ctx, filter)
|
||||
overview, err = r.getDashboardOverviewPreaggregated(ctx, filter)
|
||||
case service.OpsQueryModeAuto:
|
||||
out, err := r.getDashboardOverviewPreaggregated(ctx, filter)
|
||||
if err != nil && errors.Is(err, service.ErrOpsPreaggregatedNotPopulated) {
|
||||
return r.getDashboardOverviewRaw(ctx, filter)
|
||||
overview, err = r.getDashboardOverviewPreaggregated(ctx, filter)
|
||||
if errors.Is(err, service.ErrOpsPreaggregatedNotPopulated) {
|
||||
overview, err = r.getDashboardOverviewRaw(ctx, filter)
|
||||
}
|
||||
return out, err
|
||||
default:
|
||||
return r.getDashboardOverviewRaw(ctx, filter)
|
||||
overview, err = r.getDashboardOverviewRaw(ctx, filter)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Percentiles must use individual rates over the entire window, even in
|
||||
// preagg mode: hourly percentiles cannot be averaged into window percentiles.
|
||||
rateCtx, cancelRate := context.WithTimeout(ctx, opsRawLatencyQueryTimeout)
|
||||
defer cancelRate()
|
||||
overview.OutputTPS, err = r.queryOutputTPS(rateCtx, filter)
|
||||
if err != nil && !isQueryTimeoutErr(err) {
|
||||
return nil, err
|
||||
}
|
||||
return overview, nil
|
||||
}
|
||||
|
||||
func (r *opsRepository) queryOutputTPS(ctx context.Context, filter *service.OpsDashboardFilter) (*service.OpsOutputTPS, error) {
|
||||
join, where, args, _ := buildUsageWhere(filter, filter.StartTime.UTC(), filter.EndTime.UTC(), 1)
|
||||
query := `
|
||||
SELECT
|
||||
percentile_cont(0.05) WITHIN GROUP (ORDER BY output_tps),
|
||||
percentile_cont(0.10) WITHIN GROUP (ORDER BY output_tps),
|
||||
percentile_cont(0.50) WITHIN GROUP (ORDER BY output_tps),
|
||||
AVG(output_tps),
|
||||
COUNT(*)
|
||||
FROM (
|
||||
SELECT ul.output_tokens * 1000.0 / NULLIF(ul.duration_ms, 0) AS output_tps
|
||||
FROM usage_logs ul
|
||||
` + join + `
|
||||
` + where + `
|
||||
AND ul.output_tokens > 0 AND ul.duration_ms > 0
|
||||
AND ul.image_count = 0 AND ul.image_output_tokens = 0
|
||||
AND COALESCE(ul.billing_mode, '') <> 'image'
|
||||
` + fmt.Sprintf("AND ul.request_type IN (%d, %d, %d, %d)",
|
||||
service.RequestTypeSync, service.RequestTypeStream, service.RequestTypeWSV2, service.RequestTypeCyberBlocked) + `
|
||||
) rates`
|
||||
var stats service.OpsOutputTPS
|
||||
if err := r.db.QueryRowContext(ctx, query, args...).Scan(&stats.P5, &stats.P10, &stats.P50, &stats.Avg, &stats.SampleCount); err != nil {
|
||||
// lib/pq reports query cancellation without wrapping the context error.
|
||||
// Preserve deadline vs caller cancellation for the overview's timeout policy.
|
||||
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||
return nil, ctxErr
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &stats, nil
|
||||
}
|
||||
|
||||
func (r *opsRepository) getDashboardOverviewRaw(ctx context.Context, filter *service.OpsDashboardFilter) (*service.OpsDashboardOverview, error) {
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOpsOutputTPSDistribution(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
user := mustCreateUser(t, client, &service.User{})
|
||||
key := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID})
|
||||
account := mustCreateAccount(t, client, &service.Account{Name: "output-tps", Platform: service.PlatformAnthropic})
|
||||
group := mustCreateGroup(t, client, &service.Group{Name: "output-tps", Platform: service.PlatformOpenAI})
|
||||
otherGroup := mustCreateGroup(t, client, &service.Group{Name: "other-output-tps", Platform: service.PlatformOpenAI})
|
||||
start := time.Date(2021, 1, 17, 0, 30, 0, 0, time.UTC)
|
||||
end := start.Add(6 * time.Hour)
|
||||
t.Cleanup(func() {
|
||||
_, _ = integrationDB.ExecContext(ctx, "DELETE FROM usage_logs WHERE user_id = $1", user.ID)
|
||||
_, _ = integrationDB.ExecContext(ctx, "DELETE FROM api_keys WHERE id = $1", key.ID)
|
||||
_, _ = integrationDB.ExecContext(ctx, "DELETE FROM accounts WHERE id = $1", account.ID)
|
||||
_, _ = integrationDB.ExecContext(ctx, "DELETE FROM ops_metrics_hourly WHERE bucket_start >= $1 AND bucket_start < $2", start.Truncate(time.Hour), end)
|
||||
_, _ = integrationDB.ExecContext(ctx, "DELETE FROM groups WHERE id IN ($1, $2)", group.ID, otherGroup.ID)
|
||||
_, _ = integrationDB.ExecContext(ctx, "DELETE FROM users WHERE id = $1", user.ID)
|
||||
})
|
||||
insert := func(at time.Time, groupID *int64, output int, duration *int, requestType service.RequestType, imageCount, imageTokens int, billingMode string) {
|
||||
t.Helper()
|
||||
_, err := integrationDB.ExecContext(ctx, `INSERT INTO usage_logs
|
||||
(user_id, api_key_id, account_id, group_id, request_id, model, input_tokens, output_tokens, duration_ms, request_type, image_count, image_output_tokens, billing_mode, created_at, image_size)
|
||||
VALUES ($1,$2,$3,$4,$5,'gpt-test',100000,$6,$7,$8,$9,$10,$11,$12,CASE WHEN $9 > 0 THEN '1K' ELSE NULL END)`,
|
||||
user.ID, key.ID, account.ID, groupID, fmt.Sprintf("output-tps-%d-%d", user.ID, time.Now().UnixNano()), output, duration, requestType, imageCount, imageTokens, billingMode, at)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
// Rates 1, 10, 100, 1000: unequal durations catch ratio-of-sums mistakes.
|
||||
// Put samples in the raw head, stable hourly buckets, and raw tail.
|
||||
for _, sample := range []struct {
|
||||
output, duration int
|
||||
kind service.RequestType
|
||||
at time.Time
|
||||
}{
|
||||
{1, 1000, service.RequestTypeSync, start},
|
||||
{20, 2000, service.RequestTypeStream, start.Add(time.Hour)},
|
||||
{100, 1000, service.RequestTypeWSV2, start.Add(2 * time.Hour)},
|
||||
{500, 500, service.RequestTypeCyberBlocked, end.Add(-time.Minute)},
|
||||
} {
|
||||
insert(sample.at, &group.ID, sample.output, &sample.duration, sample.kind, 0, 0, "token")
|
||||
}
|
||||
positive, zero, negative := 1000, 0, -1
|
||||
for _, d := range []*int{nil, &zero, &negative} {
|
||||
insert(start, &group.ID, 9999, d, service.RequestTypeSync, 0, 0, "token")
|
||||
}
|
||||
for _, n := range []int{0, -1} {
|
||||
insert(start, &group.ID, n, &positive, service.RequestTypeSync, 0, 0, "token")
|
||||
}
|
||||
for _, kind := range []service.RequestType{service.RequestTypeLive, service.RequestTypeUnknown} {
|
||||
insert(start, &group.ID, 9999, &positive, kind, 0, 0, "token")
|
||||
}
|
||||
insert(start, &group.ID, 9999, &positive, service.RequestTypeSync, 1, 0, "token")
|
||||
insert(start, &group.ID, 9999, &positive, service.RequestTypeSync, 0, 1, "token")
|
||||
insert(start, &group.ID, 9999, &positive, service.RequestTypeSync, 0, 0, "image")
|
||||
insert(start.Add(-time.Second), &group.ID, 9999, &positive, service.RequestTypeSync, 0, 0, "token")
|
||||
insert(end, &group.ID, 9999, &positive, service.RequestTypeSync, 0, 0, "token")
|
||||
insert(start, &otherGroup.ID, 9999, &positive, service.RequestTypeSync, 0, 0, "token")
|
||||
insert(start, nil, 50, &positive, service.RequestTypeSync, 0, 0, "token")
|
||||
repo := NewOpsRepository(integrationDB).(*opsRepository)
|
||||
filter := &service.OpsDashboardFilter{StartTime: start, EndTime: end, Platform: " OpenAI ", GroupID: &group.ID}
|
||||
assertStats := func(stats *service.OpsOutputTPS) {
|
||||
t.Helper()
|
||||
require.NotNil(t, stats)
|
||||
require.EqualValues(t, 4, stats.SampleCount)
|
||||
require.NotNil(t, stats.P5)
|
||||
require.NotNil(t, stats.P10)
|
||||
require.NotNil(t, stats.P50)
|
||||
require.NotNil(t, stats.Avg)
|
||||
require.InDelta(t, 2.35, *stats.P5, 1e-9)
|
||||
require.InDelta(t, 3.7, *stats.P10, 1e-9)
|
||||
require.InDelta(t, 55, *stats.P50, 1e-9)
|
||||
require.InDelta(t, 277.75, *stats.Avg, 1e-9)
|
||||
}
|
||||
// Auto first exercises fallback before hourly rows exist.
|
||||
for _, mode := range []service.OpsQueryMode{service.OpsQueryModeRaw, service.OpsQueryModeAuto} {
|
||||
filter.QueryMode = mode
|
||||
overview, err := repo.GetDashboardOverview(ctx, filter)
|
||||
require.NoError(t, err)
|
||||
assertStats(overview.OutputTPS)
|
||||
}
|
||||
require.NoError(t, repo.UpsertHourlyMetrics(ctx, start.Truncate(time.Hour), end))
|
||||
for _, mode := range []service.OpsQueryMode{service.OpsQueryModePreagg, service.OpsQueryModeAuto} {
|
||||
filter.QueryMode = mode
|
||||
overview, err := repo.GetDashboardOverview(ctx, filter)
|
||||
require.NoError(t, err)
|
||||
assertStats(overview.OutputTPS)
|
||||
}
|
||||
filter.Platform = service.PlatformAnthropic
|
||||
empty, err := repo.queryOutputTPS(ctx, filter)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, &service.OpsOutputTPS{}, empty)
|
||||
// Ungrouped logs use the account platform; group platform takes precedence.
|
||||
filter.GroupID = nil
|
||||
fallback, err := repo.queryOutputTPS(ctx, filter)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 1, fallback.SampleCount)
|
||||
require.Equal(t, 50.0, *fallback.P50)
|
||||
}
|
||||
|
||||
func TestOpsOutputTPSQueryDeadline(t *testing.T) {
|
||||
tx := testTx(t)
|
||||
_, err := tx.ExecContext(context.Background(), "LOCK TABLE usage_logs IN ACCESS EXCLUSIVE MODE")
|
||||
require.NoError(t, err)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
repo := NewOpsRepository(integrationDB).(*opsRepository)
|
||||
now := time.Now().UTC()
|
||||
stats, err := repo.queryOutputTPS(ctx, &service.OpsDashboardFilter{StartTime: now.Add(-time.Hour), EndTime: now})
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
require.Nil(t, stats)
|
||||
}
|
||||
@@ -2,8 +2,15 @@ package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql/driver"
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestIsQueryTimeoutErr(t *testing.T) {
|
||||
@@ -20,3 +27,42 @@ func TestIsQueryTimeoutErr(t *testing.T) {
|
||||
t.Fatalf("wrapped context.Canceled should not be treated as query timeout")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpsOutputTPSOverviewFailurePolicy(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
err error
|
||||
wantError bool
|
||||
}{
|
||||
{"timeout keeps overview", context.DeadlineExceeded, false},
|
||||
{"cancellation propagates", context.Canceled, true},
|
||||
{"database errors propagate", errors.New("database unavailable"), true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &opsRepository{db: db}
|
||||
// The six existing raw overview queries run before the output-rate query.
|
||||
for _, count := range []int{2, 13, 6, 2, 6, 2} {
|
||||
columns := make([]string, count)
|
||||
values := make([]driver.Value, count)
|
||||
for i := range columns {
|
||||
columns[i] = fmt.Sprint(i)
|
||||
values[i] = int64(0)
|
||||
}
|
||||
mock.ExpectQuery("SELECT|WITH").WillReturnRows(sqlmock.NewRows(columns).AddRow(values...))
|
||||
}
|
||||
mock.ExpectQuery(`percentile_cont\(0.05\)`).WillReturnError(tc.err)
|
||||
start := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
overview, err := repo.GetDashboardOverview(context.Background(), &service.OpsDashboardFilter{StartTime: start, EndTime: start.Add(time.Hour), QueryMode: service.OpsQueryModeRaw})
|
||||
if tc.wantError {
|
||||
require.ErrorIs(t, err, tc.err)
|
||||
require.Nil(t, overview)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, overview)
|
||||
require.Nil(t, overview.OutputTPS)
|
||||
}
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
dbmigrations "github.com/Wei-Shaw/sub2api/migrations"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func platformConstraintExists(t *testing.T, ctx context.Context, client *dbent.Client, table, constraint string) bool {
|
||||
t.Helper()
|
||||
rows, err := client.QueryContext(ctx, `
|
||||
SELECT 1
|
||||
FROM pg_constraint c
|
||||
JOIN pg_class tbl ON tbl.oid = c.conrelid
|
||||
JOIN pg_namespace ns ON ns.oid = tbl.relnamespace
|
||||
WHERE ns.nspname = 'public' AND tbl.relname = $1 AND c.conname = $2`, table, constraint)
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = rows.Close() }()
|
||||
exists := rows.Next()
|
||||
require.NoError(t, rows.Err())
|
||||
return exists
|
||||
}
|
||||
|
||||
// 迁移 242 删除平台白名单 CHECK 约束:旧库(约束仍在)上可重复执行,执行后
|
||||
// 新登记的平台无需再做数据库迁移即可写入。
|
||||
func TestMigration242DropsPlatformCheckConstraints(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
client := tx.Client()
|
||||
|
||||
// 还原迁移 241_add_typesafe_platform 之后、242 之前的约束状态。
|
||||
for _, stmt := range []string{
|
||||
`ALTER TABLE user_platform_quotas DROP CONSTRAINT IF EXISTS user_platform_quotas_platform_check`,
|
||||
`ALTER TABLE user_platform_quotas ADD CONSTRAINT user_platform_quotas_platform_check
|
||||
CHECK (platform IN ('anthropic', 'openai', 'gemini', 'antigravity', 'grok',
|
||||
'kimi', 'zhipu', 'deepseek', 'minimax', 'opencode_go', 'typesafe'))`,
|
||||
`ALTER TABLE composite_model_routes DROP CONSTRAINT IF EXISTS composite_model_routes_target_platform_check`,
|
||||
`ALTER TABLE composite_model_routes ADD CONSTRAINT composite_model_routes_target_platform_check
|
||||
CHECK (target_platform IN ('anthropic', 'openai', 'gemini', 'antigravity', 'grok',
|
||||
'kimi', 'zhipu', 'deepseek', 'minimax', 'opencode_go', 'typesafe'))`,
|
||||
} {
|
||||
_, err := client.ExecContext(ctx, stmt)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.True(t, platformConstraintExists(t, ctx, client, "user_platform_quotas", "user_platform_quotas_platform_check"))
|
||||
require.True(t, platformConstraintExists(t, ctx, client, "composite_model_routes", "composite_model_routes_target_platform_check"))
|
||||
|
||||
migrationSQL, err := dbmigrations.FS.ReadFile("242_drop_platform_check_constraints.sql")
|
||||
require.NoError(t, err)
|
||||
for range 2 {
|
||||
_, err = client.ExecContext(ctx, string(migrationSQL))
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.False(t, platformConstraintExists(t, ctx, client, "user_platform_quotas", "user_platform_quotas_platform_check"))
|
||||
require.False(t, platformConstraintExists(t, ctx, client, "composite_model_routes", "composite_model_routes_target_platform_check"))
|
||||
// 渠道监控的 provider 约束表示探测能力,保留。
|
||||
require.True(t, platformConstraintExists(t, ctx, client, "channel_monitors", "channel_monitors_provider_check"))
|
||||
require.True(t, platformConstraintExists(t, ctx, client, "channel_monitor_request_templates", "channel_monitor_request_templates_provider_check"))
|
||||
|
||||
// 数据库层不再限制平台取值:模拟未来新登记的平台。
|
||||
userID := mustCreateUserForQuota(t, client)
|
||||
_, err = client.ExecContext(ctx, `
|
||||
INSERT INTO user_platform_quotas (user_id, platform, daily_limit_usd, daily_usage_usd, weekly_usage_usd, monthly_usage_usd, created_at, updated_at)
|
||||
VALUES ($1, 'future_platform', 1, 0, 0, 0, NOW(), NOW())`, userID)
|
||||
require.NoError(t, err)
|
||||
group := mustCreateGroup(t, client, &service.Group{Name: "platform-list-composite", Platform: service.PlatformComposite})
|
||||
_, err = client.ExecContext(ctx, `
|
||||
INSERT INTO composite_model_routes (group_id, public_model, target_platform)
|
||||
VALUES ($1, 'future-model', 'future_platform')`, group.ID)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// 删除 CHECK 约束后,平台合法性由应用层保证:原生 SQL 写入前的 repository 校验。
|
||||
func TestUserPlatformQuotaRepository_RejectsUnregisteredPlatform(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
client := tx.Client()
|
||||
|
||||
userID := mustCreateUserForQuota(t, client)
|
||||
repo := NewUserPlatformQuotaRepository(client)
|
||||
daily := 5.0
|
||||
require.NoError(t, repo.UpsertForUser(txCtx, userID, []UserPlatformQuotaRecord{
|
||||
{UserID: userID, Platform: service.PlatformOpenCodeGo, DailyLimitUSD: &daily},
|
||||
}))
|
||||
|
||||
for _, platform := range []string{"bogus", "moonshot", "Kimi", service.PlatformComposite} {
|
||||
err := repo.BulkInsertInitial(txCtx, []UserPlatformQuotaRecord{
|
||||
{UserID: userID, Platform: service.PlatformAnthropic, DailyLimitUSD: &daily},
|
||||
{UserID: userID, Platform: platform, DailyLimitUSD: &daily},
|
||||
})
|
||||
require.ErrorContains(t, err, "is not allowed", platform)
|
||||
err = repo.UpsertForUser(txCtx, userID, []UserPlatformQuotaRecord{
|
||||
{UserID: userID, Platform: platform, DailyLimitUSD: &daily},
|
||||
})
|
||||
require.ErrorContains(t, err, "is not allowed", platform)
|
||||
}
|
||||
|
||||
// 被拒绝的写入不产生任何副作用:既有记录保留、未新增任何行。
|
||||
got, err := repo.ListByUser(txCtx, userID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, got, 1)
|
||||
require.Equal(t, service.PlatformOpenCodeGo, got[0].Platform)
|
||||
|
||||
// 三档全空的记录不写入,也不参与校验(与既有 configuredRecords 语义一致)。
|
||||
require.NoError(t, repo.BulkInsertInitial(txCtx, []UserPlatformQuotaRecord{{UserID: userID, Platform: "bogus"}}))
|
||||
}
|
||||
|
||||
func TestCompositeModelRouteRepository_RejectsUnregisteredTargetPlatform(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
client := tx.Client()
|
||||
|
||||
group := mustCreateGroup(t, client, &service.Group{Name: "platform-list-routes", Platform: service.PlatformComposite})
|
||||
repo := NewCompositeModelRouteRepository(client)
|
||||
|
||||
route := &service.CompositeModelRoute{
|
||||
GroupID: group.ID, PublicModel: "glm-5", MatchType: "exact", TargetPlatform: service.PlatformZhipu,
|
||||
UpstreamModel: "glm-5", Endpoint: "any", Priority: 100, Enabled: true,
|
||||
}
|
||||
require.NoError(t, repo.Create(txCtx, route))
|
||||
|
||||
bad := &service.CompositeModelRoute{
|
||||
GroupID: group.ID, PublicModel: "bogus-model", MatchType: "exact", TargetPlatform: "bogus",
|
||||
UpstreamModel: "bogus-model", Endpoint: "any", Priority: 100, Enabled: true,
|
||||
}
|
||||
require.Error(t, repo.Create(txCtx, bad))
|
||||
|
||||
route.TargetPlatform = service.PlatformComposite
|
||||
require.Error(t, repo.Update(txCtx, route))
|
||||
|
||||
routes, err := repo.ListByGroup(txCtx, group.ID, true)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, routes, 1)
|
||||
require.Equal(t, service.PlatformZhipu, routes[0].TargetPlatform)
|
||||
}
|
||||
@@ -187,16 +187,17 @@ func (r *usageBillingRepository) applyUsageBillingEffects(ctx context.Context, t
|
||||
result.BalanceOverdrafted = !sufficient
|
||||
}
|
||||
|
||||
// Key 已不存在时跳过其自身的额度/限速计数,其余结算项不受影响。
|
||||
if cmd.APIKeyQuotaCost > 0 {
|
||||
exhausted, err := incrementUsageBillingAPIKeyQuota(ctx, tx, cmd.APIKeyID, cmd.APIKeyQuotaCost)
|
||||
if err != nil {
|
||||
if err != nil && !errors.Is(err, service.ErrAPIKeyNotFound) {
|
||||
return err
|
||||
}
|
||||
result.APIKeyQuotaExhausted = exhausted
|
||||
}
|
||||
|
||||
if cmd.APIKeyRateLimitCost > 0 {
|
||||
if err := incrementUsageBillingAPIKeyRateLimit(ctx, tx, cmd.APIKeyID, cmd.APIKeyRateLimitCost); err != nil {
|
||||
if err := incrementUsageBillingAPIKeyRateLimit(ctx, tx, cmd.APIKeyID, cmd.APIKeyRateLimitCost); err != nil && !errors.Is(err, service.ErrAPIKeyNotFound) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
@@ -162,6 +162,61 @@ func TestUsageBillingRepositoryApply_RequestFingerprintConflict(t *testing.T) {
|
||||
require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict)
|
||||
}
|
||||
|
||||
func TestUsageBillingRepositoryApply_DeletedAPIKeyStillBillsBalance(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
repo := NewUsageBillingRepository(client, integrationDB)
|
||||
|
||||
user := mustCreateUser(t, client, &service.User{
|
||||
Email: fmt.Sprintf("usage-billing-deleted-key-user-%d@example.com", time.Now().UnixNano()),
|
||||
PasswordHash: "hash",
|
||||
Balance: 100,
|
||||
})
|
||||
apiKey := mustCreateApiKey(t, client, &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-usage-billing-deleted-key-" + uuid.NewString(),
|
||||
Name: "billing-deleted-key",
|
||||
Quota: 50,
|
||||
RateLimit5h: 50,
|
||||
})
|
||||
account := mustCreateAccount(t, client, &service.Account{
|
||||
Name: "usage-billing-deleted-key-account-" + uuid.NewString(),
|
||||
Type: service.AccountTypeAPIKey,
|
||||
})
|
||||
|
||||
_, err := integrationDB.ExecContext(ctx, "UPDATE api_keys SET deleted_at = NOW() WHERE id = $1", apiKey.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
requestID := uuid.NewString()
|
||||
result, err := repo.Apply(ctx, &service.UsageBillingCommand{
|
||||
RequestID: requestID,
|
||||
APIKeyID: apiKey.ID,
|
||||
UserID: user.ID,
|
||||
AccountID: account.ID,
|
||||
AccountType: service.AccountTypeAPIKey,
|
||||
BalanceCost: 1.25,
|
||||
APIKeyQuotaCost: 1.25,
|
||||
APIKeyRateLimitCost: 1.25,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.Applied)
|
||||
require.False(t, result.APIKeyQuotaExhausted)
|
||||
|
||||
var balance float64
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT balance FROM users WHERE id = $1", user.ID).Scan(&balance))
|
||||
require.InDelta(t, 98.75, balance, 0.000001)
|
||||
|
||||
var quotaUsed, usage5h float64
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT quota_used, usage_5h FROM api_keys WHERE id = $1", apiKey.ID).Scan("aUsed, &usage5h))
|
||||
require.InDelta(t, 0, quotaUsed, 0.000001)
|
||||
require.InDelta(t, 0, usage5h, 0.000001)
|
||||
|
||||
var dedupCount int
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&dedupCount))
|
||||
require.Equal(t, 1, dedupCount)
|
||||
}
|
||||
|
||||
func TestUsageBillingRepositoryApply_UpdatesAccountQuota(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
|
||||
@@ -20,6 +20,8 @@ const (
|
||||
captureBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance\s+\+ CASE WHEN \$1 > \$2 THEN \$1 - \$2 ELSE 0 END\s+- CASE WHEN \$2 > \$1 THEN \$2 - \$1 ELSE 0 END,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$3 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance`
|
||||
releaseBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance \+ \$1,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance`
|
||||
userExistsForBillingSQL = `(?s)SELECT 1\s+FROM users\s+WHERE id = \$1 AND deleted_at IS NULL`
|
||||
apiKeyQuotaIncrementSQL = `(?s)UPDATE api_keys\s+SET quota_used = quota_used \+ \$1,.*WHERE id = \$2 AND deleted_at IS NULL\s+RETURNING`
|
||||
apiKeyRateLimitIncrementSQL = `(?s)UPDATE api_keys SET\s+usage_5h = .*WHERE id = \$2 AND deleted_at IS NULL`
|
||||
)
|
||||
|
||||
func TestDeductUsageBillingBalance_UsesSufficientBalanceGuard(t *testing.T) {
|
||||
@@ -99,6 +101,66 @@ func TestApplyUsageBillingEffects_FlagsBalanceOverdraft(t *testing.T) {
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestApplyUsageBillingEffects_DeletedAPIKeyStillBillsBalance(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
mock.ExpectBegin()
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
mock.ExpectQuery(conditionalBalanceDeductSQL).
|
||||
WithArgs(10.0, int64(42)).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(90.0))
|
||||
mock.ExpectQuery(apiKeyQuotaIncrementSQL).
|
||||
WithArgs(10.0, int64(7), service.StatusAPIKeyActive, service.StatusAPIKeyQuotaExhausted).
|
||||
WillReturnError(sql.ErrNoRows)
|
||||
mock.ExpectExec(apiKeyRateLimitIncrementSQL).
|
||||
WithArgs(10.0, int64(7)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectCommit()
|
||||
|
||||
result := &service.UsageBillingApplyResult{Applied: true}
|
||||
err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{
|
||||
UserID: 42,
|
||||
APIKeyID: 7,
|
||||
BalanceCost: 10,
|
||||
APIKeyQuotaCost: 10,
|
||||
APIKeyRateLimitCost: 10,
|
||||
}, result)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result.NewBalance)
|
||||
require.InDelta(t, 90.0, *result.NewBalance, 0.000001)
|
||||
require.False(t, result.APIKeyQuotaExhausted)
|
||||
require.NoError(t, tx.Commit())
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestApplyUsageBillingEffects_APIKeyCounterErrorStillFails(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
mock.ExpectBegin()
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
mock.ExpectExec(apiKeyRateLimitIncrementSQL).
|
||||
WithArgs(10.0, int64(7)).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{
|
||||
UserID: 42,
|
||||
APIKeyID: 7,
|
||||
APIKeyRateLimitCost: 10,
|
||||
}, &service.UsageBillingApplyResult{Applied: true})
|
||||
require.ErrorIs(t, err, sql.ErrConnDone)
|
||||
require.NoError(t, tx.Rollback())
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDeductUsageBillingBalance_ReturnsUserNotFoundWhenNoUserUpdated(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db, mock, err := sqlmock.New()
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/ent/userplatformquota"
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
)
|
||||
|
||||
@@ -95,6 +96,9 @@ func (r *userPlatformQuotaRepository) BulkInsertInitial(ctx context.Context, rec
|
||||
if len(records) == 0 {
|
||||
return nil
|
||||
}
|
||||
if err := validateQuotaPlatforms(records); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
client := clientFromContext(ctx, r.client)
|
||||
|
||||
@@ -318,6 +322,9 @@ func monthlyMaybeReset(prevUsage float64, prevStart *time.Time, cost float64, no
|
||||
// 清空某平台的限额即软删该行并放弃其累计用量;之后再次配置限额会新建行、从新窗口起算。
|
||||
func (r *userPlatformQuotaRepository) UpsertForUser(ctx context.Context, userID int64, records []UserPlatformQuotaRecord) error {
|
||||
records = configuredRecords(records)
|
||||
if err := validateQuotaPlatforms(records); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.withTx(ctx, func(txCtx context.Context, txClient *dbent.Client) error {
|
||||
platforms := make([]string, 0, len(records))
|
||||
for _, rec := range records {
|
||||
@@ -484,6 +491,17 @@ func (r *userPlatformQuotaRepository) BatchSnapshotUsage(ctx context.Context, sn
|
||||
return nil
|
||||
}
|
||||
|
||||
// validateQuotaPlatforms 在原生 SQL 写入前校验平台须为平台清单中的具体平台,
|
||||
// 取代迁移 242 删除的 user_platform_quotas_platform_check 约束(原生 SQL 不经过 ent Validate)。
|
||||
func validateQuotaPlatforms(records []UserPlatformQuotaRecord) error {
|
||||
for _, rec := range records {
|
||||
if !domain.IsConcretePlatform(rec.Platform) {
|
||||
return fmt.Errorf("user_platform_quotas: platform %q is not allowed", rec.Platform)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// configuredRecords 过滤出至少配置了一档限额的记录;三档全空的记录视为未配置。
|
||||
func configuredRecords(records []UserPlatformQuotaRecord) []UserPlatformQuotaRecord {
|
||||
out := make([]UserPlatformQuotaRecord, 0, len(records))
|
||||
|
||||
@@ -863,7 +863,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"force_email_on_third_party_signup": false,
|
||||
"default_concurrency": 5,
|
||||
"default_balance": 1.25,
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"minimax":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"opencode_go":{"daily":null,"weekly":null,"monthly":null},"typesafe":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"cline":{"daily":null,"weekly":null,"monthly":null},"command_code":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"minimax":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"opencode_go":{"daily":null,"weekly":null,"monthly":null},"typesafe":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"auth_source_default_email_platform_quotas": null,
|
||||
"auth_source_default_github_platform_quotas": null,
|
||||
"auth_source_default_google_platform_quotas": null,
|
||||
@@ -1189,7 +1189,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"purchase_subscription_url": "",
|
||||
"table_default_page_size": 20,
|
||||
"table_page_size_options": [10, 20, 50],
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"minimax":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"opencode_go":{"daily":null,"weekly":null,"monthly":null},"typesafe":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"cline":{"daily":null,"weekly":null,"monthly":null},"command_code":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"minimax":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"opencode_go":{"daily":null,"weekly":null,"monthly":null},"typesafe":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"auth_source_default_email_platform_quotas": null,
|
||||
"auth_source_default_github_platform_quotas": null,
|
||||
"auth_source_default_google_platform_quotas": null,
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/requestmodel"
|
||||
@@ -46,22 +47,15 @@ func RegisterGatewayRoutes(
|
||||
groupModelAllowlist := middleware.GroupModelAllowlist()
|
||||
|
||||
isOpenAIResponsesCompatibleGatewayPlatform := func(c *gin.Context) bool {
|
||||
switch getGroupPlatform(c) {
|
||||
case service.PlatformOpenAI, service.PlatformGrok,
|
||||
service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek,
|
||||
service.PlatformMiniMax, service.PlatformOpenCodeGo:
|
||||
// 国产 OpenAI 兼容供应商与 openai/grok 一样经 OpenAI 网关转发。
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
// openai、grok 与多协议 API Key 供应商经 OpenAI 网关转发(平台清单)。
|
||||
return domain.UsesOpenAIGateway(getGroupPlatform(c))
|
||||
}
|
||||
countTokensHandler := func(c *gin.Context) {
|
||||
switch getGroupPlatform(c) {
|
||||
case service.PlatformOpenAI, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax, service.PlatformOpenCodeGo:
|
||||
h.OpenAIGateway.CountTokens(c)
|
||||
case service.PlatformGrok:
|
||||
switch platform := getGroupPlatform(c); {
|
||||
case platform == service.PlatformGrok:
|
||||
h.OpenAIGateway.GrokCountTokens(c)
|
||||
case domain.UsesOpenAIGateway(platform):
|
||||
h.OpenAIGateway.CountTokens(c)
|
||||
default:
|
||||
h.Gateway.CountTokens(c)
|
||||
}
|
||||
|
||||
@@ -303,7 +303,7 @@ func (a *Account) IsCNProvider() bool {
|
||||
// openai/grok 原生走 OpenAI 网关;国产供应商同为 OpenAI Chat Completions
|
||||
// 兼容上游,也经 OpenAI 网关转发。OpenCode 同样经 OpenAI 网关按模型分流。
|
||||
func (a *Account) IsOpenAICompatible() bool {
|
||||
return a != nil && (a.Platform == PlatformOpenAI || a.Platform == PlatformGrok || a.IsCNProvider() || a.IsOpenCodeGo())
|
||||
return a != nil && domain.UsesOpenAIGateway(a.Platform)
|
||||
}
|
||||
|
||||
func (a *Account) GeminiOAuthType() string {
|
||||
@@ -1384,7 +1384,7 @@ func (a *Account) IsOpenAIApiKey() bool {
|
||||
// 适用 openai、国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)与 OpenCode Go;
|
||||
// grok 走 GetGrokBaseURL,此处对 grok 返回 "" 以保持原有行为。
|
||||
func (a *Account) GetOpenAIBaseURL() string {
|
||||
if !a.IsOpenAI() && !a.IsCNProvider() && !a.IsOpenCodeGo() {
|
||||
if !a.IsOpenAI() && !a.IsMultiProtocolAPIKey() {
|
||||
return ""
|
||||
}
|
||||
if a.IsMultiProtocolAPIKey() && a.IsAdaptiveAPIProtocol() {
|
||||
@@ -1399,27 +1399,11 @@ func (a *Account) GetOpenAIBaseURL() string {
|
||||
return baseURL
|
||||
}
|
||||
}
|
||||
// 平台默认 base_url:CN 供应商按 account_mode 选择 payg / coding 默认值。
|
||||
switch a.Platform {
|
||||
case PlatformKimi:
|
||||
if a.GetAccountMode() == AccountModeCoding {
|
||||
return DefaultKimiCodingBaseURL
|
||||
}
|
||||
return DefaultKimiPayGBaseURL
|
||||
case PlatformZhipu:
|
||||
if a.GetAccountMode() == AccountModeCoding {
|
||||
return DefaultZhipuCodingBaseURL
|
||||
}
|
||||
return DefaultZhipuPayGBaseURL
|
||||
case PlatformDeepseek:
|
||||
return DefaultDeepseekBaseURL
|
||||
case PlatformMiniMax:
|
||||
return DefaultMiniMaxBaseURL
|
||||
case PlatformOpenCodeGo:
|
||||
return a.openCodeDefaultChatBaseURL()
|
||||
default:
|
||||
return "https://api.openai.com"
|
||||
// 平台默认 base_url:多协议供应商按 account_mode 查 provider profile。
|
||||
if baseURL := a.defaultProviderBaseURL(APIProtocolChatCompletions); baseURL != "" {
|
||||
return baseURL
|
||||
}
|
||||
return "https://api.openai.com"
|
||||
}
|
||||
|
||||
// GetAccountMode 返回国产供应商账号的接入模式(payg / coding);非国产供应商或未设置时
|
||||
@@ -1460,7 +1444,8 @@ func (a *Account) GetAPIProtocol() string {
|
||||
case APIProtocolChatCompletions:
|
||||
return APIProtocolChatCompletions
|
||||
}
|
||||
if a.IsOpenCodeGo() {
|
||||
// 按模型分流的供应商(多模型聚合平台)未显式配置时默认 adaptive。
|
||||
if a.routesByModel() {
|
||||
return APIProtocolAdaptive
|
||||
}
|
||||
return APIProtocolChatCompletions
|
||||
@@ -1470,15 +1455,7 @@ func (a *Account) GetAPIProtocol() string {
|
||||
// DeepSeek 官方为 /responses(无 /v1);Kimi 按量付费与 Coding Plan 均为
|
||||
// /v1/responses(moonshot.cn / kimi.com/coding);MiniMax 为 /v1/responses。
|
||||
func (a *Account) SupportsNativeCNResponses() bool {
|
||||
if a == nil {
|
||||
return false
|
||||
}
|
||||
switch a.Platform {
|
||||
case PlatformDeepseek, PlatformKimi, PlatformMiniMax, PlatformOpenCodeGo:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return a.providerSupportsProtocol(APIProtocolResponses)
|
||||
}
|
||||
|
||||
// UsesNativeCNResponses 报告当前账号是否应按原生 Responses 协议转发
|
||||
@@ -1519,48 +1496,7 @@ func (a *Account) GetCNProtocolBaseURL(protocol string) string {
|
||||
}
|
||||
}
|
||||
}
|
||||
return a.defaultCNProtocolBaseURL(protocol)
|
||||
}
|
||||
|
||||
func (a *Account) defaultCNProtocolBaseURL(protocol string) string {
|
||||
switch protocol {
|
||||
case APIProtocolAnthropic:
|
||||
switch a.Platform {
|
||||
case PlatformKimi:
|
||||
if a.GetAccountMode() == AccountModeCoding {
|
||||
return DefaultKimiCodingAnthropicBaseURL
|
||||
}
|
||||
return DefaultKimiPayGAnthropicBaseURL
|
||||
case PlatformZhipu:
|
||||
return DefaultZhipuAnthropicBaseURL
|
||||
case PlatformDeepseek:
|
||||
return DefaultDeepseekAnthropicBaseURL
|
||||
case PlatformMiniMax:
|
||||
return DefaultMiniMaxAnthropicBaseURL
|
||||
case PlatformOpenCodeGo:
|
||||
return a.openCodeDefaultAnthropicBaseURL()
|
||||
}
|
||||
case APIProtocolChatCompletions, APIProtocolResponses:
|
||||
switch a.Platform {
|
||||
case PlatformKimi:
|
||||
if a.GetAccountMode() == AccountModeCoding {
|
||||
return DefaultKimiCodingBaseURL
|
||||
}
|
||||
return DefaultKimiPayGBaseURL
|
||||
case PlatformZhipu:
|
||||
if a.GetAccountMode() == AccountModeCoding {
|
||||
return DefaultZhipuCodingBaseURL
|
||||
}
|
||||
return DefaultZhipuPayGBaseURL
|
||||
case PlatformDeepseek:
|
||||
return DefaultDeepseekBaseURL
|
||||
case PlatformMiniMax:
|
||||
return DefaultMiniMaxBaseURL
|
||||
case PlatformOpenCodeGo:
|
||||
return a.openCodeDefaultChatBaseURL()
|
||||
}
|
||||
}
|
||||
return ""
|
||||
return a.defaultProviderBaseURL(protocol)
|
||||
}
|
||||
|
||||
// IsAnthropicProtocol 报告账号是否以原生 Anthropic 协议接入上游
|
||||
@@ -1584,23 +1520,7 @@ func (a *Account) GetAnthropicProtocolBaseURL() string {
|
||||
return baseURL
|
||||
}
|
||||
}
|
||||
switch a.Platform {
|
||||
case PlatformKimi:
|
||||
if a.GetAccountMode() == AccountModeCoding {
|
||||
return DefaultKimiCodingAnthropicBaseURL
|
||||
}
|
||||
return DefaultKimiPayGAnthropicBaseURL
|
||||
case PlatformZhipu:
|
||||
return DefaultZhipuAnthropicBaseURL
|
||||
case PlatformDeepseek:
|
||||
return DefaultDeepseekAnthropicBaseURL
|
||||
case PlatformMiniMax:
|
||||
return DefaultMiniMaxAnthropicBaseURL
|
||||
case PlatformOpenCodeGo:
|
||||
return a.openCodeDefaultAnthropicBaseURL()
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
return a.defaultProviderBaseURL(APIProtocolAnthropic)
|
||||
}
|
||||
|
||||
// GetOpenAIFormatBaseURL 返回供 OpenAI 格式端点(/v1/models、/v1/chat/completions
|
||||
@@ -1612,26 +1532,10 @@ func (a *Account) GetOpenAIFormatBaseURL() string {
|
||||
if a == nil || !a.IsAnthropicProtocol() {
|
||||
return a.GetOpenAIBaseURL()
|
||||
}
|
||||
switch a.Platform {
|
||||
case PlatformKimi:
|
||||
if a.GetAccountMode() == AccountModeCoding {
|
||||
return DefaultKimiCodingBaseURL
|
||||
}
|
||||
return DefaultKimiPayGBaseURL
|
||||
case PlatformZhipu:
|
||||
if a.GetAccountMode() == AccountModeCoding {
|
||||
return DefaultZhipuCodingBaseURL
|
||||
}
|
||||
return DefaultZhipuPayGBaseURL
|
||||
case PlatformDeepseek:
|
||||
return DefaultDeepseekBaseURL
|
||||
case PlatformMiniMax:
|
||||
return DefaultMiniMaxBaseURL
|
||||
case PlatformOpenCodeGo:
|
||||
return a.openCodeDefaultChatBaseURL()
|
||||
default:
|
||||
return a.GetOpenAIBaseURL()
|
||||
if baseURL := a.defaultProviderBaseURL(APIProtocolChatCompletions); baseURL != "" {
|
||||
return baseURL
|
||||
}
|
||||
return a.GetOpenAIBaseURL()
|
||||
}
|
||||
|
||||
// GetCNAPIKey 返回国产 OpenAI 兼容供应商账号的 api_key 凭据(kimi/zhipu/deepseek)。
|
||||
@@ -1643,7 +1547,8 @@ func (a *Account) GetCNAPIKey() string {
|
||||
return a.GetCredential("api_key")
|
||||
}
|
||||
|
||||
// GetCodingPlanProvider 根据 base_url 识别 Coding Plan 供应商(kimi / zhipu / minimax),
|
||||
// GetCodingPlanProvider 根据 base_url 识别 Coding Plan 供应商(kimi / zhipu / minimax;
|
||||
// OpenCode Go 订阅与官方主机上的 Command Code、Cline 账号按平台识别),
|
||||
// 用于路由到对应的额度查询端点。非 coding 模式或无法识别时返回空串。
|
||||
// 只认官方域名:自定义中转不得把第三方 Key 发往厂商官方额度端点。
|
||||
func (a *Account) GetCodingPlanProvider() string {
|
||||
@@ -1653,6 +1558,18 @@ func (a *Account) GetCodingPlanProvider() string {
|
||||
if a.IsOpenCodeGoPlan() {
|
||||
return PlatformOpenCodeGo
|
||||
}
|
||||
if a.IsCommandCode() {
|
||||
if a.commandCodeUsageSupported() {
|
||||
return PlatformCommandCode
|
||||
}
|
||||
return ""
|
||||
}
|
||||
if a.IsCline() {
|
||||
if a.clineAccountAPISupported() {
|
||||
return PlatformCline
|
||||
}
|
||||
return ""
|
||||
}
|
||||
if a.GetAccountMode() != AccountModeCoding {
|
||||
return ""
|
||||
}
|
||||
|
||||
@@ -77,10 +77,10 @@ func (a *Account) IsHeaderOverrideEligible() bool {
|
||||
if a == nil {
|
||||
return false
|
||||
}
|
||||
switch a.Platform {
|
||||
case PlatformAnthropic, PlatformOpenAI, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo:
|
||||
switch {
|
||||
case a.Platform == PlatformAnthropic, a.Platform == PlatformOpenAI, IsMultiProtocolAPIKeyProvider(a.Platform):
|
||||
return a.Type == AccountTypeAPIKey
|
||||
case PlatformGrok:
|
||||
case a.Platform == PlatformGrok:
|
||||
return a.Type == AccountTypeAPIKey || a.Type == AccountTypeOAuth
|
||||
default:
|
||||
return false
|
||||
|
||||
@@ -58,7 +58,7 @@ func EvaluateAccountSchedulingThreshold(account *Account, thresholds map[string]
|
||||
winner = pickLatestResetSchedulingCandidate(anthropicThresholdCandidates(account), threshold, now)
|
||||
case PlatformGrok:
|
||||
winner = pickLatestResetSchedulingCandidate(grokThresholdCandidates(account), threshold, now)
|
||||
case PlatformKimi, PlatformZhipu, PlatformMiniMax, PlatformOpenCodeGo:
|
||||
case PlatformKimi, PlatformZhipu, PlatformMiniMax, PlatformOpenCodeGo, PlatformCommandCode:
|
||||
winner = pickLatestResetSchedulingCandidate(cnProviderThresholdCandidates(account, decision.Platform), threshold, now)
|
||||
default:
|
||||
return decision
|
||||
@@ -353,25 +353,33 @@ func grokThresholdCandidates(account *Account) []*accountSchedulingThresholdCand
|
||||
}
|
||||
}
|
||||
|
||||
// cnProviderThresholdCandidates 读取国产供应商 Coding Plan 账号的 5h / weekly 滚动窗口
|
||||
// 用量快照(由 CNProviderQuotaService 写入 account.Extra,键形如
|
||||
// <provider>_5h_used_percent / <provider>_weekly_reset_at)。payg 账号无此快照,
|
||||
// cnProviderThresholdCandidates 读取国产供应商 Coding Plan、OpenCode Go 与 Command Code
|
||||
// 账号的 5h / weekly(及月度)窗口用量快照(由 CNProviderQuotaService 写入 account.Extra,
|
||||
// 键形如 <provider>_5h_used_percent / <provider>_weekly_reset_at)。payg 账号无此快照,
|
||||
// 候选为空 → 不触发阈值停调(余额型走余额检测)。与 openai 的快照驱动停调一致:
|
||||
// 仅当用量超阈值且窗口尚未重置时才停调。
|
||||
func cnProviderThresholdCandidates(account *Account, provider string) []*accountSchedulingThresholdCandidate {
|
||||
if account == nil || len(account.Extra) == 0 {
|
||||
return nil
|
||||
}
|
||||
if provider == PlatformCommandCode && !commandCodeWindowsBinding(account.Extra) {
|
||||
return nil
|
||||
}
|
||||
candidates := []*accountSchedulingThresholdCandidate{
|
||||
cnThresholdCandidate(account.Extra, provider, "5h"),
|
||||
cnThresholdCandidate(account.Extra, provider, "weekly"),
|
||||
}
|
||||
if provider == PlatformOpenCodeGo {
|
||||
if cnQuotaHasMonthlyWindow(provider) {
|
||||
candidates = append(candidates, cnThresholdCandidate(account.Extra, provider, "monthly"))
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
// cnQuotaHasMonthlyWindow 报告供应商的额度快照是否含月度窗口(OpenCode Go、Command Code)。
|
||||
func cnQuotaHasMonthlyWindow(provider string) bool {
|
||||
return provider == PlatformOpenCodeGo || provider == PlatformCommandCode
|
||||
}
|
||||
|
||||
func cnThresholdCandidate(extra map[string]any, provider, window string) *accountSchedulingThresholdCandidate {
|
||||
var usedKey, resetKey string
|
||||
switch window {
|
||||
|
||||
@@ -522,10 +522,12 @@ func (s *AccountService) TestCredentials(ctx context.Context, id int64) error {
|
||||
case PlatformTypeSafe:
|
||||
// TypeSafe credentials are API keys; inference failures drive health and cooldown state.
|
||||
return nil
|
||||
case PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo:
|
||||
// 国产 OpenAI 兼容供应商与 OpenCode:凭证为 API Key,实际可用性经余额/额度探测与转发路径验证。
|
||||
return nil
|
||||
default:
|
||||
if IsMultiProtocolAPIKeyProvider(account.Platform) {
|
||||
// 多协议 API Key 供应商(国产厂商与聚合平台):凭证为 API Key,实际可用性
|
||||
// 经余额/额度探测与转发路径验证。
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("unsupported platform: %s", account.Platform)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/logredact"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const accountTestBackgroundKey = "account_test_background"
|
||||
|
||||
// Bind attribution before lookup so missing accounts are identifiable too. The
|
||||
// request logger retains HTTP request IDs or the caller's background correlation.
|
||||
func initAccountTestLogger(c *gin.Context, accountID int64, modelID, mode string) {
|
||||
source := "http"
|
||||
if c.GetBool(accountTestBackgroundKey) {
|
||||
source = "background"
|
||||
}
|
||||
ctx := c.Request.Context()
|
||||
l := logger.FromContext(ctx).With(
|
||||
zap.String("component", "service.account_test"),
|
||||
zap.Int64("account_id", accountID),
|
||||
zap.String("test_id", uuid.NewString()),
|
||||
zap.String("test_source", source),
|
||||
zap.String("requested_model", modelID),
|
||||
zap.String("test_mode", normalizeAccountTestMode(mode)),
|
||||
)
|
||||
c.Request = c.Request.WithContext(logger.IntoContext(ctx, l))
|
||||
}
|
||||
|
||||
func bindAccountTestPlatform(c *gin.Context, account *Account) {
|
||||
ctx := c.Request.Context()
|
||||
l := logger.FromContext(ctx).With(
|
||||
zap.String("platform", account.Platform),
|
||||
zap.String("account_type", account.Type),
|
||||
)
|
||||
c.Request = c.Request.WithContext(logger.IntoContext(ctx, l))
|
||||
}
|
||||
|
||||
func logAccountTestError(c *gin.Context, message string) {
|
||||
// Redact known credential fields before bounding text; truncating first can
|
||||
// cut a quoted credential in half and prevent the redactor from matching it.
|
||||
message = logredact.RedactText(message, "api_key", "authorization")
|
||||
logger.FromContext(c.Request.Context()).Warn("account_test.failed",
|
||||
zap.String("error", truncateString(message, 2048)),
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/zap"
|
||||
"go.uber.org/zap/zaptest/observer"
|
||||
)
|
||||
|
||||
type accountTestLoggingRepo struct {
|
||||
AccountRepository
|
||||
account *Account
|
||||
}
|
||||
|
||||
func (r *accountTestLoggingRepo) GetByID(context.Context, int64) (*Account, error) {
|
||||
if r.account == nil {
|
||||
return nil, errors.New("missing")
|
||||
}
|
||||
return r.account, nil
|
||||
}
|
||||
|
||||
func TestAccountTestFailureLogAttribution(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, background := range []bool{false, true} {
|
||||
for _, missing := range []bool{false, true} {
|
||||
core, logs := observer.New(zap.WarnLevel)
|
||||
ctx := logger.IntoContext(context.Background(), zap.New(core).With(zap.String("request_id", "req-test")))
|
||||
repo := &accountTestLoggingRepo{}
|
||||
if !missing {
|
||||
repo.account = &Account{ID: 77, Platform: PlatformAnthropic, Type: AccountTypeAPIKey}
|
||||
}
|
||||
svc := &AccountTestService{accountRepo: repo}
|
||||
if background {
|
||||
result, err := svc.RunTestBackground(ctx, 77, "test-model")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "failed", result.Status)
|
||||
require.NotEmpty(t, result.ErrorMessage)
|
||||
} else {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest("POST", "/accounts/77/test", nil).WithContext(ctx)
|
||||
err := svc.TestAccountConnection(c, 77, "test-model", "private prompt", "")
|
||||
require.Error(t, err)
|
||||
require.Contains(t, rec.Body.String(), `"type":"error"`)
|
||||
}
|
||||
entries := logs.FilterMessage("account_test.failed").All()
|
||||
require.Len(t, entries, 1)
|
||||
fields := entries[0].ContextMap()
|
||||
require.Equal(t, int64(77), fields["account_id"])
|
||||
require.Equal(t, "req-test", fields["request_id"])
|
||||
require.Equal(t, "test-model", fields["requested_model"])
|
||||
require.NotEmpty(t, fields["test_id"])
|
||||
if background {
|
||||
require.Equal(t, "background", fields["test_source"])
|
||||
} else {
|
||||
require.Equal(t, "http", fields["test_source"])
|
||||
}
|
||||
if !missing {
|
||||
require.Equal(t, PlatformAnthropic, fields["platform"])
|
||||
}
|
||||
require.NotContains(t, fields, "prompt")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccountTestErrorLogRedaction(t *testing.T) {
|
||||
core, logs := observer.New(zap.WarnLevel)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("POST", "/test", nil).WithContext(logger.IntoContext(context.Background(), zap.New(core)))
|
||||
initAccountTestLogger(c, 1, "", "")
|
||||
message := `API returned 403: {"api_key":"secret-key","access_token":"secret-token","message":"denied"}`
|
||||
err := (&AccountTestService{}).sendErrorAndEnd(c, message)
|
||||
require.EqualError(t, err, message)
|
||||
entry := logs.FilterMessage("account_test.failed").All()[0]
|
||||
logged, ok := entry.ContextMap()["error"].(string)
|
||||
require.True(t, ok)
|
||||
require.NotContains(t, logged, "secret-key")
|
||||
require.NotContains(t, logged, "secret-token")
|
||||
require.Contains(t, logged, "denied")
|
||||
logAccountTestError(c, strings.Repeat("x", 10000))
|
||||
last, ok := logs.All()[1].ContextMap()["error"].(string)
|
||||
require.True(t, ok)
|
||||
require.Less(t, len(last), 2100)
|
||||
}
|
||||
@@ -359,6 +359,7 @@ func createTestPayload(modelID string) (map[string]any, error) {
|
||||
// mode is optional - "compact" routes OpenAI accounts to the /responses/compact probe path
|
||||
// opts is optional media (image/audio data URLs for real generation / STT).
|
||||
func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int64, modelID string, prompt string, mode string, opts ...AccountTestOptions) error {
|
||||
initAccountTestLogger(c, accountID, modelID, mode)
|
||||
ctx := c.Request.Context()
|
||||
testOpts := firstAccountTestOptions(opts)
|
||||
|
||||
@@ -367,6 +368,7 @@ func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, "Account not found")
|
||||
}
|
||||
bindAccountTestPlatform(c, account)
|
||||
|
||||
// Synthetic UI load-test accounts exercise the real SSE parsing and modal
|
||||
// interactions, but intentionally do not send their placeholder credentials
|
||||
@@ -383,7 +385,8 @@ func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int
|
||||
}
|
||||
|
||||
// Route to platform-specific test method
|
||||
if account.IsCNProvider() {
|
||||
// 按入站协议分流的多协议供应商(国产厂商等):按账号协议选测试路径。
|
||||
if account.RoutesProtocolByInbound() {
|
||||
switch account.GetAPIProtocol() {
|
||||
case APIProtocolAdaptive:
|
||||
return s.testCNProviderAdaptiveConnection(c, account, modelID, prompt)
|
||||
@@ -412,8 +415,9 @@ func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int
|
||||
return s.routeAntigravityTest(c, account, modelID, prompt)
|
||||
}
|
||||
|
||||
if account.IsOpenCodeGo() {
|
||||
return s.testOpenCodeGoAccountConnection(c, account, modelID, prompt)
|
||||
// 按模型分流的多模型聚合平台(OpenCode、Command Code 等)。
|
||||
if account.routesByModel() {
|
||||
return s.testModelRoutedAccountConnection(c, account, modelID, prompt)
|
||||
}
|
||||
|
||||
if account.IsTypeSafe() {
|
||||
@@ -423,36 +427,42 @@ func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int
|
||||
return s.testClaudeAccountConnection(c, account, modelID)
|
||||
}
|
||||
|
||||
// testOpenCodeGoAccountConnection probes the native endpoint for the selected
|
||||
// model. Adaptive accounts (the default) follow OpenCodeGoModelProtocol:
|
||||
// testModelRoutedAccountConnection probes the native endpoint for the selected
|
||||
// model on providers that route by model (see ProviderRoutingByModel). Adaptive
|
||||
// accounts (the default) follow the provider's protocol rules, e.g. OpenCode Go:
|
||||
// grok/gpt/muse-spark → Responses, minimax/qwen → Anthropic, everything else
|
||||
// (including deepseek-v4-flash) → Chat Completions. A pinned api_protocol
|
||||
// overrides that catalog. Falling through to the generic Claude tester used
|
||||
// credentials.base_url + /v1/messages?beta=true, which 404s as HTML on
|
||||
// https://opencode.ai/zen/go/v1/v1/messages.
|
||||
func (s *AccountTestService) testOpenCodeGoAccountConnection(c *gin.Context, account *Account, modelID string, prompt string) error {
|
||||
func (s *AccountTestService) testModelRoutedAccountConnection(c *gin.Context, account *Account, modelID string, prompt string) error {
|
||||
testModelID := strings.TrimSpace(modelID)
|
||||
if testModelID == "" {
|
||||
testModelID = DefaultOpenCodeGoTestModel
|
||||
testModelID = account.providerDefaultTestModel()
|
||||
}
|
||||
if testModelID == "" {
|
||||
testModelID = openai.DefaultTestModel
|
||||
}
|
||||
testModelID = account.GetMappedModel(testModelID)
|
||||
proto := account.GetAPIProtocol()
|
||||
switch proto {
|
||||
case APIProtocolChatCompletions, APIProtocolAnthropic, APIProtocolResponses:
|
||||
default:
|
||||
proto = openCodeGoNativeProtocol(account, testModelID)
|
||||
if account.IsOpenCodeGo() && IsOpenCodeUnsupportedModel(testModelID) {
|
||||
return fmt.Errorf("model %q is not supported on OpenCode standard gateway (gemini models require Google SDK endpoint, jev models require System One endpoint)", testModelID)
|
||||
}
|
||||
switch proto {
|
||||
// 与网关同一判定(含上游模型目录);测试没有入站协议,取模型的首选协议。
|
||||
protocol := account.resolveModelRoutedProtocol(testModelID)
|
||||
if s.openaiGatewayService != nil {
|
||||
protocol = s.openaiGatewayService.resolveUpstreamProtocolFor(c.Request.Context(), account, "", testModelID)
|
||||
}
|
||||
switch protocol {
|
||||
case APIProtocolAnthropic:
|
||||
return s.testCNProviderAnthropicConnection(c, account, testModelID)
|
||||
case APIProtocolResponses:
|
||||
return s.testOpenCodeGoResponsesConnection(c, account, testModelID)
|
||||
return s.testModelRoutedResponsesConnection(c, account, testModelID)
|
||||
default:
|
||||
return s.testCNProviderChatCompletionsConnection(c, account, testModelID, prompt)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AccountTestService) testOpenCodeGoResponsesConnection(c *gin.Context, account *Account, testModelID string) error {
|
||||
func (s *AccountTestService) testModelRoutedResponsesConnection(c *gin.Context, account *Account, testModelID string) error {
|
||||
authToken := strings.TrimSpace(account.GetOpenAIProtocolAPIKey())
|
||||
if authToken == "" {
|
||||
return s.sendErrorAndEnd(c, "No API key available")
|
||||
@@ -468,6 +478,9 @@ func (s *AccountTestService) testOpenCodeGoResponsesConnection(c *gin.Context, a
|
||||
|
||||
func (s *AccountTestService) testCNProviderChatCompletionsConnection(c *gin.Context, account *Account, modelID string, prompt string) error {
|
||||
testModelID := strings.TrimSpace(modelID)
|
||||
if testModelID == "" {
|
||||
testModelID = account.providerDefaultTestModel()
|
||||
}
|
||||
if testModelID == "" {
|
||||
testModelID = openai.DefaultTestModel
|
||||
}
|
||||
@@ -3271,7 +3284,7 @@ func (s *AccountTestService) sendEvent(c *gin.Context, event TestEvent) {
|
||||
|
||||
// sendErrorAndEnd sends an error event and ends the stream
|
||||
func (s *AccountTestService) sendErrorAndEnd(c *gin.Context, errorMsg string) error {
|
||||
log.Printf("Account test error: %s", errorMsg)
|
||||
logAccountTestError(c, errorMsg)
|
||||
s.sendEvent(c, TestEvent{Type: "error", Error: errorMsg})
|
||||
return fmt.Errorf("%s", errorMsg)
|
||||
}
|
||||
@@ -3284,6 +3297,7 @@ func (s *AccountTestService) RunTestBackground(ctx context.Context, accountID in
|
||||
w := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(w)
|
||||
ginCtx.Request = (&http.Request{}).WithContext(ctx)
|
||||
ginCtx.Set(accountTestBackgroundKey, true)
|
||||
|
||||
testErr := s.TestAccountConnection(ginCtx, accountID, modelID, "", AccountTestModeDefault)
|
||||
|
||||
|
||||
@@ -18,14 +18,19 @@ import (
|
||||
const accountTestSuppressCompletionContextKey = "account_test_suppress_completion"
|
||||
|
||||
// testCNProviderAdaptiveConnection verifies every native endpoint used by an
|
||||
// adaptive CN-provider account. Zhipu uses Chat Completions plus Anthropic;
|
||||
// DeepSeek and Kimi additionally use their native Responses endpoints.
|
||||
// adaptive account of a provider that routes by inbound protocol. Zhipu uses
|
||||
// Chat Completions plus Anthropic; DeepSeek and Kimi additionally use their
|
||||
// native Responses endpoints. Endpoints the provider profile does not offer are
|
||||
// skipped.
|
||||
func (s *AccountTestService) testCNProviderAdaptiveConnection(c *gin.Context, account *Account, modelID string, prompt string) error {
|
||||
testModelID := strings.TrimSpace(modelID)
|
||||
if testModelID == "" {
|
||||
testModelID = openai.DefaultTestModel
|
||||
requestedModelID := strings.TrimSpace(modelID)
|
||||
if requestedModelID == "" {
|
||||
requestedModelID = account.providerDefaultTestModel()
|
||||
}
|
||||
testModelID = account.GetMappedModel(testModelID)
|
||||
if requestedModelID == "" {
|
||||
requestedModelID = openai.DefaultTestModel
|
||||
}
|
||||
testModelID := account.GetMappedModel(requestedModelID)
|
||||
|
||||
authToken := strings.TrimSpace(account.GetOpenAIProtocolAPIKey())
|
||||
if authToken == "" {
|
||||
@@ -36,12 +41,14 @@ func (s *AccountTestService) testCNProviderAdaptiveConnection(c *gin.Context, ac
|
||||
// completion events until every native adaptive endpoint has passed.
|
||||
c.Set(accountTestSuppressCompletionContextKey, true)
|
||||
defer c.Set(accountTestSuppressCompletionContextKey, false)
|
||||
if err := s.testCNProviderChatCompletionsConnection(c, account, modelID, prompt); err != nil {
|
||||
if err := s.testCNProviderChatCompletionsConnection(c, account, requestedModelID, prompt); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := s.testCNProviderAdaptiveAnthropicConnection(c, account, testModelID, authToken); err != nil {
|
||||
return err
|
||||
if account.providerSupportsProtocol(APIProtocolAnthropic) {
|
||||
if err := s.testCNProviderAdaptiveAnthropicConnection(c, account, testModelID, authToken); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if account.SupportsNativeCNResponses() {
|
||||
@@ -237,10 +244,10 @@ func (s *AccountTestService) testCNProviderAnthropicConnection(c *gin.Context, a
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Anthropic base URL: %s", err.Error()))
|
||||
}
|
||||
if hint := cnAnthropicBaseURLMisconfigHint(baseURL); hint != "" {
|
||||
if hint := cnAnthropicBaseURLMisconfigHint(baseURL, account.routesByModel()); hint != "" {
|
||||
return s.sendErrorAndEnd(c, hint)
|
||||
}
|
||||
apiURL := strings.TrimRight(baseURL, "/") + "/v1/messages"
|
||||
apiURL := nativeAnthropicMessagesURL(account, baseURL)
|
||||
|
||||
c.Writer.Header().Set("Content-Type", "text/event-stream")
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
@@ -295,7 +302,9 @@ func (s *AccountTestService) testCNProviderAnthropicConnection(c *gin.Context, a
|
||||
// endpoint (paas path, version segment, or chat/completions / responses
|
||||
// suffix). The naive {base}/v1/messages join would 404 (e.g.
|
||||
// .../api/paas/v4/v1/messages) with no hint about the actual misconfiguration.
|
||||
func cnAnthropicBaseURLMisconfigHint(baseURL string) string {
|
||||
// versionAware is set for providers whose join is version-aware (see
|
||||
// nativeAnthropicMessagesURL), where a trailing version segment is valid.
|
||||
func cnAnthropicBaseURLMisconfigHint(baseURL string, versionAware bool) string {
|
||||
parsed, err := url.Parse(strings.TrimSpace(baseURL))
|
||||
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
|
||||
return ""
|
||||
@@ -307,7 +316,7 @@ func cnAnthropicBaseURLMisconfigHint(baseURL string) string {
|
||||
openAICompatShaped := strings.Contains(path, "/paas/") ||
|
||||
strings.HasSuffix(path, "/chat/completions") ||
|
||||
strings.HasSuffix(path, "/responses") ||
|
||||
openAIBaseURLHasVersionSuffix(path)
|
||||
(!versionAware && openAIBaseURLHasVersionSuffix(path))
|
||||
if !openAICompatShaped {
|
||||
return ""
|
||||
}
|
||||
|
||||
@@ -323,7 +323,7 @@ func (s *adminServiceImpl) DuplicateAccount(ctx context.Context, id int64, actor
|
||||
if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := NormalizeOpenCodeGoProtocolRulesCredentials(input.Credentials); err != nil {
|
||||
if err := NormalizeProtocolRulesCredentials(input.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
duplicate, err := buildAccountForCreate(input, accountExtra)
|
||||
@@ -522,7 +522,7 @@ func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccou
|
||||
if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := NormalizeOpenCodeGoProtocolRulesCredentials(input.Credentials); err != nil {
|
||||
if err := NormalizeProtocolRulesCredentials(input.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Never persist ephemeral SSO/password secrets after OAuth conversion.
|
||||
@@ -655,7 +655,7 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
|
||||
if err := NormalizeHeaderOverrideCredentials(account.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := NormalizeOpenCodeGoProtocolRulesCredentials(account.Credentials); err != nil {
|
||||
if err := NormalizeProtocolRulesCredentials(account.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Strip SSO/password residue that must never sit next to OAuth tokens.
|
||||
@@ -1102,7 +1102,7 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp
|
||||
if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := NormalizeOpenCodeGoProtocolRulesCredentials(input.Credentials); err != nil {
|
||||
if err := NormalizeProtocolRulesCredentials(input.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Bulk may mix platforms; always drop ephemeral SSO/password keys (cookie
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
@@ -319,10 +320,13 @@ func defaultAllowImageGenerationForPlatform(platform string) bool {
|
||||
func compositeDefaultModelsListCandidateIDs() []string {
|
||||
seen := make(map[string]struct{})
|
||||
ids := make([]string, 0)
|
||||
// TypeSafe stays out of the static composite candidates (jev-latest only works
|
||||
// through /v1/systemone); groups with TypeSafe accounts still get it from the
|
||||
// account model mappings collected by GetGroupModelsListCandidates.
|
||||
for _, platform := range []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo} {
|
||||
for _, platform := range domain.CompositePrecedencePlatformIDs() {
|
||||
// TypeSafe stays out of the static composite candidates (jev-latest only works
|
||||
// through /v1/systemone); groups with TypeSafe accounts still get it from the
|
||||
// account model mappings collected by GetGroupModelsListCandidates.
|
||||
if platform == PlatformTypeSafe {
|
||||
continue
|
||||
}
|
||||
for _, id := range defaultModelsListCandidateIDs(platform) {
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
|
||||
@@ -21,7 +21,7 @@ func (a *Account) GetAnthropicAPIKeyAuthScheme() string {
|
||||
if a == nil || a.Type != AccountTypeAPIKey {
|
||||
return AnthropicAPIKeyAuthSchemeXAPIKey
|
||||
}
|
||||
if a.Platform != PlatformAnthropic && !a.IsCNProvider() {
|
||||
if a.Platform != PlatformAnthropic && !a.RoutesProtocolByInbound() {
|
||||
return AnthropicAPIKeyAuthSchemeXAPIKey
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// anthropicSSEWithNegativeBlockIndex is a buffered Anthropic stream in which a
|
||||
// malformed upstream sends content_block_delta with a negative index. The
|
||||
// buffered aggregators only bounded the index from above, so indexing
|
||||
// finalResp.Content with it panicked.
|
||||
func anthropicSSEWithNegativeBlockIndex() string {
|
||||
return strings.Join([]string{
|
||||
`event: message_start`,
|
||||
`data: {"type":"message_start","message":{"id":"msg_neg","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4.5","usage":{"input_tokens":10}}}`,
|
||||
``,
|
||||
`event: content_block_start`,
|
||||
`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
|
||||
``,
|
||||
`event: content_block_delta`,
|
||||
`data: {"type":"content_block_delta","index":-1,"delta":{"type":"text_delta","text":"IGNORED"}}`,
|
||||
``,
|
||||
`event: content_block_delta`,
|
||||
`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}`,
|
||||
``,
|
||||
`event: message_delta`,
|
||||
`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5}}`,
|
||||
``,
|
||||
`event: message_stop`,
|
||||
`data: {"type":"message_stop"}`,
|
||||
``,
|
||||
}, "\n")
|
||||
}
|
||||
|
||||
func TestAnthropicBufferedAggregators_IgnoreNegativeBlockIndex(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
run func(resp *http.Response, c *gin.Context) error
|
||||
}{
|
||||
{
|
||||
name: "gateway chat completions",
|
||||
run: func(resp *http.Response, c *gin.Context) error {
|
||||
_, err := (&GatewayService{}).handleCCBufferedFromAnthropic(resp, c, "gpt-5", "claude-sonnet-4.5", nil, time.Now())
|
||||
return err
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "gateway responses",
|
||||
run: func(resp *http.Response, c *gin.Context) error {
|
||||
_, err := (&GatewayService{}).handleResponsesBufferedStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now(), apicompat.ResponsesClientToolMapping{})
|
||||
return err
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "openai native anthropic chat completions",
|
||||
run: func(resp *http.Response, c *gin.Context) error {
|
||||
_, err := newNativeAnthropicHangTestService(5).handleCCBufferedFromNativeAnthropic(resp, c, "glm-4.7", "glm-4.7", "glm-4.7", nil, time.Now())
|
||||
return err
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "openai native anthropic responses",
|
||||
run: func(resp *http.Response, c *gin.Context) error {
|
||||
_, err := newNativeAnthropicHangTestService(5).handleResponsesBufferedFromNativeAnthropic(resp, c, "glm-4.7", "glm-4.7", "glm-4.7", nil, time.Now(), apicompat.ResponsesClientToolMapping{})
|
||||
return err
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(anthropicSSEWithNegativeBlockIndex())),
|
||||
}
|
||||
|
||||
var err error
|
||||
require.NotPanics(t, func() { err = tt.run(resp, c) })
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, rec.Body.String(), "Hello")
|
||||
require.NotContains(t, rec.Body.String(), "IGNORED")
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1937,9 +1937,9 @@ func (s *AuthService) snapshotPlatformQuotaDefaults(ctx context.Context, userID
|
||||
return nil
|
||||
}
|
||||
// 平台配额快照是 best-effort(fail-open):必须脱离调用方事务执行。
|
||||
// 否则某平台违反 user_platform_quotas 的 CHECK 约束(如尚未进约束的新平台)会让
|
||||
// 整个调用方事务被 Postgres 标记 aborted,把"无关紧要的默认配额快照"放大成
|
||||
// "整笔注册失败"(OAuth pending 路径曾因此 500 → 清 cookie → 404)。
|
||||
// 否则写入时的任何数据库错误(历史上是新平台尚未加入 user_platform_quotas 的
|
||||
// CHECK 约束)都会让整个调用方事务被 Postgres 标记 aborted,把"无关紧要的默认
|
||||
// 配额快照"放大成"整笔注册失败"(OAuth pending 路径曾因此 500 → 清 cookie → 404)。
|
||||
ctx = dbent.WithoutTx(ctx)
|
||||
// 仅为至少配置了一档限额的平台建行:user_platform_quotas 中不存在的行等价于不限额,
|
||||
// 三档全空的记录不携带任何可执行的限额。
|
||||
|
||||
@@ -297,7 +297,7 @@ func callProvider(ctx context.Context, provider, endpoint, apiKey, model, prompt
|
||||
return "", "", 0, err
|
||||
}
|
||||
headers := mergeHeaders(adapter.buildHeaders(apiKey), opts)
|
||||
full := joinURL(endpoint, adapter.buildPath(model))
|
||||
full := joinURL(endpoint, monitorRequestPath(provider, endpoint, adapter, model))
|
||||
respBytes, status, err := postRawJSON(ctx, full, body, headers)
|
||||
if err != nil {
|
||||
return "", "", status, err
|
||||
@@ -556,6 +556,37 @@ func postRawJSON(ctx context.Context, fullURL string, payload []byte, headers ma
|
||||
return respBody, resp.StatusCode, nil
|
||||
}
|
||||
|
||||
// monitorRequestPath 返回探测请求路径。智谱按 endpoint 区分:
|
||||
// - 已带 /paas/v4(含 Coding Plan 的 /api/coding/paas/v4):只追加 /chat/completions
|
||||
// - 官方域名根地址:/api/paas/v4/chat/completions
|
||||
// - 中转站 / 本站网关:只暴露 OpenAI 兼容的 /v1/chat/completions,与 Kimi / DeepSeek 一致
|
||||
func monitorRequestPath(provider, endpoint string, adapter providerAdapter, model string) string {
|
||||
if provider != MonitorProviderZhipu {
|
||||
return adapter.buildPath(model)
|
||||
}
|
||||
u, err := url.Parse(strings.TrimSpace(endpoint))
|
||||
if err != nil {
|
||||
return providerOpenAIPath
|
||||
}
|
||||
if strings.Contains(u.EscapedPath(), "/paas/v4") {
|
||||
return "/chat/completions"
|
||||
}
|
||||
if isZhipuOfficialHost(u) {
|
||||
return adapter.buildPath(model)
|
||||
}
|
||||
return providerOpenAIPath
|
||||
}
|
||||
|
||||
func isZhipuOfficialHost(u *url.URL) bool {
|
||||
host := strings.ToLower(u.Hostname())
|
||||
for _, official := range []string{"bigmodel.cn", "z.ai"} {
|
||||
if host == official || strings.HasSuffix(host, "."+official) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// joinURL 保留 base 的上游路径前缀,并避免重复追加已有的 API 路径前缀。
|
||||
// 使用 EscapedPath 匹配完整路径段,避免把 hostname 或编码斜杠当作路径。
|
||||
func joinURL(base, path string) string {
|
||||
|
||||
@@ -84,6 +84,9 @@ func TestCallProvider_BasePath(t *testing.T) {
|
||||
{"openai responses version", MonitorProviderOpenAI, MonitorAPIModeResponses, "/relay/v1", "/relay/v1/responses"},
|
||||
{"gemini version", MonitorProviderGemini, "", "/relay/v1beta", "/relay/v1beta/models/test-model:generateContent"},
|
||||
{"zhipu version", MonitorProviderZhipu, "", "/api/paas/v4", "/api/paas/v4/chat/completions"},
|
||||
{"zhipu coding plan", MonitorProviderZhipu, "", "/api/coding/paas/v4", "/api/coding/paas/v4/chat/completions"},
|
||||
{"zhipu relay origin", MonitorProviderZhipu, "", "", "/v1/chat/completions"},
|
||||
{"zhipu relay version", MonitorProviderZhipu, "", "/relay/v1", "/relay/v1/chat/completions"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, _, status, err := callProvider(context.Background(), tc.provider,
|
||||
@@ -94,3 +97,19 @@ func TestCallProvider_BasePath(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMonitorRequestPath_Zhipu(t *testing.T) {
|
||||
adapter := providerAdapters[MonitorProviderZhipu]
|
||||
for endpoint, want := range map[string]string{
|
||||
"https://open.bigmodel.cn": providerZhipuPath,
|
||||
"https://api.z.ai": providerZhipuPath,
|
||||
"https://open.bigmodel.cn/api/coding/paas/v4": "/chat/completions",
|
||||
"https://birdapi.up.railway.app": providerOpenAIPath,
|
||||
"https://notbigmodel.cn": providerOpenAIPath,
|
||||
"https://relay.example.com/v1": providerOpenAIPath,
|
||||
} {
|
||||
require.Equal(t, want, monitorRequestPath(MonitorProviderZhipu, endpoint, adapter, "glm-5.3-flash"), endpoint)
|
||||
}
|
||||
require.Equal(t, providerOpenAIPath,
|
||||
monitorRequestPath(MonitorProviderDeepseek, "https://birdapi.up.railway.app", providerAdapters[MonitorProviderDeepseek], "m"))
|
||||
}
|
||||
|
||||
@@ -963,6 +963,9 @@ func ChannelMonitorV2HealthForWithThresholds(metrics ChannelMonitorV2Metric, thr
|
||||
Overall: "unknown", ErrorRate: "unknown", TTFT: "unknown", Cache: "unknown",
|
||||
MinimumSample: thresholds.MinimumSample, Thresholds: thresholds,
|
||||
}
|
||||
if metrics.RequestCount < result.MinimumSample {
|
||||
return result
|
||||
}
|
||||
|
||||
type scored struct {
|
||||
score float64
|
||||
|
||||
@@ -252,6 +252,81 @@ func TestChannelMonitorV2HealthBlendsErrorTTFTAndCache(t *testing.T) {
|
||||
require.Nil(t, health.Score)
|
||||
}
|
||||
|
||||
func TestChannelMonitorV2HealthRequiresRequestSamples(t *testing.T) {
|
||||
for _, requestCount := range []int64{0, 8, 9, 49} {
|
||||
t.Run(fmt.Sprintf("requests_%d", requestCount), func(t *testing.T) {
|
||||
thresholds := DefaultChannelMonitorV2HealthThresholds()
|
||||
thresholds.WarningCacheRate = 0.85
|
||||
thresholds.CriticalCacheRate = 0.60
|
||||
metrics := ChannelMonitorV2Metric{SuccessRequests: requestCount, RequestCount: requestCount}
|
||||
if requestCount > 0 {
|
||||
metrics.CacheRateDenominator = 448
|
||||
}
|
||||
health := ChannelMonitorV2HealthForWithThresholds(metrics, thresholds)
|
||||
require.Equal(t, ChannelMonitorV2Health{
|
||||
Overall: "unknown", ErrorRate: "unknown", TTFT: "unknown", Cache: "unknown",
|
||||
MinimumSample: 50, Thresholds: NormalizeChannelMonitorV2HealthThresholds(thresholds),
|
||||
}, health)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChannelMonitorV2HealthMinimumSampleBoundaries(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
minimumSample int64
|
||||
wantMinimum int64
|
||||
}{
|
||||
{name: "configured", minimumSample: 20, wantMinimum: 20},
|
||||
{name: "configured_one", minimumSample: 1, wantMinimum: 1},
|
||||
{name: "normalized_zero", minimumSample: 0, wantMinimum: 50},
|
||||
{name: "normalized_negative", minimumSample: -1, wantMinimum: 50},
|
||||
{name: "normalized_maximum", minimumSample: 10001, wantMinimum: 10000},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
thresholds := DefaultChannelMonitorV2HealthThresholds()
|
||||
thresholds.MinimumSample = test.minimumSample
|
||||
thresholds.WarningCacheRate = 0.85
|
||||
thresholds.CriticalCacheRate = 0.60
|
||||
p50 := int64(1000)
|
||||
metrics := ChannelMonitorV2Metric{
|
||||
SuccessRequests: test.wantMinimum - 1, RequestCount: test.wantMinimum - 1,
|
||||
CacheRateDenominator: 20000,
|
||||
TTFT: ChannelMonitorV2Latency{SampleCount: 20000, P50Ms: &p50},
|
||||
}
|
||||
health := ChannelMonitorV2HealthForWithThresholds(metrics, thresholds)
|
||||
require.Equal(t, ChannelMonitorV2Health{
|
||||
Overall: "unknown", ErrorRate: "unknown", TTFT: "unknown", Cache: "unknown",
|
||||
MinimumSample: test.wantMinimum, Thresholds: NormalizeChannelMonitorV2HealthThresholds(thresholds),
|
||||
}, health)
|
||||
|
||||
metrics.RequestCount = test.wantMinimum
|
||||
metrics.SuccessRequests = test.wantMinimum
|
||||
health = ChannelMonitorV2HealthForWithThresholds(metrics, thresholds)
|
||||
require.Equal(t, test.wantMinimum, health.MinimumSample)
|
||||
require.Equal(t, "healthy", health.ErrorRate)
|
||||
require.Equal(t, "healthy", health.TTFT)
|
||||
require.Equal(t, "critical", health.Cache)
|
||||
require.NotNil(t, health.ErrorRateScore)
|
||||
require.NotNil(t, health.TTFTScore)
|
||||
require.NotNil(t, health.CacheScore)
|
||||
require.NotNil(t, health.Score)
|
||||
require.InDelta(t, 80.0, *health.Score, 0.01)
|
||||
require.Equal(t, "healthy", health.Overall)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChannelMonitorV2HealthSufficientHighErrorIsCritical(t *testing.T) {
|
||||
health := ChannelMonitorV2HealthFor(ChannelMonitorV2Metric{
|
||||
RequestCount: 50, ErrorRequests: 50, ErrorRate: 1,
|
||||
})
|
||||
require.Equal(t, "critical", health.ErrorRate)
|
||||
require.Equal(t, "critical", health.Overall)
|
||||
require.NotNil(t, health.Score)
|
||||
require.Zero(t, *health.Score)
|
||||
}
|
||||
|
||||
func TestChannelMonitorV2HealthLeavesMissingTTFTUnknown(t *testing.T) {
|
||||
health := ChannelMonitorV2HealthFor(ChannelMonitorV2Metric{
|
||||
RequestCount: 200,
|
||||
@@ -452,6 +527,9 @@ func TestSnapshotRedactsPublicConfigPolicyFields(t *testing.T) {
|
||||
Metrics: ChannelMonitorV2Metric{RequestCount: 100, ErrorRate: 0.1, SuccessRate: 0.9, RPM: 5},
|
||||
},
|
||||
}
|
||||
repo.snap.Health = ChannelMonitorV2HealthFor(repo.snap.Metrics)
|
||||
wantHealth := repo.snap.Health
|
||||
require.NotNil(t, wantHealth.Score)
|
||||
svc := NewChannelMonitorV2Service(repo)
|
||||
snap, err := svc.Snapshot(context.Background(), ChannelMonitorV2Filter{}, false)
|
||||
require.NoError(t, err)
|
||||
@@ -461,6 +539,7 @@ func TestSnapshotRedactsPublicConfigPolicyFields(t *testing.T) {
|
||||
require.Empty(t, snap.Config.Platforms[0].Models)
|
||||
require.Equal(t, 300, snap.Config.RefreshIntervalSeconds)
|
||||
require.Zero(t, snap.Metrics.RequestCount)
|
||||
require.Equal(t, wantHealth, snap.Health)
|
||||
require.InDelta(t, 0.1, snap.Metrics.ErrorRate, 0.0001)
|
||||
require.Zero(t, snap.Metrics.RPM)
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
@@ -368,7 +369,7 @@ func isPlatformPricingMatch(groupPlatform, pricingPlatform string) bool {
|
||||
// fallback used before a request target has been resolved.
|
||||
func matchingPlatforms(groupPlatform string) []string {
|
||||
if groupPlatform == PlatformComposite {
|
||||
return []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe}
|
||||
return domain.CompositePrecedencePlatformIDs()
|
||||
}
|
||||
return []string{groupPlatform}
|
||||
}
|
||||
|
||||
@@ -2127,7 +2127,7 @@ func TestMatchingPlatforms(t *testing.T) {
|
||||
{"anthropic returns itself", PlatformAnthropic, []string{PlatformAnthropic}},
|
||||
{"gemini returns itself", PlatformGemini, []string{PlatformGemini}},
|
||||
{"openai returns itself", PlatformOpenAI, []string{PlatformOpenAI}},
|
||||
{"composite returns concrete platforms", PlatformComposite, []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe}},
|
||||
{"composite returns concrete platforms", PlatformComposite, []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe, PlatformCommandCode, PlatformCline}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Cline API(https://docs.cline.bot/api/overview):OpenAI 兼容的多模型聚合端点,
|
||||
// 基址 https://api.cline.bot/api/v1,只提供 Chat Completions(Cline 自己的客户端同样
|
||||
// 只走该端点),模型名带厂商前缀(如 anthropic/claude-sonnet-4-6)。Responses 与
|
||||
// Anthropic 入站由网关转换为 Chat Completions。
|
||||
//
|
||||
// 同一个 API Key 下有两个互不影响的额度来源(下称"钱包"),由模型名决定扣哪一个:
|
||||
// - 积分:cline-pass/* 与 cline-free/* 以外的模型按量扣账户积分,余额可查;
|
||||
// - ClinePass 订阅:cline-pass/* 模型,受 5 小时 / 7 天 / 30 天三档滚动用量上限限制,
|
||||
// 各档已用百分比与重置时间可查(/api/v1/users/me/plan/usage-limits);
|
||||
// - cline-free/* 免费模型不扣任何一边,只有逐模型的每日额度。
|
||||
//
|
||||
// 一边用完只影响该边的模型,所以超限冷却按钱包落在模型级冷却键上(见
|
||||
// clineWalletRateLimitKey),不停整个账号,也不需要管理员声明账号类型。
|
||||
|
||||
const (
|
||||
// DefaultClineTestModel 是连接测试的默认模型(按量积分,价格低、长期在售)。
|
||||
DefaultClineTestModel = "deepseek/deepseek-v4-flash"
|
||||
// DefaultClinePassTestModel 是订阅了 ClinePass 的账号连接测试的默认模型。
|
||||
DefaultClinePassTestModel = "cline-pass/glm-5.3-flash"
|
||||
|
||||
clineAPIHost = "api.cline.bot"
|
||||
|
||||
clinePassModelPrefix = "cline-pass/"
|
||||
clineFreeModelPrefix = "cline-free/"
|
||||
|
||||
// clinePassRateLimitKey / clineCreditsRateLimitKey 是两个钱包的模型级冷却键:
|
||||
// 命中后该钱包的全部模型都不再调度到该账号,另一边不受影响。
|
||||
clinePassRateLimitKey = "cline:pass"
|
||||
clineCreditsRateLimitKey = "cline:credits"
|
||||
)
|
||||
|
||||
// IsCline 报告账号是否为 Cline 平台账号。
|
||||
func (a *Account) IsCline() bool {
|
||||
return a != nil && a.Platform == PlatformCline
|
||||
}
|
||||
|
||||
// clineWalletRateLimitKey 返回模型所属钱包的冷却键;免费模型不属于任何钱包,返回空串。
|
||||
func clineWalletRateLimitKey(model string) string {
|
||||
model = strings.ToLower(strings.TrimSpace(model))
|
||||
switch {
|
||||
case model == "":
|
||||
return ""
|
||||
case strings.HasPrefix(model, clinePassModelPrefix):
|
||||
return clinePassRateLimitKey
|
||||
case strings.HasPrefix(model, clineFreeModelPrefix):
|
||||
return ""
|
||||
default:
|
||||
return clineCreditsRateLimitKey
|
||||
}
|
||||
}
|
||||
|
||||
// isOfficialClineHost 报告 URL 是否指向官方 Cline API 主机。
|
||||
func isOfficialClineHost(target string) bool {
|
||||
parsed, err := url.Parse(strings.TrimSpace(target))
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return strings.EqualFold(parsed.Scheme, "https") && strings.EqualFold(parsed.Hostname(), clineAPIHost)
|
||||
}
|
||||
|
||||
// clineAccountAPISupported 报告账号能否查询积分余额与 ClinePass 用量:API Key 账号且
|
||||
// 推理端点指向官方主机。自定义中转的 Key 不发往官方账号接口。
|
||||
func (a *Account) clineAccountAPISupported() bool {
|
||||
return a.IsCline() && a.Type == AccountTypeAPIKey && isOfficialClineHost(a.GetOpenAIBaseURL())
|
||||
}
|
||||
|
||||
// clinePassSubscribed 报告最近一次用量探测是否显示账号订阅了 ClinePass(有窗口快照)。
|
||||
func (a *Account) clinePassSubscribed() bool {
|
||||
if a == nil || a.Extra == nil {
|
||||
return false
|
||||
}
|
||||
for _, suffix := range []string{cnExtraSuffix5hUsed, cnExtraSuffixWeeklyUsed, cnExtraSuffixMonthlyUsed} {
|
||||
if value, ok := a.Extra[cnExtraKey(PlatformCline, suffix)]; ok && value != nil {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,648 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// clineTestAccount 只带 api_key:端点与测试模型全部来自 provider profile。
|
||||
func clineTestAccount(id int64) *Account {
|
||||
return &Account{
|
||||
ID: id,
|
||||
Name: "cline",
|
||||
Platform: PlatformCline,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{"api_key": "cline_test_key", "api_protocol": APIProtocolAdaptive},
|
||||
}
|
||||
}
|
||||
|
||||
// withClinePassSnapshot 模拟用量探测写入的 ClinePass 窗口快照。
|
||||
func withClinePassSnapshot(account *Account, extra map[string]any) *Account {
|
||||
account.Extra = map[string]any{"cline_5h_used_percent": 10.0, "cline_weekly_used_percent": 20.0, "cline_monthly_used_percent": 30.0}
|
||||
for key, value := range extra {
|
||||
account.Extra[key] = value
|
||||
}
|
||||
return account
|
||||
}
|
||||
|
||||
func TestClineProfileSingleModeChatCompletionsOnly(t *testing.T) {
|
||||
profile := LookupProviderProfile(PlatformCline)
|
||||
require.NotNil(t, profile)
|
||||
require.Len(t, profile.Modes, 1)
|
||||
|
||||
account := clineTestAccount(1)
|
||||
require.True(t, account.IsMultiProtocolAPIKey())
|
||||
require.True(t, account.RoutesProtocolByInbound())
|
||||
require.False(t, account.routesByModel())
|
||||
require.Equal(t, APIProtocolAdaptive, account.GetAPIProtocol())
|
||||
require.Equal(t, DefaultClineBaseURL, account.GetOpenAIBaseURL())
|
||||
require.False(t, account.SupportsNativeCNResponses())
|
||||
require.False(t, account.providerSupportsProtocol(APIProtocolAnthropic))
|
||||
for _, inbound := range []string{APIProtocolChatCompletions, APIProtocolResponses, APIProtocolAnthropic} {
|
||||
require.Equal(t, APIProtocolChatCompletions, resolveUpstreamProtocol(account, inbound, "anthropic/claude-sonnet-4-6", nil), inbound)
|
||||
}
|
||||
|
||||
require.False(t, account.clinePassSubscribed())
|
||||
require.Equal(t, DefaultClineTestModel, account.providerDefaultTestModel())
|
||||
subscribed := withClinePassSnapshot(clineTestAccount(2), nil)
|
||||
require.True(t, subscribed.clinePassSubscribed())
|
||||
require.Equal(t, DefaultClinePassTestModel, subscribed.providerDefaultTestModel())
|
||||
// 取消订阅后探测把窗口写成 nil。
|
||||
cancelled := withClinePassSnapshot(clineTestAccount(3), map[string]any{
|
||||
"cline_5h_used_percent": nil, "cline_weekly_used_percent": nil, "cline_monthly_used_percent": nil,
|
||||
})
|
||||
require.False(t, cancelled.clinePassSubscribed())
|
||||
}
|
||||
|
||||
// 经真实入口:三种入站都转换 / 直通到 Chat Completions,模型名带厂商前缀原样发往上游。
|
||||
func TestClineGatewayForwardsToChatCompletions(t *testing.T) {
|
||||
for _, ingress := range routingMatrixIngresses() {
|
||||
upstream := &httpUpstreamRecorder{err: errors.New("stop after capture")}
|
||||
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
|
||||
body := routingMatrixCase{ingress: ingress, model: "anthropic/claude-sonnet-4-6"}.body()
|
||||
_ = ingress.forward(svc, adaptiveProtocolTestContext(ingress.path, body), clineTestAccount(3), body)
|
||||
|
||||
require.NotEmpty(t, upstream.requests, ingress.name)
|
||||
require.Equal(t, "https://api.cline.bot/api/v1/chat/completions", upstream.requests[len(upstream.requests)-1].URL.String(), ingress.name)
|
||||
require.Equal(t, "Bearer cline_test_key", upstream.requests[len(upstream.requests)-1].Header.Get("Authorization"), ingress.name)
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, "messages").Exists(), ingress.name)
|
||||
require.Equal(t, "anthropic/claude-sonnet-4-6", gjson.GetBytes(upstream.lastBody, "model").String(), ingress.name)
|
||||
}
|
||||
}
|
||||
|
||||
// 连接测试未指定模型时:订阅了 ClinePass 用订阅模型(不消耗积分),否则用按量模型;
|
||||
// ClinePass 钱包冷却中也用按量模型;自适应与固定 Chat Completions 一致。
|
||||
func TestAccountTestService_ClineDefaultTestModel(t *testing.T) {
|
||||
cases := []struct {
|
||||
subscribed bool
|
||||
passCooling bool
|
||||
apiProtocol string
|
||||
want string
|
||||
}{
|
||||
{false, false, "", DefaultClineTestModel},
|
||||
{true, false, "", DefaultClinePassTestModel},
|
||||
{false, false, APIProtocolChatCompletions, DefaultClineTestModel},
|
||||
{true, false, APIProtocolChatCompletions, DefaultClinePassTestModel},
|
||||
{true, true, "", DefaultClineTestModel},
|
||||
}
|
||||
for i, tc := range cases {
|
||||
account := clineTestAccount(int64(600 + i))
|
||||
if tc.subscribed {
|
||||
withClinePassSnapshot(account, nil)
|
||||
}
|
||||
if tc.passCooling {
|
||||
now := time.Now()
|
||||
setAccountModelRateLimitSnapshot(account, clinePassRateLimitKey, now.Add(time.Hour), clinePassLimitReason, now)
|
||||
}
|
||||
if tc.apiProtocol != "" {
|
||||
account.Credentials["api_protocol"] = tc.apiProtocol
|
||||
}
|
||||
svc, upstream := adaptiveCNAccountTestService(account, adaptiveCNChatTestResponse())
|
||||
c, recorder := newTestContext()
|
||||
|
||||
err := svc.TestAccountConnection(c, account.ID, "", "hi", AccountTestModeDefault)
|
||||
|
||||
require.NoError(t, err, tc)
|
||||
require.Len(t, upstream.requests, 1, tc)
|
||||
require.Equal(t, "https://api.cline.bot/api/v1/chat/completions", upstream.requests[0].URL.String(), tc)
|
||||
require.Equal(t, tc.want, gjson.GetBytes(upstream.lastBody, "model").String(), tc)
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`, tc)
|
||||
}
|
||||
}
|
||||
|
||||
// 响应体样例取自 Cline 客户端的判定文案与官方错误文档。
|
||||
const (
|
||||
clineInsufficientCreditsBody = `{"error":{"code":"insufficient_credits","message":"Insufficient credits","details":{"current_balance":0}}}`
|
||||
clineSpendLimitBody = `{"error":{"code":"SPEND_LIMIT_EXCEEDED","message":"Organization spend limit exceeded"}}`
|
||||
clineInferenceCapBody = `{"error":{"code":"INFERENCE_CAP_ERROR","message":"Inference cap reached"}}`
|
||||
clinePass5hBody = `{"error":{"code":429,"message":"You have reached your 5-hour ClinePass limit. Please try again later."}}`
|
||||
clinePassWeeklyBody = `{"error":{"code":429,"message":"You have reached your weekly ClinePass limit. Please try again later."}}`
|
||||
clinePassMonthlyBody = `{"error":{"code":429,"message":"You have reached your monthly ClinePass limit. Please try again later."}}`
|
||||
clineNotSubscribedBody = `{"error":{"code":403,"message":"The user is not subscribed to required model plan"}}`
|
||||
clineOrgPassBody = `{"error":{"code":403,"message":"Organization accounts cannot use individual model inference subscriptions"}}`
|
||||
clineFreeModelLimitBody = `{"error":{"code":429,"message":"Free limit reached on model cline-free/kimi-k3. Try again in 3h 20m"}}`
|
||||
clineRateLimitBody = `{"error":{"code":429,"message":"Rate limit exceeded"}}`
|
||||
)
|
||||
|
||||
func TestClassifyClineError(t *testing.T) {
|
||||
cases := map[string]clineErrorKind{
|
||||
clineInsufficientCreditsBody: clineCreditsExhausted,
|
||||
clineSpendLimitBody: clineSpendLimit,
|
||||
clineInferenceCapBody: clineInferenceCap,
|
||||
clinePass5hBody: clinePassLimit,
|
||||
clinePassWeeklyBody: clinePassLimit,
|
||||
clineNotSubscribedBody: clinePassUnavailable,
|
||||
clineOrgPassBody: clinePassUnavailable,
|
||||
clineFreeModelLimitBody: clineFreeModelLimit,
|
||||
clineRateLimitBody: clineErrorNone,
|
||||
`{"error":{"message":"You have reached your limit. Please try again later."}}`: clineErrorNone,
|
||||
}
|
||||
for body, want := range cases {
|
||||
kind, _ := parseClineError([]byte(body))
|
||||
require.Equal(t, want, kind, body)
|
||||
}
|
||||
|
||||
require.Equal(t, "5h", clinePassWindow([]byte(clinePass5hBody)))
|
||||
require.Equal(t, "weekly", clinePassWindow([]byte(clinePassWeeklyBody)))
|
||||
require.Equal(t, "monthly", clinePassWindow([]byte(clinePassMonthlyBody)))
|
||||
require.Equal(t, 3*time.Hour+20*time.Minute, clineFreeModelResetAfter([]byte(clineFreeModelLimitBody)))
|
||||
require.Zero(t, clineFreeModelResetAfter([]byte(`Free limit reached on model x. Try again in 3 weeks`)))
|
||||
}
|
||||
|
||||
// ClinePass 超限冷却到快照中该窗口的重置时间;没有可用快照时按窗口估计。
|
||||
func TestClinePassLimitCooldown(t *testing.T) {
|
||||
now := time.Now()
|
||||
require.Equal(t, clineLimitRecheck, clinePassLimitCooldown(clineTestAccount(1), "5h", now))
|
||||
require.Equal(t, clinePassWeeklyRecheck, clinePassLimitCooldown(clineTestAccount(1), "weekly", now))
|
||||
require.Equal(t, clinePassMonthlyRecheck, clinePassLimitCooldown(clineTestAccount(1), "monthly", now))
|
||||
|
||||
account := withClinePassSnapshot(clineTestAccount(2), map[string]any{
|
||||
"cline_weekly_reset_at": now.Add(50 * time.Hour).UTC().Format(time.RFC3339),
|
||||
"cline_5h_reset_at": now.Add(-time.Minute).UTC().Format(time.RFC3339),
|
||||
})
|
||||
require.InDelta(t, (50 * time.Hour).Seconds(), clinePassLimitCooldown(account, "weekly", now).Seconds(), 1)
|
||||
require.Equal(t, clineLimitRecheck, clinePassLimitCooldown(account, "5h", now), "past reset falls back to the estimate")
|
||||
}
|
||||
|
||||
// cline-pass/* 与其余付费模型各自落在钱包冷却键上,免费模型不属于任何钱包;调度按
|
||||
// 请求模型检查所属钱包。
|
||||
func TestClineWalletRateLimitKeys(t *testing.T) {
|
||||
require.Equal(t, clinePassRateLimitKey, clineWalletRateLimitKey("cline-pass/glm-5.3"))
|
||||
require.Equal(t, clinePassRateLimitKey, clineWalletRateLimitKey(" Cline-Pass/GLM-5.3 "))
|
||||
require.Equal(t, clineCreditsRateLimitKey, clineWalletRateLimitKey("deepseek/deepseek-v4-flash"))
|
||||
require.Empty(t, clineWalletRateLimitKey("cline-free/kimi-k3"))
|
||||
require.Empty(t, clineWalletRateLimitKey(""))
|
||||
|
||||
account := clineTestAccount(1)
|
||||
setAccountModelRateLimitSnapshot(account, clinePassRateLimitKey, time.Now().Add(time.Hour), clinePassLimitReason, time.Now())
|
||||
ctx := context.Background()
|
||||
require.True(t, account.isModelRateLimitedWithContext(ctx, "cline-pass/glm-5.3"))
|
||||
require.False(t, account.isModelRateLimitedWithContext(ctx, "deepseek/deepseek-v4-flash"))
|
||||
require.False(t, account.isModelRateLimitedWithContext(ctx, "cline-free/kimi-k3"))
|
||||
|
||||
setAccountModelRateLimitSnapshot(account, clineCreditsRateLimitKey, time.Now().Add(time.Hour), clineCreditsReason, time.Now())
|
||||
require.True(t, account.isModelRateLimitedWithContext(ctx, "deepseek/deepseek-v4-flash"))
|
||||
require.False(t, account.isModelRateLimitedWithContext(ctx, "cline-free/kimi-k3"))
|
||||
require.Equal(t, clineCreditsReason, account.modelRateLimitReason(clineCreditsRateLimitKey))
|
||||
|
||||
// 其他平台不受 Cline 钱包键影响。
|
||||
other := &Account{Platform: PlatformDeepseek, Extra: account.Extra}
|
||||
require.False(t, other.isModelRateLimitedWithContext(ctx, "cline-pass/glm-5.3"))
|
||||
}
|
||||
|
||||
// 已识别的错误只冷却请求模型所属的钱包(或免费模型本身),不停整个账号、不置 error。
|
||||
func TestHandleUpstreamError_ClineCoolsDownWallet(t *testing.T) {
|
||||
paid := "deepseek/deepseek-v4-flash"
|
||||
pass := "cline-pass/glm-5.3"
|
||||
now := time.Now()
|
||||
snapshot := withClinePassSnapshot(clineTestAccount(1), map[string]any{
|
||||
"cline_5h_reset_at": now.Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
})
|
||||
cases := []struct {
|
||||
name string
|
||||
account *Account
|
||||
model string
|
||||
status int
|
||||
body string
|
||||
scope string
|
||||
wait time.Duration
|
||||
reason string
|
||||
}{
|
||||
{"pass weekly", clineTestAccount(10), pass, http.StatusTooManyRequests, clinePassWeeklyBody, clinePassRateLimitKey, clinePassWeeklyRecheck, clinePassLimitReason},
|
||||
{"pass 5h from snapshot", snapshot, pass, http.StatusTooManyRequests, clinePass5hBody, clinePassRateLimitKey, 2 * time.Hour, clinePassLimitReason},
|
||||
{"not subscribed", clineTestAccount(11), pass, http.StatusForbidden, clineNotSubscribedBody, clinePassRateLimitKey, clineLimitRecheck, clinePassUnavailableReason},
|
||||
{"credits", clineTestAccount(12), paid, http.StatusPaymentRequired, clineInsufficientCreditsBody, clineCreditsRateLimitKey, cnBalanceCheckCooldown(nil), clineCreditsReason},
|
||||
{"unrecognized 402", clineTestAccount(13), paid, http.StatusPaymentRequired, `{"error":{"message":"Payment required"}}`, clineCreditsRateLimitKey, cnBalanceCheckCooldown(nil), clineCreditsReason},
|
||||
{"spend limit", clineTestAccount(14), paid, http.StatusTooManyRequests, clineSpendLimitBody, clineCreditsRateLimitKey, clineLimitRecheck, clineSpendLimitReason},
|
||||
{"inference cap on paid model", clineTestAccount(15), paid, http.StatusTooManyRequests, clineInferenceCapBody, clineCreditsRateLimitKey, clineLimitRecheck, clineInferenceCapReason},
|
||||
{"inference cap on pass model", clineTestAccount(16), pass, http.StatusTooManyRequests, clineInferenceCapBody, clinePassRateLimitKey, clineLimitRecheck, clineInferenceCapReason},
|
||||
{"unknown model uses the error's wallet", clineTestAccount(17), "", http.StatusForbidden, clineOrgPassBody, clinePassRateLimitKey, clineLimitRecheck, clinePassUnavailableReason},
|
||||
{"free model", clineTestAccount(18), "cline-free/kimi-k3", http.StatusTooManyRequests, clineFreeModelLimitBody, "cline-free/kimi-k3", 3*time.Hour + 20*time.Minute, clineFreeModelLimitReason},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
ctx := context.Background()
|
||||
if tc.model != "" {
|
||||
ctx = withTempUnschedulableModel(ctx, []string{tc.model})
|
||||
}
|
||||
repo := &commandCodeRateLimitRepo{}
|
||||
|
||||
shouldDisable := NewRateLimitService(repo, nil, &config.Config{}, nil, nil).HandleUpstreamError(ctx, tc.account,
|
||||
tc.status, http.Header{}, []byte(tc.body))
|
||||
|
||||
require.False(t, shouldDisable, tc.name)
|
||||
require.Equal(t, 0, repo.tempCalls, tc.name)
|
||||
require.Equal(t, 0, repo.setErrorCalls, tc.name)
|
||||
require.Equal(t, 0, repo.rateLimitedCalls, tc.name)
|
||||
require.Len(t, repo.modelLimits, 1, tc.name)
|
||||
require.WithinDuration(t, time.Now().Add(tc.wait), repo.modelLimits[tc.scope], time.Minute, tc.name)
|
||||
require.True(t, strings.HasPrefix(repo.reasons[tc.scope], tc.reason), tc.name+": "+repo.reasons[tc.scope])
|
||||
require.True(t, tc.account.isRateLimitActiveForKey(tc.scope), tc.name)
|
||||
}
|
||||
|
||||
// 普通频率限制交回默认 429 逻辑。
|
||||
plain := &commandCodeRateLimitRepo{}
|
||||
NewRateLimitService(plain, nil, &config.Config{}, nil, nil).HandleUpstreamError(withTempUnschedulableModel(context.Background(), []string{paid}),
|
||||
clineTestAccount(30), http.StatusTooManyRequests, http.Header{}, []byte(clineRateLimitBody))
|
||||
require.Empty(t, plain.modelLimits)
|
||||
require.Equal(t, 0, plain.tempCalls)
|
||||
}
|
||||
|
||||
// clineModelLimitFailRepo 让钱包冷却写入失败。
|
||||
type clineModelLimitFailRepo struct {
|
||||
commandCodeRateLimitRepo
|
||||
}
|
||||
|
||||
func (r *clineModelLimitFailRepo) SetModelRateLimit(context.Context, int64, string, time.Time, ...string) error {
|
||||
return errors.New("db unavailable")
|
||||
}
|
||||
|
||||
// 积分冷却写入失败时 402 也不落入默认分支永久置 error:探测只覆盖激活账号,置 error 后
|
||||
// 充值也无法自动恢复。
|
||||
func TestHandleUpstreamError_Cline402KeepsAccountWhenCooldownWriteFails(t *testing.T) {
|
||||
ctx := withTempUnschedulableModel(context.Background(), []string{"deepseek/deepseek-v4-flash"})
|
||||
for _, body := range []string{clineInsufficientCreditsBody, `{"error":{"message":"Payment required"}}`} {
|
||||
repo := &clineModelLimitFailRepo{}
|
||||
|
||||
shouldDisable := NewRateLimitService(repo, nil, &config.Config{}, nil, nil).HandleUpstreamError(ctx, clineTestAccount(19),
|
||||
http.StatusPaymentRequired, http.Header{}, []byte(body))
|
||||
|
||||
require.False(t, shouldDisable, body)
|
||||
require.Equal(t, 0, repo.setErrorCalls, body)
|
||||
require.Equal(t, 0, repo.tempCalls, body)
|
||||
}
|
||||
}
|
||||
|
||||
// clineAccountRepo 为余额 / 用量探测与周期检测提供 Cline 账号,并记录写入。
|
||||
type clineAccountRepo struct {
|
||||
AccountRepository
|
||||
accounts []Account
|
||||
mu sync.Mutex
|
||||
extraWrites map[int64][]map[string]any
|
||||
modelLimits map[int64]map[string]time.Time
|
||||
reasons map[int64]map[string]string
|
||||
paused []int64
|
||||
}
|
||||
|
||||
func (r *clineAccountRepo) ListByPlatform(_ context.Context, platform string) ([]Account, error) {
|
||||
if platform == PlatformCline {
|
||||
return r.accounts, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *clineAccountRepo) GetByID(_ context.Context, id int64) (*Account, error) {
|
||||
for i := range r.accounts {
|
||||
if r.accounts[i].ID == id {
|
||||
return &r.accounts[i], nil
|
||||
}
|
||||
}
|
||||
return nil, errors.New("not found")
|
||||
}
|
||||
|
||||
func (r *clineAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.extraWrites == nil {
|
||||
r.extraWrites = map[int64][]map[string]any{}
|
||||
}
|
||||
r.extraWrites[id] = append(r.extraWrites[id], updates)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *clineAccountRepo) SetModelRateLimit(_ context.Context, id int64, scope string, resetAt time.Time, reason ...string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.modelLimits == nil {
|
||||
r.modelLimits = map[int64]map[string]time.Time{}
|
||||
r.reasons = map[int64]map[string]string{}
|
||||
}
|
||||
if r.modelLimits[id] == nil {
|
||||
r.modelLimits[id] = map[string]time.Time{}
|
||||
r.reasons[id] = map[string]string{}
|
||||
}
|
||||
r.modelLimits[id][scope] = resetAt
|
||||
if len(reason) > 0 {
|
||||
r.reasons[id][scope] = reason[0]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *clineAccountRepo) SetTempUnschedulable(_ context.Context, id int64, _ time.Time, _ string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.paused = append(r.paused, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
const (
|
||||
clineUsageLimitsPath = "/api/v1/users/me/plan/usage-limits"
|
||||
clineUsageLimitsBody = `{"data":{"limits":[{"type":"five_hour","percentUsed":12.5,"resetsAt":"2099-01-01T05:00:00.123456Z"},{"type":"weekly","percentUsed":40},{"type":"monthly","percentUsed":0}]},"success":true}`
|
||||
)
|
||||
|
||||
func newClineAccountUpstream(usageLimits commandCodeAlphaResponse, balance string) *commandCodeAlphaUpstream {
|
||||
return &commandCodeAlphaUpstream{responses: map[string]commandCodeAlphaResponse{
|
||||
clineUsageLimitsPath: usageLimits,
|
||||
"/api/v1/users/me": {status: http.StatusOK, body: `{"success":true,"data":{"id":"u-1","email":"a@example.com"}}`},
|
||||
"/api/v1/users/u-1/balance": {status: http.StatusOK, body: balance},
|
||||
}}
|
||||
}
|
||||
|
||||
// balance 是百万分之一积分(1 积分 = 1 美元),按 Cline 的展示精度保留 4 位小数。
|
||||
func TestClineBalanceProbeConvertsMicrocredits(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
balance string
|
||||
want float64
|
||||
}{
|
||||
{`{"balance":12340000,"userId":"u-1"}`, 12.34},
|
||||
{`{"success":true,"data":{"balance":"500000","userId":"u-1"}}`, 0.5},
|
||||
{`{"balance":1234567,"userId":"u-1"}`, 1.2346},
|
||||
} {
|
||||
account := clineTestAccount(40)
|
||||
upstream := newClineAccountUpstream(commandCodeAlphaResponse{}, tc.balance)
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
svc := NewCNProviderBalanceService(repo, nil, upstream, &config.Config{})
|
||||
|
||||
result, err := svc.QueryBalance(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.Success)
|
||||
require.True(t, result.Persisted)
|
||||
require.InDelta(t, tc.want, result.Balance, 1e-9)
|
||||
require.Equal(t, "USD", result.Currency)
|
||||
require.Len(t, upstream.requests, 2)
|
||||
require.Equal(t, "Bearer cline_test_key", upstream.requests[0].Header.Get("Authorization"))
|
||||
require.InDelta(t, tc.want, repo.extraWrites[0]["cline_balance"], 1e-9)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClineBalanceProbeFailures(t *testing.T) {
|
||||
account := clineTestAccount(41)
|
||||
upstream := &commandCodeAlphaUpstream{responses: map[string]commandCodeAlphaResponse{
|
||||
"/api/v1/users/me": {status: http.StatusUnauthorized, body: `{"error":"Unauthorized"}`},
|
||||
}}
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
result, err := NewCNProviderBalanceService(repo, nil, upstream, &config.Config{}).QueryBalance(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.False(t, result.Success)
|
||||
require.Equal(t, http.StatusUnauthorized, result.StatusCode)
|
||||
require.Contains(t, result.Error, "HTTP 401")
|
||||
require.Empty(t, repo.extraWrites)
|
||||
|
||||
envelope := &commandCodeAlphaUpstream{responses: map[string]commandCodeAlphaResponse{
|
||||
"/api/v1/users/me": {status: http.StatusOK, body: `{"success":false,"error":"token expired"}`},
|
||||
}}
|
||||
result, err = NewCNProviderBalanceService(&cnBalanceProbeRepo{account: account}, nil, envelope, &config.Config{}).QueryBalance(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.False(t, result.Success)
|
||||
require.Contains(t, result.Error, "token expired")
|
||||
}
|
||||
|
||||
// 余额与用量只查官方主机上的 API Key 账号:自定义中转的 Key 不发往官方账号接口。
|
||||
func TestClineAccountAPIRequiresOfficialHost(t *testing.T) {
|
||||
require.NoError(t, validatePayGAccount(clineTestAccount(50)))
|
||||
require.NoError(t, validateCodingPlanAccount(clineTestAccount(50)))
|
||||
require.Equal(t, PlatformCline, clineTestAccount(50).GetCodingPlanProvider())
|
||||
|
||||
relay := clineTestAccount(52)
|
||||
relay.Credentials["base_url"] = "https://relay.example.com/v1"
|
||||
relay.Credentials["api_protocol"] = APIProtocolChatCompletions
|
||||
require.Error(t, validatePayGAccount(relay))
|
||||
require.Error(t, validateCodingPlanAccount(relay))
|
||||
require.Empty(t, relay.GetCodingPlanProvider())
|
||||
}
|
||||
|
||||
func clineUsageTestConfig() *config.Config {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.CNProviders.BalanceThreshold = 0.5
|
||||
return cfg
|
||||
}
|
||||
|
||||
// 一次探测刷新三档窗口与积分余额;两个钱包都可用时不写冷却。
|
||||
func TestClineUsageProbeWritesWindowsAndBalance(t *testing.T) {
|
||||
account := clineTestAccount(60)
|
||||
repo := &clineAccountRepo{accounts: []Account{*account}}
|
||||
upstream := newClineAccountUpstream(commandCodeAlphaResponse{status: http.StatusOK, body: clineUsageLimitsBody}, `{"success":true,"data":{"balance":12340000}}`)
|
||||
|
||||
result, err := NewCNProviderQuotaService(repo, nil, upstream, clineUsageTestConfig()).QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.Success)
|
||||
require.True(t, result.CredentialValid)
|
||||
require.True(t, result.Persisted)
|
||||
require.Equal(t, "ClinePass", result.PlanLevel)
|
||||
require.Equal(t, []CNQuotaTier{
|
||||
{Window: "5h", UsedPercent: 12.5, ResetAt: "2099-01-01T05:00:00Z"},
|
||||
{Window: "weekly", UsedPercent: 40},
|
||||
{Window: "monthly", UsedPercent: 0},
|
||||
}, result.Tiers)
|
||||
require.NotNil(t, result.Balance)
|
||||
require.InDelta(t, 12.34, result.Balance.Balance, 1e-9)
|
||||
require.Equal(t, "Bearer cline_test_key", upstream.requests[0].Header.Get("Authorization"))
|
||||
|
||||
written := repo.extraWrites[60][0]
|
||||
require.Equal(t, 12.5, written["cline_5h_used_percent"])
|
||||
require.Equal(t, "2099-01-01T05:00:00Z", written["cline_5h_reset_at"])
|
||||
require.Equal(t, 40.0, written["cline_weekly_used_percent"])
|
||||
require.Contains(t, written, "cline_weekly_reset_at")
|
||||
require.Nil(t, written["cline_weekly_reset_at"], "a window that has not started has no reset time")
|
||||
require.InDelta(t, 12.34, written["cline_balance"], 1e-9)
|
||||
require.Empty(t, repo.modelLimits)
|
||||
}
|
||||
|
||||
// 窗口用满冷却 ClinePass 到重置时间;积分低于阈值冷却积分钱包;都不停整个账号。
|
||||
func TestClineUsageProbeCoolsDownExhaustedWallets(t *testing.T) {
|
||||
resetAt := time.Now().Add(72 * time.Hour).UTC().Truncate(time.Second)
|
||||
body := `{"data":{"limits":[{"type":"five_hour","percentUsed":30},{"type":"weekly","percentUsed":100,"resetsAt":"` +
|
||||
resetAt.Format(time.RFC3339) + `"}]},"success":true}`
|
||||
account := clineTestAccount(61)
|
||||
repo := &clineAccountRepo{accounts: []Account{*account}}
|
||||
upstream := newClineAccountUpstream(commandCodeAlphaResponse{status: http.StatusOK, body: body}, `{"success":true,"data":{"balance":200000}}`)
|
||||
|
||||
_, err := NewCNProviderQuotaService(repo, nil, upstream, clineUsageTestConfig()).QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, repo.modelLimits[61][clinePassRateLimitKey].Equal(resetAt))
|
||||
require.Equal(t, clinePassLimitReason+": weekly window used up", repo.reasons[61][clinePassRateLimitKey])
|
||||
require.WithinDuration(t, time.Now().Add(cnBalanceCheckCooldown(nil)), repo.modelLimits[61][clineCreditsRateLimitKey], time.Minute)
|
||||
require.True(t, strings.HasPrefix(repo.reasons[61][clineCreditsRateLimitKey], clineCreditsLowReason), repo.reasons[61][clineCreditsRateLimitKey])
|
||||
require.Empty(t, repo.paused)
|
||||
}
|
||||
|
||||
// 没有订阅(404 或空列表):清掉旧窗口快照,ClinePass 模型不调度到该账号。
|
||||
func TestClineUsageProbeWithoutSubscription(t *testing.T) {
|
||||
for _, usage := range []commandCodeAlphaResponse{
|
||||
{status: http.StatusNotFound, body: `{"success":false,"error":"no plan history found for user"}`},
|
||||
{status: http.StatusOK, body: `{"success":true,"data":{"limits":[]}}`},
|
||||
} {
|
||||
account := withClinePassSnapshot(clineTestAccount(62), nil)
|
||||
repo := &clineAccountRepo{accounts: []Account{*account}}
|
||||
upstream := newClineAccountUpstream(usage, `{"success":true,"data":{"balance":5000000}}`)
|
||||
|
||||
result, err := NewCNProviderQuotaService(repo, nil, upstream, clineUsageTestConfig()).QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.Success)
|
||||
require.Empty(t, result.Tiers)
|
||||
require.Empty(t, result.PlanLevel)
|
||||
written := repo.extraWrites[62][0]
|
||||
for _, key := range []string{"cline_5h_used_percent", "cline_weekly_used_percent", "cline_monthly_used_percent"} {
|
||||
require.Contains(t, written, key)
|
||||
require.Nil(t, written[key])
|
||||
}
|
||||
require.True(t, strings.HasPrefix(repo.reasons[62][clinePassRateLimitKey], clinePassUnavailableReason))
|
||||
require.NotContains(t, repo.modelLimits[62], clineCreditsRateLimitKey)
|
||||
}
|
||||
}
|
||||
|
||||
// 数据显示钱包恢复可用时解除冷却;组织花费上限无从探测,不提前解除。
|
||||
func TestClineUsageProbeClearsRecoveredWallets(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
creditsReason string
|
||||
cleared bool
|
||||
}{
|
||||
{clineCreditsReason + ": Insufficient credits", true},
|
||||
{clineCreditsLowReason + ": balance 0.1 USD below threshold 0.50", true},
|
||||
{clineSpendLimitReason + ": Organization spend limit exceeded", false},
|
||||
} {
|
||||
account := clineTestAccount(63)
|
||||
now := time.Now()
|
||||
setAccountModelRateLimitSnapshot(account, clinePassRateLimitKey, now.Add(6*time.Hour), clinePassLimitReason, now)
|
||||
setAccountModelRateLimitSnapshot(account, clineCreditsRateLimitKey, now.Add(time.Hour), tc.creditsReason, now)
|
||||
repo := &clineAccountRepo{accounts: []Account{*account}}
|
||||
upstream := newClineAccountUpstream(commandCodeAlphaResponse{status: http.StatusOK, body: clineUsageLimitsBody}, `{"success":true,"data":{"balance":5000000}}`)
|
||||
|
||||
_, err := NewCNProviderQuotaService(repo, nil, upstream, clineUsageTestConfig()).QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.False(t, repo.modelLimits[63][clinePassRateLimitKey].After(time.Now()), tc.creditsReason)
|
||||
_, touched := repo.modelLimits[63][clineCreditsRateLimitKey]
|
||||
require.Equal(t, tc.cleared, touched, tc.creditsReason)
|
||||
if tc.cleared {
|
||||
require.False(t, repo.modelLimits[63][clineCreditsRateLimitKey].After(time.Now()), tc.creditsReason)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ClinePass 同样只解除探测能证明已恢复的冷却(窗口用满、未订阅):组织账户、所需订阅
|
||||
// 档位与推理额度上限反映不到用量接口里,提前解除会让每个检测周期都重放一次失败请求。
|
||||
func TestClineUsageProbeClearsOnlyProvablePassCooldowns(t *testing.T) {
|
||||
ctx := withTempUnschedulableModel(context.Background(), []string{"cline-pass/glm-5.3"})
|
||||
balance := `{"success":true,"data":{"balance":5000000}}`
|
||||
for _, tc := range []struct {
|
||||
status int
|
||||
body string
|
||||
cleared bool
|
||||
}{
|
||||
{http.StatusTooManyRequests, clinePassWeeklyBody, true},
|
||||
{http.StatusForbidden, clineNotSubscribedBody, false},
|
||||
{http.StatusForbidden, clineOrgPassBody, false},
|
||||
{http.StatusTooManyRequests, clineInferenceCapBody, false},
|
||||
} {
|
||||
account := clineTestAccount(65)
|
||||
NewRateLimitService(&commandCodeRateLimitRepo{}, nil, &config.Config{}, nil, nil).HandleUpstreamError(ctx, account,
|
||||
tc.status, http.Header{}, []byte(tc.body))
|
||||
require.True(t, account.isRateLimitActiveForKey(clinePassRateLimitKey), tc.body)
|
||||
repo := &clineAccountRepo{accounts: []Account{*account}}
|
||||
upstream := newClineAccountUpstream(commandCodeAlphaResponse{status: http.StatusOK, body: clineUsageLimitsBody}, balance)
|
||||
|
||||
_, err := NewCNProviderQuotaService(repo, nil, upstream, clineUsageTestConfig()).QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
_, touched := repo.modelLimits[65][clinePassRateLimitKey]
|
||||
require.Equal(t, tc.cleared, touched, tc.body)
|
||||
}
|
||||
|
||||
// 探测发现未订阅设的冷却,订阅后的下一次探测解除。
|
||||
repo := &clineAccountRepo{accounts: []Account{*clineTestAccount(66)}}
|
||||
unsubscribed := newClineAccountUpstream(commandCodeAlphaResponse{status: http.StatusNotFound,
|
||||
body: `{"success":false,"error":"no plan history found for user"}`}, balance)
|
||||
_, err := NewCNProviderQuotaService(repo, nil, unsubscribed, clineUsageTestConfig()).QueryUsage(context.Background(), 66)
|
||||
require.NoError(t, err)
|
||||
require.True(t, repo.modelLimits[66][clinePassRateLimitKey].After(time.Now()))
|
||||
|
||||
subscribed := newClineAccountUpstream(commandCodeAlphaResponse{status: http.StatusOK, body: clineUsageLimitsBody}, balance)
|
||||
_, err = NewCNProviderQuotaService(repo, nil, subscribed, clineUsageTestConfig()).QueryUsage(context.Background(), 66)
|
||||
require.NoError(t, err)
|
||||
require.False(t, repo.modelLimits[66][clinePassRateLimitKey].After(time.Now()))
|
||||
}
|
||||
|
||||
// 用量接口鉴权失败:探测失败但不写快照、不动冷却。
|
||||
func TestClineUsageProbeAuthFailure(t *testing.T) {
|
||||
account := clineTestAccount(64)
|
||||
repo := &clineAccountRepo{accounts: []Account{*account}}
|
||||
upstream := newClineAccountUpstream(commandCodeAlphaResponse{status: http.StatusUnauthorized, body: `{"error":"Unauthorized"}`}, `{}`)
|
||||
|
||||
result, err := NewCNProviderQuotaService(repo, nil, upstream, clineUsageTestConfig()).QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.False(t, result.Success)
|
||||
require.False(t, result.CredentialValid)
|
||||
require.Equal(t, http.StatusUnauthorized, result.StatusCode)
|
||||
require.Empty(t, repo.extraWrites)
|
||||
require.Empty(t, repo.modelLimits)
|
||||
}
|
||||
|
||||
// 周期检测探测官方主机上的全部激活账号(含手动停调的,钱包冷却与调度开关无关),
|
||||
// 只冷却钱包,不停整个账号。
|
||||
func TestCNProviderBalanceCheckRunOnceProbesClineUsage(t *testing.T) {
|
||||
official := *clineTestAccount(1)
|
||||
manualOff := *clineTestAccount(2)
|
||||
manualOff.Schedulable = false
|
||||
relay := *clineTestAccount(3)
|
||||
relay.Credentials["base_url"] = "https://relay.example.com/v1"
|
||||
relay.Credentials["api_protocol"] = APIProtocolChatCompletions
|
||||
|
||||
repo := &clineAccountRepo{accounts: []Account{official, manualOff, relay}}
|
||||
upstream := newClineAccountUpstream(commandCodeAlphaResponse{status: http.StatusOK, body: clineUsageLimitsBody}, `{"success":true,"data":{"balance":200000}}`)
|
||||
cfg := clineUsageTestConfig()
|
||||
svc := &CNProviderBalanceCheckService{
|
||||
accountRepo: repo,
|
||||
balanceService: NewCNProviderBalanceService(repo, nil, upstream, cfg),
|
||||
quotaService: NewCNProviderQuotaService(repo, nil, upstream, cfg),
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
svc.runOnce()
|
||||
|
||||
require.Len(t, upstream.requests, 6, "usage-limits + me + balance for each official account")
|
||||
require.Len(t, repo.extraWrites, 2)
|
||||
require.Contains(t, repo.modelLimits[1], clineCreditsRateLimitKey)
|
||||
require.Contains(t, repo.modelLimits[2], clineCreditsRateLimitKey)
|
||||
require.Empty(t, repo.paused)
|
||||
}
|
||||
|
||||
// 同一轮检测里每个账号走自己平台登记的检测函数:积分同样低于阈值,Command Code 停整个
|
||||
// 账号,Cline 不停(钱包冷却在探测中处理)。
|
||||
func TestCNProviderBalanceCheckRunOnceDispatchesAccountProbesByPlatform(t *testing.T) {
|
||||
commandCode := *commandCodeUsageAccount()
|
||||
commandCode.ID, commandCode.Schedulable = 1, true
|
||||
cline := *clineTestAccount(2)
|
||||
repo := &commandCodeCheckRepo{accounts: []Account{commandCode, cline}}
|
||||
prober := &commandCodeBalanceProber{balance: 0.2}
|
||||
svc := &CNProviderBalanceCheckService{accountRepo: repo, quotaService: prober, cfg: clineUsageTestConfig()}
|
||||
|
||||
svc.runOnce()
|
||||
|
||||
require.ElementsMatch(t, []int64{1, 2}, prober.probed)
|
||||
require.Equal(t, []int64{1}, repo.paused)
|
||||
}
|
||||
|
||||
// 官方主机不做上游计费探测:中转站的计费接口在官方 API 上不存在,只会把 Key 发往官方主机。
|
||||
func TestClineOfficialHostSkipsUpstreamBillingProbe(t *testing.T) {
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI(DefaultClineBaseURL))
|
||||
require.False(t, upstreamBillingProbeTargetIsOfficialAPI("https://relay.example.com/v1"))
|
||||
}
|
||||
@@ -0,0 +1,368 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// Cline 积分余额与 ClinePass 用量来自账号接口(Cline 控制台与客户端同源),账号 API Key
|
||||
// 可直接访问(Bearer <key>;实测带 workos: 前缀反而 401):
|
||||
//
|
||||
// GET /api/v1/users/me -> {id, email, ...}
|
||||
// GET /api/v1/users/{id}/balance -> {balance, userId}
|
||||
// GET /api/v1/users/me/plan/usage-limits -> {limits:[{type, percentUsed, resetsAt?}]}
|
||||
//
|
||||
// 响应包在 {success, error, data} 信封里。balance 单位为百万分之一积分(1 积分 =
|
||||
// 1 美元;Cline CLI 的 normalizeCreditBalance 除以 1e6,扩展先除 100 再除 10000),
|
||||
// 按 Cline 的展示精度保留 4 位小数。usage-limits 的 type 为 five_hour / weekly /
|
||||
// monthly(套餐里对应最近 5 小时 / 7 天 / 30 天的滚动用量上限),percentUsed 为 0-100,
|
||||
// 窗口未开始时没有 resetsAt;未订阅时返回 404(no plan history found for user)或空列表。
|
||||
|
||||
const clineAPIBase = "https://" + clineAPIHost
|
||||
|
||||
var errClineInvalidResponse = errors.New("invalid cline account response")
|
||||
|
||||
// clineHTTPError 是账号接口的非 2xx 响应。
|
||||
type clineHTTPError struct {
|
||||
status int
|
||||
body string
|
||||
}
|
||||
|
||||
func (e *clineHTTPError) Error() string {
|
||||
return fmt.Sprintf("HTTP %d: %s", e.status, e.body)
|
||||
}
|
||||
|
||||
type clineAccountClient struct {
|
||||
cfg *config.Config
|
||||
httpUpstream HTTPUpstream
|
||||
account *Account
|
||||
apiKey string
|
||||
proxyURL string
|
||||
// errPrefix 是应用错误码前缀(CN_BALANCE / CN_QUOTA),与调用方服务的错误码一致。
|
||||
errPrefix string
|
||||
}
|
||||
|
||||
func newClineAccountClient(cfg *config.Config, httpUpstream HTTPUpstream, account *Account, proxyURL, errPrefix string) *clineAccountClient {
|
||||
return &clineAccountClient{
|
||||
cfg: cfg,
|
||||
httpUpstream: httpUpstream,
|
||||
account: account,
|
||||
apiKey: strings.TrimSpace(account.GetCNAPIKey()),
|
||||
proxyURL: proxyURL,
|
||||
errPrefix: errPrefix,
|
||||
}
|
||||
}
|
||||
|
||||
// get 请求一个账号接口,返回拆掉 {success, error, data} 信封后的 JSON。
|
||||
func (c *clineAccountClient) get(ctx context.Context, path string) (gjson.Result, error) {
|
||||
validatedURL, err := cnValidateProbeURL(c.cfg, clineAPIBase+path)
|
||||
if err != nil {
|
||||
return gjson.Result{}, infraerrors.New(http.StatusForbidden, c.errPrefix+"_URL_REJECTED", err.Error())
|
||||
}
|
||||
callCtx, cancel := context.WithTimeout(ctx, cnQuotaUpstreamTimeout)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(callCtx, http.MethodGet, validatedURL, nil)
|
||||
if err != nil {
|
||||
return gjson.Result{}, infraerrors.Newf(http.StatusInternalServerError, c.errPrefix+"_REQUEST_BUILD_FAILED", "build request: %v", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
c.account.ApplyHeaderOverrides(req.Header)
|
||||
|
||||
resp, err := c.httpUpstream.Do(req, c.proxyURL, c.account.ID, maxInt(c.account.Concurrency, 1))
|
||||
if err != nil {
|
||||
return gjson.Result{}, fmt.Errorf("upstream request failed: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, cnQuotaMaxBodyBytes))
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return gjson.Result{}, &clineHTTPError{status: resp.StatusCode, body: truncate(strings.TrimSpace(string(body)), 240)}
|
||||
}
|
||||
if !gjson.ValidBytes(body) {
|
||||
return gjson.Result{}, fmt.Errorf("%w: %s", errClineInvalidResponse, path)
|
||||
}
|
||||
parsed := gjson.ParseBytes(body)
|
||||
if success := parsed.Get("success"); success.Exists() {
|
||||
if !success.Bool() {
|
||||
return gjson.Result{}, &clineHTTPError{status: resp.StatusCode, body: truncate(parsed.Get("error").String(), 240)}
|
||||
}
|
||||
return parsed.Get("data"), nil
|
||||
}
|
||||
return parsed, nil
|
||||
}
|
||||
|
||||
// balance 查询积分余额(美元)。
|
||||
func (c *clineAccountClient) balance(ctx context.Context) (float64, error) {
|
||||
me, err := c.get(ctx, "/api/v1/users/me")
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
userID := strings.TrimSpace(me.Get("id").String())
|
||||
if userID == "" {
|
||||
return 0, fmt.Errorf("%w: missing user id", errClineInvalidResponse)
|
||||
}
|
||||
balance, err := c.get(ctx, "/api/v1/users/"+url.PathEscape(userID)+"/balance")
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
micro := balance.Get("balance")
|
||||
if !micro.Exists() || (micro.Type != gjson.Number && micro.Type != gjson.String) {
|
||||
return 0, fmt.Errorf("%w: missing balance", errClineInvalidResponse)
|
||||
}
|
||||
return math.Round(micro.Float()/100) / 10000, nil
|
||||
}
|
||||
|
||||
// passLimits 查询 ClinePass 三档窗口;subscribed=false 表示没有订阅。
|
||||
func (c *clineAccountClient) passLimits(ctx context.Context) (tiers []CNQuotaTier, subscribed bool, err error) {
|
||||
data, err := c.get(ctx, "/api/v1/users/me/plan/usage-limits")
|
||||
if err != nil {
|
||||
var httpErr *clineHTTPError
|
||||
if errors.As(err, &httpErr) && httpErr.status == http.StatusNotFound {
|
||||
return nil, false, nil
|
||||
}
|
||||
return nil, false, err
|
||||
}
|
||||
limits := data.Get("limits")
|
||||
if !limits.IsArray() {
|
||||
return nil, false, fmt.Errorf("%w: missing limits", errClineInvalidResponse)
|
||||
}
|
||||
windows := map[string]string{"five_hour": "5h", "weekly": "weekly", "monthly": "monthly"}
|
||||
for _, item := range limits.Array() {
|
||||
window, ok := windows[strings.ToLower(strings.TrimSpace(item.Get("type").String()))]
|
||||
used := item.Get("percentUsed")
|
||||
if !ok || used.Type != gjson.Number {
|
||||
continue
|
||||
}
|
||||
tier := CNQuotaTier{Window: window, UsedPercent: used.Float()}
|
||||
if resetAt, err := time.Parse(time.RFC3339Nano, strings.TrimSpace(item.Get("resetsAt").String())); err == nil {
|
||||
tier.ResetAt = resetAt.UTC().Format(time.RFC3339)
|
||||
}
|
||||
tiers = append(tiers, tier)
|
||||
}
|
||||
return tiers, len(tiers) > 0, nil
|
||||
}
|
||||
|
||||
// clineFetchFailure 把取数错误转换为探测失败的状态码与文案;ok=false 表示应作为错误返回。
|
||||
func clineFetchFailure(err error) (status int, message string, ok bool) {
|
||||
var httpErr *clineHTTPError
|
||||
switch {
|
||||
case errors.As(err, &httpErr):
|
||||
return httpErr.status, fmt.Sprintf("API error (HTTP %d): %s", httpErr.status, httpErr.body), true
|
||||
case errors.Is(err, errClineInvalidResponse):
|
||||
return http.StatusOK, "Invalid response", true
|
||||
default:
|
||||
return 0, "", false
|
||||
}
|
||||
}
|
||||
|
||||
func clineApplicationError(err error, errPrefix string) error {
|
||||
if isApplicationError(err) {
|
||||
return err
|
||||
}
|
||||
return infraerrors.Newf(http.StatusBadGateway, errPrefix+"_REQUEST_FAILED", "%v", err)
|
||||
}
|
||||
|
||||
func clineBalanceResult(usd float64, now time.Time) *CNProviderBalanceResult {
|
||||
return &CNProviderBalanceResult{
|
||||
Provider: PlatformCline,
|
||||
Success: true,
|
||||
Balance: usd,
|
||||
Currency: "USD",
|
||||
Balances: []CNProviderBalanceEntry{{Currency: "USD", Balance: usd}},
|
||||
Available: true,
|
||||
StatusCode: http.StatusOK,
|
||||
FetchedAt: now.Unix(),
|
||||
}
|
||||
}
|
||||
|
||||
// queryClineBalance 只查积分余额(美元)并落余额快照。
|
||||
func (s *CNProviderBalanceService) queryClineBalance(ctx context.Context, account *Account) (*CNProviderBalanceResult, error) {
|
||||
client := newClineAccountClient(s.cfg, s.httpUpstream, account, s.resolveProxyURL(ctx, account), "CN_BALANCE")
|
||||
if client.apiKey == "" {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_BALANCE_NO_APIKEY", "account api_key is empty")
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
usd, err := client.balance(ctx)
|
||||
if err != nil {
|
||||
status, message, ok := clineFetchFailure(err)
|
||||
if !ok {
|
||||
return nil, clineApplicationError(err, "CN_BALANCE")
|
||||
}
|
||||
return &CNProviderBalanceResult{Provider: PlatformCline, StatusCode: status, FetchedAt: now.Unix(), Available: true,
|
||||
Error: "Balance query failed: " + message}, nil
|
||||
}
|
||||
result := clineBalanceResult(usd, now)
|
||||
if err := s.accountRepo.UpdateExtra(ctx, account.ID, cnBalanceExtraUpdates(PlatformCline, result, now)); err != nil {
|
||||
slog.Warn("cn_balance_persist_failed", "account_id", account.ID, "provider", PlatformCline, "error", err)
|
||||
} else {
|
||||
result.Persisted = true
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// queryClineUsage 查询 ClinePass 窗口与积分余额,落快照,并按结果设置或解除两个钱包的
|
||||
// 冷却(见 applyClineWallets)。余额查询失败不影响窗口结果。
|
||||
func (s *CNProviderQuotaService) queryClineUsage(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
|
||||
client := newClineAccountClient(s.cfg, s.httpUpstream, account, s.resolveProxyURL(ctx, account), "CN_QUOTA")
|
||||
if client.apiKey == "" {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NO_APIKEY", "account api_key is empty")
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
result := &CNProviderQuotaProbeResult{
|
||||
Provider: PlatformCline,
|
||||
Source: "cline_usage_limits",
|
||||
FetchedAt: now.Unix(),
|
||||
}
|
||||
tiers, subscribed, err := client.passLimits(ctx)
|
||||
if err != nil {
|
||||
status, message, ok := clineFetchFailure(err)
|
||||
if !ok {
|
||||
return nil, clineApplicationError(err, "CN_QUOTA")
|
||||
}
|
||||
result.StatusCode = status
|
||||
result.Error = message
|
||||
result.CredentialValid = status != http.StatusUnauthorized && status != http.StatusForbidden
|
||||
return result, nil
|
||||
}
|
||||
result.StatusCode = http.StatusOK
|
||||
result.Success = true
|
||||
result.CredentialValid = true
|
||||
result.Tiers = tiers
|
||||
if subscribed {
|
||||
result.PlanLevel = "ClinePass"
|
||||
}
|
||||
|
||||
updates := cnQuotaExtraUpdates(PlatformCline, tiers, now)
|
||||
present := make(map[string]CNQuotaTier, len(tiers))
|
||||
for _, tier := range tiers {
|
||||
present[tier.Window] = tier
|
||||
}
|
||||
// 缺席的窗口写 nil:取消订阅后不沿用旧快照(clinePassSubscribed 据此判断);窗口尚未
|
||||
// 开始时没有重置时间,同样清掉上一次的值。
|
||||
for _, window := range []struct{ name, used, reset string }{
|
||||
{"5h", cnExtraSuffix5hUsed, cnExtraSuffix5hReset},
|
||||
{"weekly", cnExtraSuffixWeeklyUsed, cnExtraSuffixWeeklyReset},
|
||||
{"monthly", cnExtraSuffixMonthlyUsed, cnExtraSuffixMonthlyReset},
|
||||
} {
|
||||
tier, ok := present[window.name]
|
||||
if !ok {
|
||||
updates[cnExtraKey(PlatformCline, window.used)] = nil
|
||||
}
|
||||
if tier.ResetAt == "" {
|
||||
updates[cnExtraKey(PlatformCline, window.reset)] = nil
|
||||
}
|
||||
}
|
||||
if usd, err := client.balance(ctx); err != nil {
|
||||
slog.Warn("cline_balance_probe_failed", "account_id", account.ID, "error", err)
|
||||
} else {
|
||||
result.Balance = clineBalanceResult(usd, now)
|
||||
for key, value := range cnBalanceExtraUpdates(PlatformCline, result.Balance, now) {
|
||||
updates[key] = value
|
||||
}
|
||||
}
|
||||
if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil {
|
||||
slog.Warn("cn_quota_persist_failed", "account_id", account.ID, "provider", PlatformCline, "error", err)
|
||||
} else {
|
||||
result.Persisted = true
|
||||
if result.Balance != nil {
|
||||
result.Balance.Persisted = true
|
||||
}
|
||||
}
|
||||
s.applyClineWallets(ctx, account, subscribed, tiers, result.Balance, now)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
const (
|
||||
clineCreditsLowReason = "cline_credits_low"
|
||||
// clinePassNoSubscriptionReason 是探测发现未订阅时的原因;与请求侧的未订阅 / 组织账户
|
||||
// 共用 clinePassUnavailableReason 前缀,单独区分出探测能证明已恢复的这一种。
|
||||
clinePassNoSubscriptionReason = clinePassUnavailableReason + ": no ClinePass subscription"
|
||||
)
|
||||
|
||||
// applyClineWallets 按探测结果设置或解除两个钱包的冷却:
|
||||
// - ClinePass:未订阅,或有窗口用满(≥100%)时冷却到该窗口的重置时间;可用时只解除
|
||||
// 窗口用满与未订阅的冷却(组织账户、所需订阅档位与推理额度上限反映不到用量接口里,
|
||||
// 仍等到期后由请求重新确认);
|
||||
// - 积分:余额低于阈值时冷却,恢复后只解除积分类原因的冷却(组织花费上限无从探测,
|
||||
// 仍等到期后由请求重新确认)。
|
||||
//
|
||||
// 没有重置时间时冷却 2 个检测周期,由下一次探测续期或解除。
|
||||
func (s *CNProviderQuotaService) applyClineWallets(ctx context.Context, account *Account, subscribed bool, tiers []CNQuotaTier, balance *CNProviderBalanceResult, now time.Time) {
|
||||
fallback := now.Add(cnBalanceCheckCooldown(s.cfg))
|
||||
|
||||
switch until, reason := clinePassWalletState(subscribed, tiers, now, fallback); {
|
||||
case until != nil:
|
||||
s.setClineWalletLimit(ctx, account, clinePassRateLimitKey, *until, reason, now)
|
||||
case account.isRateLimitActiveForKey(clinePassRateLimitKey):
|
||||
reason := account.modelRateLimitReason(clinePassRateLimitKey)
|
||||
if strings.HasPrefix(reason, clinePassLimitReason) || strings.HasPrefix(reason, clinePassNoSubscriptionReason) {
|
||||
s.setClineWalletLimit(ctx, account, clinePassRateLimitKey, now, "", now)
|
||||
}
|
||||
}
|
||||
|
||||
if balance == nil {
|
||||
return
|
||||
}
|
||||
threshold := 0.0
|
||||
if s.cfg != nil {
|
||||
threshold = s.cfg.Gateway.CNProviders.BalanceThreshold
|
||||
}
|
||||
switch low := balance.Balance <= 0 || balance.Balance < threshold; {
|
||||
case low:
|
||||
s.setClineWalletLimit(ctx, account, clineCreditsRateLimitKey, fallback,
|
||||
fmt.Sprintf("%s: balance %.4g USD below threshold %.2f", clineCreditsLowReason, balance.Balance, threshold), now)
|
||||
case account.isRateLimitActiveForKey(clineCreditsRateLimitKey):
|
||||
reason := account.modelRateLimitReason(clineCreditsRateLimitKey)
|
||||
if strings.HasPrefix(reason, clineCreditsReason) || strings.HasPrefix(reason, clineCreditsLowReason) {
|
||||
s.setClineWalletLimit(ctx, account, clineCreditsRateLimitKey, now, "", now)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// clinePassWalletState 返回 ClinePass 钱包应冷却到的时间;nil 表示可用。
|
||||
func clinePassWalletState(subscribed bool, tiers []CNQuotaTier, now, fallback time.Time) (*time.Time, string) {
|
||||
if !subscribed {
|
||||
return &fallback, clinePassNoSubscriptionReason
|
||||
}
|
||||
var until *time.Time
|
||||
var exhausted []string
|
||||
for _, tier := range tiers {
|
||||
if tier.UsedPercent < 100 {
|
||||
continue
|
||||
}
|
||||
exhausted = append(exhausted, tier.Window)
|
||||
resetAt := fallback
|
||||
if parsed := parseSchedulingResetAt(tier.ResetAt); parsed != nil && parsed.After(now) && parsed.Sub(now) <= clinePassMaxCooldown {
|
||||
resetAt = *parsed
|
||||
}
|
||||
if until == nil || resetAt.After(*until) {
|
||||
until = &resetAt
|
||||
}
|
||||
}
|
||||
if until == nil {
|
||||
return nil, ""
|
||||
}
|
||||
return until, clinePassLimitReason + ": " + strings.Join(exhausted, ", ") + " window used up"
|
||||
}
|
||||
|
||||
// setClineWalletLimit 写钱包冷却;until 不晚于 now 时即解除。
|
||||
func (s *CNProviderQuotaService) setClineWalletLimit(ctx context.Context, account *Account, scope string, until time.Time, reason string, now time.Time) {
|
||||
setAccountModelRateLimitSnapshot(account, scope, until, reason, now)
|
||||
if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, scope, until, reason); err != nil {
|
||||
slog.Warn("cline_wallet_limit_set_failed", "account_id", account.ID, "scope", scope, "error", err)
|
||||
return
|
||||
}
|
||||
slog.Info("cline_wallet_limit", "account_id", account.ID, "scope", scope, "until", until.UTC(), "cleared", !until.After(now))
|
||||
}
|
||||
@@ -27,6 +27,8 @@ const cnQuotaProbeConcurrency = 4
|
||||
// 克隆自 AccountExpiryService 的 Start/Stop/runOnce + ticker 骨架。
|
||||
// 余额探测仅覆盖有公开余额端点的 kimi / deepseek;智谱无余额端点,仅靠响应式 429/402。
|
||||
// 额度探测覆盖 kimi / zhipu 的 coding plan 账号(deepseek 无 coding 套餐)。
|
||||
// Command Code 账号一次探测同时刷新窗口快照与积分余额,积分低于阈值时同样临时停调。
|
||||
// Cline 账号同样一次探测刷新 ClinePass 窗口与积分,只冷却用完的一边(见 cline_usage.go)。
|
||||
type CNProviderBalanceCheckService struct {
|
||||
accountRepo AccountRepository
|
||||
balanceService *CNProviderBalanceService
|
||||
@@ -154,7 +156,8 @@ func (s *CNProviderBalanceCheckService) runOnce() {
|
||||
}
|
||||
|
||||
// 预算按工作量放大:4 并发 × 15s/批 + payg 每账号 5s,下限 30s 上限 300s。
|
||||
batches := (len(quotaTargets) + cnQuotaProbeConcurrency - 1) / cnQuotaProbeConcurrency
|
||||
accountTargets := s.collectAccountProbeTargets()
|
||||
batches := (len(quotaTargets) + len(accountTargets) + cnQuotaProbeConcurrency - 1) / cnQuotaProbeConcurrency
|
||||
timeout := 30*time.Second + time.Duration(batches)*15*time.Second + time.Duration(len(paygTargets))*5*time.Second
|
||||
if timeout > 300*time.Second {
|
||||
timeout = 300 * time.Second
|
||||
@@ -188,11 +191,106 @@ func (s *CNProviderBalanceCheckService) runOnce() {
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
if len(accountTargets) > 0 {
|
||||
sem := make(chan struct{}, cnQuotaProbeConcurrency)
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
for _, target := range accountTargets {
|
||||
wg.Add(1)
|
||||
go func(t accountProbeTarget) {
|
||||
defer wg.Done()
|
||||
sem <- struct{}{}
|
||||
defer func() { <-sem }()
|
||||
switch t.check(ctx, t.account, threshold) {
|
||||
case cnBalancePaused:
|
||||
mu.Lock()
|
||||
paused++
|
||||
mu.Unlock()
|
||||
case cnBalanceCleared:
|
||||
mu.Lock()
|
||||
cleared++
|
||||
mu.Unlock()
|
||||
}
|
||||
}(target)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
if paused > 0 || cleared > 0 {
|
||||
log.Printf("[CNBalance] paused=%d cleared=%d (threshold=%.2f)", paused, cleared, threshold)
|
||||
}
|
||||
}
|
||||
|
||||
// accountProbeCheck 是走账号接口探测的平台的周期检测函数。
|
||||
type accountProbeCheck func(ctx context.Context, account *Account, threshold float64) cnBalanceCheckOutcome
|
||||
|
||||
// accountProbeTarget 是一个待检测账号与其平台的检测函数。
|
||||
type accountProbeTarget struct {
|
||||
account *Account
|
||||
check accountProbeCheck
|
||||
}
|
||||
|
||||
// collectAccountProbeTargets 选出走账号接口探测的平台上、官方主机上激活的账号,并配上
|
||||
// 该平台的检测函数。一次探测同时刷新窗口快照与积分余额:
|
||||
// - Command Code:手动可调度的账号再按积分余额停调 / 恢复(窗口快照供阈值停调读取);
|
||||
// - Cline:探测中按结果设置或解除两个钱包的冷却,不停整个账号。
|
||||
func (s *CNProviderBalanceCheckService) collectAccountProbeTargets() []accountProbeTarget {
|
||||
if s.quotaService == nil {
|
||||
return nil
|
||||
}
|
||||
var targets []accountProbeTarget
|
||||
for _, p := range []struct {
|
||||
platform string
|
||||
supported func(*Account) bool
|
||||
check accountProbeCheck
|
||||
}{
|
||||
{PlatformCommandCode, (*Account).commandCodeUsageSupported, s.checkCommandCode},
|
||||
{PlatformCline, (*Account).clineAccountAPISupported, s.checkCline},
|
||||
} {
|
||||
accounts, err := s.accountRepo.ListByPlatform(context.Background(), p.platform)
|
||||
if err != nil {
|
||||
log.Printf("[CNBalance] list %s accounts failed: %v", p.platform, err)
|
||||
continue
|
||||
}
|
||||
for i := range accounts {
|
||||
if account := &accounts[i]; account.IsActive() && p.supported(account) {
|
||||
targets = append(targets, accountProbeTarget{account: account, check: p.check})
|
||||
}
|
||||
}
|
||||
}
|
||||
return targets
|
||||
}
|
||||
|
||||
// checkCommandCode 探测 Command Code 窗口与积分;手动可调度的账号再按积分余额停调 / 恢复。
|
||||
func (s *CNProviderBalanceCheckService) checkCommandCode(ctx context.Context, account *Account, threshold float64) cnBalanceCheckOutcome {
|
||||
result, err := s.quotaService.QueryUsage(ctx, account.ID)
|
||||
if err != nil {
|
||||
log.Printf("[CNBalance] command code usage account %d failed: %v", account.ID, err)
|
||||
return cnBalanceNoChange
|
||||
}
|
||||
if result == nil || !result.Success || result.Balance == nil {
|
||||
if result != nil && result.Error != "" {
|
||||
log.Printf("[CNBalance] command code usage account %d error: %s", account.ID, result.Error)
|
||||
}
|
||||
return cnBalanceNoChange
|
||||
}
|
||||
if !account.Schedulable {
|
||||
return cnBalanceNoChange
|
||||
}
|
||||
return s.applyBalanceResult(ctx, account, result.Balance, threshold)
|
||||
}
|
||||
|
||||
// checkCline 探测 Cline 窗口与积分;钱包冷却在探测中按结果设置或解除,不停整个账号。
|
||||
func (s *CNProviderBalanceCheckService) checkCline(ctx context.Context, account *Account, _ float64) cnBalanceCheckOutcome {
|
||||
result, err := s.quotaService.QueryUsage(ctx, account.ID)
|
||||
if err != nil {
|
||||
log.Printf("[CNBalance] cline usage account %d failed: %v", account.ID, err)
|
||||
} else if result != nil && !result.Success && result.Error != "" {
|
||||
log.Printf("[CNBalance] cline usage account %d error: %s", account.ID, result.Error)
|
||||
}
|
||||
return cnBalanceNoChange
|
||||
}
|
||||
|
||||
// probeQuota 探测单个 coding plan 账号的滚动窗口用量并落 extra 快照。
|
||||
// 不在此处做停调/恢复决策:调度阈值评估读取快照统一判定(含暂停账号的续停)。
|
||||
func (s *CNProviderBalanceCheckService) probeQuota(ctx context.Context, accountID int64, platform string) {
|
||||
@@ -224,7 +322,11 @@ func (s *CNProviderBalanceCheckService) checkOne(ctx context.Context, account *A
|
||||
if err != nil || result == nil || !result.Success {
|
||||
return cnBalanceNoChange
|
||||
}
|
||||
return s.applyBalanceResult(ctx, account, result, threshold)
|
||||
}
|
||||
|
||||
// applyBalanceResult 按余额探测结果停调或恢复账号。
|
||||
func (s *CNProviderBalanceCheckService) applyBalanceResult(ctx context.Context, account *Account, result *CNProviderBalanceResult, threshold float64) cnBalanceCheckOutcome {
|
||||
// 双币种(deepseek CNY+USD)任一币种余额达标即可继续调度;仅当全部低于
|
||||
// 阈值(或不可用)才停调。
|
||||
low := !result.Available || allCNBalancesBelowThreshold(result, threshold)
|
||||
@@ -274,15 +376,5 @@ func allCNBalancesBelowThreshold(result *CNProviderBalanceResult, threshold floa
|
||||
|
||||
// cooldown 返回临时停调持续时长(= 2× 检测周期),与响应式 402/429 路径一致。
|
||||
func (s *CNProviderBalanceCheckService) cooldown() time.Duration {
|
||||
minutes := 10
|
||||
if s.cfg != nil {
|
||||
if cfgMin := s.cfg.Gateway.CNProviders.BalanceCheckIntervalMinutes; cfgMin > 0 {
|
||||
minutes = cfgMin
|
||||
}
|
||||
}
|
||||
cooldown := time.Duration(minutes) * time.Minute * 2
|
||||
if cooldown < time.Minute {
|
||||
cooldown = 10 * time.Minute
|
||||
}
|
||||
return cooldown
|
||||
return cnBalanceCheckCooldown(s.cfg)
|
||||
}
|
||||
|
||||
@@ -123,6 +123,12 @@ func (s *CNProviderBalanceService) QueryBalanceForAccount(ctx context.Context, a
|
||||
}
|
||||
|
||||
func (s *CNProviderBalanceService) queryBalanceForAccount(ctx context.Context, account *Account) (*CNProviderBalanceResult, error) {
|
||||
if account.IsCommandCode() {
|
||||
return s.queryCommandCodeBalance(ctx, account)
|
||||
}
|
||||
if account.IsCline() {
|
||||
return s.queryClineBalance(ctx, account)
|
||||
}
|
||||
provider := account.Platform
|
||||
if provider != PlatformKimi && provider != PlatformDeepseek {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_BALANCE_NO_ENDPOINT", "account provider has no balance endpoint")
|
||||
@@ -219,22 +225,7 @@ func (s *CNProviderBalanceService) queryBalanceForAccount(ctx context.Context, a
|
||||
result.Available = available
|
||||
result.Success = true
|
||||
|
||||
balanceUpdates := make([]any, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
balanceUpdates = append(balanceUpdates, map[string]any{
|
||||
"currency": entry.Currency,
|
||||
"balance": entry.Balance,
|
||||
})
|
||||
}
|
||||
updates := map[string]any{
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixBalance): result.Balance,
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixCurrency): result.Currency,
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixAvailable): available,
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixUpdated): now.Format(time.RFC3339),
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixBalances): balanceUpdates,
|
||||
// 余额探测成功即清除响应式 402/429 写下的 balance_low 标记。
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixLow): false,
|
||||
}
|
||||
updates := cnBalanceExtraUpdates(provider, result, now)
|
||||
if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil {
|
||||
slog.Warn("cn_balance_persist_failed", "account_id", account.ID, "provider", provider, "error", err)
|
||||
} else {
|
||||
@@ -243,6 +234,26 @@ func (s *CNProviderBalanceService) queryBalanceForAccount(ctx context.Context, a
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// cnBalanceExtraUpdates 构造 provider 维度的余额快照更新。
|
||||
func cnBalanceExtraUpdates(provider string, result *CNProviderBalanceResult, now time.Time) map[string]any {
|
||||
balanceUpdates := make([]any, 0, len(result.Balances))
|
||||
for _, entry := range result.Balances {
|
||||
balanceUpdates = append(balanceUpdates, map[string]any{
|
||||
"currency": entry.Currency,
|
||||
"balance": entry.Balance,
|
||||
})
|
||||
}
|
||||
return map[string]any{
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixBalance): result.Balance,
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixCurrency): result.Currency,
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixAvailable): result.Available,
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixUpdated): now.Format(time.RFC3339),
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixBalances): balanceUpdates,
|
||||
// 余额探测成功即清除响应式 402/429 写下的 balance_low 标记。
|
||||
cnExtraKey(provider, cnBalanceExtraSuffixLow): false,
|
||||
}
|
||||
}
|
||||
|
||||
// loadPayGAccount 加载 payg 模式的国产供应商账号(余额仅对 payg 有意义;coding 走额度)。
|
||||
func (s *CNProviderBalanceService) loadPayGAccount(ctx context.Context, accountID int64) (*Account, error) {
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
@@ -261,6 +272,20 @@ func validatePayGAccount(account *Account) error {
|
||||
if account == nil {
|
||||
return infraerrors.New(http.StatusNotFound, "CN_BALANCE_ACCOUNT_NOT_FOUND", "account not found")
|
||||
}
|
||||
// Cline 积分余额只查官方主机。
|
||||
if account.IsCline() {
|
||||
if !account.clineAccountAPISupported() {
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_BALANCE_NO_ENDPOINT", "cline balance is only available for api key accounts on the official host")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
// Command Code 积分余额与套餐窗口同源,只查官方主机。
|
||||
if account.IsCommandCode() {
|
||||
if !account.commandCodeUsageSupported() {
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_BALANCE_NO_ENDPOINT", "command code balance is only available for api key accounts on the official host")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if !account.IsCNProvider() {
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_BALANCE_INVALID_PLATFORM", "account is not a CN provider account")
|
||||
}
|
||||
|
||||
@@ -62,9 +62,13 @@ type CNProviderQuotaProbeResult struct {
|
||||
FetchedAt int64 `json:"fetched_at"`
|
||||
Persisted bool `json:"persisted"`
|
||||
Error string `json:"error,omitempty"`
|
||||
|
||||
// Balance 为同一次探测得到的余额(Command Code 积分与窗口同源),其余供应商为空。
|
||||
Balance *CNProviderBalanceResult `json:"balance,omitempty"`
|
||||
}
|
||||
|
||||
// CNProviderQuotaService 探测 Kimi / Zhipu Coding Plan 的滚动窗口用量。
|
||||
// CNProviderQuotaService 探测 Kimi / Zhipu / MiniMax Coding Plan、OpenCode Go、
|
||||
// Command Code 与 Cline(ClinePass)的滚动窗口用量。
|
||||
type CNProviderQuotaService struct {
|
||||
accountRepo AccountRepository
|
||||
proxyRepo ProxyRepository
|
||||
@@ -131,6 +135,12 @@ func (s *CNProviderQuotaService) QueryUsageForAccount(ctx context.Context, accou
|
||||
|
||||
func (s *CNProviderQuotaService) queryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
|
||||
provider := account.GetCodingPlanProvider()
|
||||
if provider == PlatformCommandCode {
|
||||
return s.queryCommandCodeUsage(ctx, account)
|
||||
}
|
||||
if provider == PlatformCline {
|
||||
return s.queryClineUsage(ctx, account)
|
||||
}
|
||||
if provider != PlatformKimi && provider != PlatformZhipu && provider != PlatformMiniMax && provider != PlatformOpenCodeGo {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "account is not a kimi/zhipu/minimax coding plan or opencode go account")
|
||||
}
|
||||
@@ -290,6 +300,18 @@ func validateCodingPlanAccount(account *Account) error {
|
||||
if account.IsOpenCodeGoPlan() {
|
||||
return nil
|
||||
}
|
||||
if account.IsCommandCode() {
|
||||
if !account.commandCodeUsageSupported() {
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "command code usage is only available for api key accounts on the official host")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if account.IsCline() {
|
||||
if !account.clineAccountAPISupported() {
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "cline usage is only available for api key accounts on the official host")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if account.IsOpenCodeGo() {
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "opencode zen accounts have no subscription quota window")
|
||||
}
|
||||
|
||||
@@ -668,8 +668,9 @@ func TestAdaptiveProtocolBaseURLs(t *testing.T) {
|
||||
}{
|
||||
{"kimi payg", PlatformKimi, AccountModePayG, DefaultKimiPayGBaseURL, DefaultKimiPayGAnthropicBaseURL, DefaultKimiPayGBaseURL},
|
||||
{"kimi coding", PlatformKimi, AccountModeCoding, DefaultKimiCodingBaseURL, DefaultKimiCodingAnthropicBaseURL, DefaultKimiCodingBaseURL},
|
||||
{"zhipu payg", PlatformZhipu, AccountModePayG, DefaultZhipuPayGBaseURL, DefaultZhipuAnthropicBaseURL, DefaultZhipuPayGBaseURL},
|
||||
{"zhipu coding", PlatformZhipu, AccountModeCoding, DefaultZhipuCodingBaseURL, DefaultZhipuAnthropicBaseURL, DefaultZhipuCodingBaseURL},
|
||||
// 智谱没有原生 Responses 端点:profile 不登记 responses 条目,默认基址为空。
|
||||
{"zhipu payg", PlatformZhipu, AccountModePayG, DefaultZhipuPayGBaseURL, DefaultZhipuAnthropicBaseURL, ""},
|
||||
{"zhipu coding", PlatformZhipu, AccountModeCoding, DefaultZhipuCodingBaseURL, DefaultZhipuAnthropicBaseURL, ""},
|
||||
{"deepseek", PlatformDeepseek, AccountModePayG, DefaultDeepseekBaseURL, DefaultDeepseekAnthropicBaseURL, DefaultDeepseekBaseURL},
|
||||
{"minimax payg", PlatformMiniMax, AccountModePayG, DefaultMiniMaxBaseURL, DefaultMiniMaxAnthropicBaseURL, DefaultMiniMaxBaseURL},
|
||||
{"minimax coding", PlatformMiniMax, AccountModeCoding, DefaultMiniMaxBaseURL, DefaultMiniMaxAnthropicBaseURL, DefaultMiniMaxBaseURL},
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
package service
|
||||
|
||||
// Command Code Provider API(https://commandcode.ai/docs/provider):同一 API Key 下
|
||||
// 按模型分流到三种原生端点,基址 https://api.commandcode.ai/provider/v1:
|
||||
// - Claude 系模型只在 /messages(Anthropic Messages)提供,发往其他端点返回 400;
|
||||
// - OpenAI 与开源模型提供 /chat/completions,多数也提供 /responses。
|
||||
// 每个模型支持的端点由 /models 的 supported_endpoints 给出(公开接口),网关据此在
|
||||
// 模型支持入站协议时同协议直通(见 model_protocol_catalog.go);目录就绪前使用下面的
|
||||
// 内置规则。模型名可带厂商前缀(如 deepseek/deepseek-v4-flash),原样发往上游。
|
||||
|
||||
// DefaultCommandCodeTestModel 是管理端连接测试未指定模型时使用的模型(Chat Completions)。
|
||||
const DefaultCommandCodeTestModel = "deepseek/deepseek-v4-flash"
|
||||
|
||||
// DefaultCommandCodeProtocolRules 是 Command Code 的内置分流表:Claude 只走 Anthropic;
|
||||
// GPT 同时支持 Responses 与 Chat Completions(入站同协议直通,其余入站首选 Responses);
|
||||
// 其余模型只保证 Chat Completions,是否支持 Responses 以上游模型目录为准。
|
||||
func DefaultCommandCodeProtocolRules() []ProtocolRule {
|
||||
return []ProtocolRule{
|
||||
{Pattern: "claude-*", Protocol: APIProtocolAnthropic},
|
||||
{Pattern: "gpt-*", Protocol: APIProtocolResponses, Protocols: []string{APIProtocolResponses, APIProtocolChatCompletions}},
|
||||
}
|
||||
}
|
||||
|
||||
// IsCommandCode 报告账号是否为 Command Code 平台账号。
|
||||
func (a *Account) IsCommandCode() bool {
|
||||
return a != nil && a.Platform == PlatformCommandCode
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// commandCodeTestAccount 只带 api_key:端点与分流规则全部来自 provider profile。
|
||||
func commandCodeTestAccount(id int64) *Account {
|
||||
return &Account{
|
||||
ID: id,
|
||||
Name: "cc",
|
||||
Platform: PlatformCommandCode,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{"api_key": "user_test_key"},
|
||||
}
|
||||
}
|
||||
|
||||
func TestCommandCodeProfileRoutesByModel(t *testing.T) {
|
||||
account := commandCodeTestAccount(1)
|
||||
require.True(t, account.IsMultiProtocolAPIKey())
|
||||
require.True(t, account.routesByModel())
|
||||
require.False(t, account.RoutesProtocolByInbound())
|
||||
require.Equal(t, APIProtocolAdaptive, account.GetAPIProtocol())
|
||||
require.Equal(t, DefaultCommandCodeBaseURL, account.GetOpenAIBaseURL())
|
||||
require.Equal(t, DefaultCommandCodeAnthropicBaseURL, account.GetCNProtocolBaseURL(APIProtocolAnthropic))
|
||||
require.Equal(t, DefaultCommandCodeBaseURL, account.GetCNProtocolBaseURL(APIProtocolResponses))
|
||||
|
||||
// 未取到模型目录时按内置规则:Claude 只走 Anthropic;GPT 支持 Responses 与 Chat
|
||||
// Completions(同协议直通,Anthropic 入站走首选 Responses);其余只保证 Chat Completions。
|
||||
for _, inbound := range []string{APIProtocolChatCompletions, APIProtocolResponses, APIProtocolAnthropic} {
|
||||
require.Equal(t, APIProtocolAnthropic, resolveUpstreamProtocol(account, inbound, "claude-sonnet-4-6", nil), inbound)
|
||||
require.Equal(t, APIProtocolChatCompletions, resolveUpstreamProtocol(account, inbound, "deepseek/deepseek-v4-flash", nil), inbound)
|
||||
}
|
||||
require.Equal(t, APIProtocolResponses, resolveUpstreamProtocol(account, APIProtocolResponses, "gpt-5.5", nil))
|
||||
require.Equal(t, APIProtocolChatCompletions, resolveUpstreamProtocol(account, APIProtocolChatCompletions, "gpt-5.5", nil))
|
||||
require.Equal(t, APIProtocolResponses, resolveUpstreamProtocol(account, APIProtocolAnthropic, "gpt-5.5", nil))
|
||||
|
||||
// 显式协议优先于内置规则。
|
||||
account.Credentials["api_protocol"] = APIProtocolChatCompletions
|
||||
require.Equal(t, APIProtocolChatCompletions, resolveUpstreamProtocol(account, APIProtocolAnthropic, "claude-sonnet-4-6", nil))
|
||||
}
|
||||
|
||||
func TestCommandCodeNativeAnthropicTargetURL(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{}
|
||||
targetURL, err := svc.nativeAnthropicTargetURL(commandCodeTestAccount(2))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://api.commandcode.ai/provider/v1/messages", targetURL)
|
||||
}
|
||||
|
||||
func TestAccountTestService_CommandCodeClaudeUsesAnthropicMessages(t *testing.T) {
|
||||
account := commandCodeTestAccount(501)
|
||||
svc, upstream := adaptiveCNAccountTestService(account, adaptiveCNAnthropicTestResponse())
|
||||
c, recorder := newTestContext()
|
||||
|
||||
err := svc.TestAccountConnection(c, account.ID, "claude-sonnet-4-6", "hi", AccountTestModeDefault)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, upstream.requests, 1)
|
||||
require.Equal(t, "https://api.commandcode.ai/provider/v1/messages", upstream.requests[0].URL.String())
|
||||
require.Equal(t, "user_test_key", upstream.requests[0].Header.Get("x-api-key"))
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
|
||||
}
|
||||
|
||||
// 连接测试与转发共用端点拼接:按模型分流的平台,Anthropic 基址带 /v1 也合法。
|
||||
func TestAccountTestService_CommandCodeAnthropicBaseWithVersion(t *testing.T) {
|
||||
account := commandCodeTestAccount(504)
|
||||
account.Credentials["api_base_urls"] = map[string]any{APIProtocolAnthropic: "https://api.commandcode.ai/provider/v1"}
|
||||
svc, upstream := adaptiveCNAccountTestService(account, adaptiveCNAnthropicTestResponse())
|
||||
c, recorder := newTestContext()
|
||||
|
||||
err := svc.TestAccountConnection(c, account.ID, "claude-sonnet-4-6", "hi", AccountTestModeDefault)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, upstream.requests, 1)
|
||||
require.Equal(t, "https://api.commandcode.ai/provider/v1/messages", upstream.requests[0].URL.String())
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
|
||||
|
||||
target, err := (&OpenAIGatewayService{cfg: rawChatCompletionsTestConfig()}).nativeAnthropicTargetURL(account)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, upstream.requests[0].URL.String(), target)
|
||||
}
|
||||
|
||||
// 连接测试与网关同一判定:模型目录决定协议时,测试也按目录选择端点。
|
||||
func TestAccountTestService_CommandCodeUsesModelCatalog(t *testing.T) {
|
||||
base := fmt.Sprintf("http://cc-test-%d.example/provider/v1", time.Now().UnixNano())
|
||||
account := commandCodeTestAccount(505)
|
||||
account.Credentials["api_base_urls"] = map[string]any{
|
||||
APIProtocolChatCompletions: base,
|
||||
APIProtocolResponses: base,
|
||||
}
|
||||
url := buildOpenAIModelsURL(base)
|
||||
upstreamModelProtocols.store(modelProtocolCatalogKey(account, url),
|
||||
map[string][]string{"vendor/responses-only": {APIProtocolResponses}}, nil, time.Now())
|
||||
svc, upstream := adaptiveCNAccountTestService(account, adaptiveCNResponsesTestResponse())
|
||||
svc.openaiGatewayService = &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
|
||||
c, recorder := newTestContext()
|
||||
|
||||
err := svc.TestAccountConnection(c, account.ID, "vendor/responses-only", "hi", AccountTestModeDefault)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, upstream.requests, 1)
|
||||
require.Equal(t, base+"/responses", upstream.requests[0].URL.String())
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
|
||||
}
|
||||
|
||||
func TestAccountTestService_CommandCodeGPTUsesResponses(t *testing.T) {
|
||||
account := commandCodeTestAccount(502)
|
||||
svc, upstream := adaptiveCNAccountTestService(account, adaptiveCNResponsesTestResponse())
|
||||
c, recorder := newTestContext()
|
||||
|
||||
err := svc.TestAccountConnection(c, account.ID, "gpt-5.5", "hi", AccountTestModeDefault)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, upstream.requests, 1)
|
||||
require.Equal(t, "https://api.commandcode.ai/provider/v1/responses", upstream.requests[0].URL.String())
|
||||
require.Equal(t, "Bearer user_test_key", upstream.requests[0].Header.Get("Authorization"))
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
|
||||
}
|
||||
|
||||
func TestAccountTestService_CommandCodeEmptyModelUsesProfileDefault(t *testing.T) {
|
||||
account := commandCodeTestAccount(503)
|
||||
svc, upstream := adaptiveCNAccountTestService(account, adaptiveCNChatTestResponse())
|
||||
c, recorder := newTestContext()
|
||||
|
||||
err := svc.TestAccountConnection(c, account.ID, "", "hi", AccountTestModeDefault)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, upstream.requests, 1)
|
||||
require.Equal(t, "https://api.commandcode.ai/provider/v1/chat/completions", upstream.requests[0].URL.String())
|
||||
require.Equal(t, DefaultCommandCodeTestModel, gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
|
||||
}
|
||||
|
||||
// 仅在 profile 登记、按入站协议分流的新供应商:adaptive 测试走原生端点探测,
|
||||
// 跳过 profile 未提供的 Anthropic 端点,不再落入 Claude 测试路径。
|
||||
func TestAccountTestService_ByInboundProviderSkipsMissingNativeEndpoints(t *testing.T) {
|
||||
const platform = "test_inbound_only_chat"
|
||||
providerProfiles[platform] = &ProviderProfile{
|
||||
Platform: platform, DefaultMode: AccountModePayG, Routing: ProviderRoutingByInbound,
|
||||
DefaultTestModel: "inbound-test-model",
|
||||
Modes: map[string]ProviderEndpoints{AccountModePayG: {BaseURLs: map[string]string{
|
||||
APIProtocolChatCompletions: "https://inbound.example.com/v1",
|
||||
}}},
|
||||
}
|
||||
t.Cleanup(func() { delete(providerProfiles, platform) })
|
||||
|
||||
account := &Account{
|
||||
ID: 601, Name: "inbound", Platform: platform, Type: AccountTypeAPIKey, Status: StatusActive, Concurrency: 1,
|
||||
Credentials: map[string]any{"api_key": "sk-inbound", "api_protocol": APIProtocolAdaptive},
|
||||
}
|
||||
svc, upstream := adaptiveCNAccountTestService(account, adaptiveCNChatTestResponse())
|
||||
c, recorder := newTestContext()
|
||||
|
||||
err := svc.TestAccountConnection(c, account.ID, "", "hi", AccountTestModeDefault)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, upstream.requests, 1)
|
||||
require.Equal(t, "https://inbound.example.com/v1/chat/completions", upstream.requests[0].URL.String())
|
||||
require.Equal(t, "inbound-test-model", gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
|
||||
}
|
||||
|
||||
// 官方主机不做上游计费探测:中转站的计费接口在官方 API 上不存在,只会把 Key 发往官方主机。
|
||||
func TestCommandCodeOfficialHostSkipsUpstreamBillingProbe(t *testing.T) {
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI(DefaultCommandCodeBaseURL))
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI(DefaultCommandCodeAnthropicBaseURL))
|
||||
}
|
||||
@@ -0,0 +1,511 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
)
|
||||
|
||||
// Command Code 的积分与用量窗口来自 /alpha 下四个只读接口(官方 CLI /usage 同源,
|
||||
// 社区插件实测复核),全部使用账号 API Key(Bearer):
|
||||
//
|
||||
// GET /alpha/whoami -> {user, org};个人账户 org 为 null
|
||||
// GET /alpha/billing/credits[?orgId=] -> {credits:{monthlyCredits, purchasedCredits, freeCredits},
|
||||
// windowLimits:{limited, fiveHour, weekly}}
|
||||
// GET /alpha/billing/subscriptions[?orgId=] -> {data:{planId, currentPeriodStart, currentPeriodEnd}}
|
||||
// GET /alpha/usage/summary[?orgId=][&since=] -> {totalCost, totalMonthlyCredits}
|
||||
//
|
||||
// 积分单位为美元。windowLimits 与 credits 平级;fiveHour / weekly 的 used、cap 为积分,
|
||||
// resetAt 为 epoch 毫秒;limited=false(按量计费)时没有滚动窗口。月度窗口上游只给
|
||||
// 剩余量(credits.monthlyCredits),上限按「本周期已用 + 剩余」还原。个人账户不能带
|
||||
// orgId 参数(空值返回 400),只有组织账户才带。
|
||||
//
|
||||
// 订阅套餐的滚动窗口只限制套餐内积分:有充值积分时请求跳过窗口检查,此时窗口快照
|
||||
// 不参与阈值停调与 429 冷却(见 commandCodeWindowsBinding),已有的窗口冷却在探测到
|
||||
// 充值积分后解除(见 clearCommandCodeWindowCooldown)。
|
||||
|
||||
const (
|
||||
commandCodeAPIBase = "https://api.commandcode.ai"
|
||||
|
||||
// commandCodeExtraSuffixPurchasedCredits 记录充值积分,供阈值停调与 429 冷却判断
|
||||
// 滚动窗口是否生效。
|
||||
commandCodeExtraSuffixPurchasedCredits = "purchased_credits"
|
||||
)
|
||||
|
||||
var errCommandCodeInvalidResponse = errors.New("invalid command code usage response")
|
||||
|
||||
// commandCodePlanNames 与官方 CLI 的 planId 展示名一致;未收录的原样展示。
|
||||
var commandCodePlanNames = map[string]string{
|
||||
"individual-go": "Go",
|
||||
"individual-goat": "GOAT",
|
||||
"individual-pro": "Pro",
|
||||
"individual-pro-v1": "Pro",
|
||||
"individual-provider": "Provider",
|
||||
"individual-max": "Max",
|
||||
"individual-ultra": "Ultra",
|
||||
"teams-pro": "Teams Pro",
|
||||
}
|
||||
|
||||
// commandCodeUsageSupported 报告账号能否查询 Command Code 用量:API Key 账号且推理
|
||||
// 端点指向官方主机。自定义中转的 Key 不发往官方用量接口。
|
||||
func (a *Account) commandCodeUsageSupported() bool {
|
||||
return a.IsCommandCode() && a.Type == AccountTypeAPIKey && isOfficialCommandCodeHost(a.GetOpenAIBaseURL())
|
||||
}
|
||||
|
||||
// commandCodeNumber 兼容数字与数字字符串。
|
||||
type commandCodeNumber float64
|
||||
|
||||
func (n *commandCodeNumber) UnmarshalJSON(data []byte) error {
|
||||
data = bytes.TrimSpace(data)
|
||||
if len(data) == 0 || bytes.Equal(data, []byte("null")) {
|
||||
*n = 0
|
||||
return nil
|
||||
}
|
||||
if data[0] == '"' {
|
||||
var s string
|
||||
if err := json.Unmarshal(data, &s); err != nil {
|
||||
return err
|
||||
}
|
||||
v, err := strconv.ParseFloat(strings.TrimSpace(s), 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*n = commandCodeNumber(v)
|
||||
return nil
|
||||
}
|
||||
var v float64
|
||||
if err := json.Unmarshal(data, &v); err != nil {
|
||||
return err
|
||||
}
|
||||
*n = commandCodeNumber(v)
|
||||
return nil
|
||||
}
|
||||
|
||||
type commandCodeWindow struct {
|
||||
Used commandCodeNumber `json:"used"`
|
||||
Cap commandCodeNumber `json:"cap"`
|
||||
ResetAt commandCodeNumber `json:"resetAt"`
|
||||
}
|
||||
|
||||
type commandCodeCredits struct {
|
||||
MonthlyCredits commandCodeNumber `json:"monthlyCredits"`
|
||||
PurchasedCredits commandCodeNumber `json:"purchasedCredits"`
|
||||
FreeCredits commandCodeNumber `json:"freeCredits"`
|
||||
}
|
||||
|
||||
type commandCodeCreditsPayload struct {
|
||||
Credits *commandCodeCredits `json:"credits"`
|
||||
WindowLimits struct {
|
||||
Limited bool `json:"limited"`
|
||||
FiveHour *commandCodeWindow `json:"fiveHour"`
|
||||
Weekly *commandCodeWindow `json:"weekly"`
|
||||
} `json:"windowLimits"`
|
||||
}
|
||||
|
||||
type commandCodeSubscriptionPayload struct {
|
||||
Data struct {
|
||||
PlanID string `json:"planId"`
|
||||
CurrentPeriodStart string `json:"currentPeriodStart"`
|
||||
CurrentPeriodEnd string `json:"currentPeriodEnd"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
type commandCodeSummaryPayload struct {
|
||||
TotalCost *commandCodeNumber `json:"totalCost"`
|
||||
TotalMonthlyCredits *commandCodeNumber `json:"totalMonthlyCredits"`
|
||||
}
|
||||
|
||||
// commandCodeUsage 是一次取数的汇总;subscription / summary 取数失败时为 nil,并置
|
||||
// periodFetchFailed:此时无从判断月度窗口是否存在,不能当作没有月度窗口。
|
||||
type commandCodeUsage struct {
|
||||
credits commandCodeCreditsPayload
|
||||
subscription *commandCodeSubscriptionPayload
|
||||
summary *commandCodeSummaryPayload
|
||||
periodFetchFailed bool
|
||||
}
|
||||
|
||||
// commandCodeHTTPError 是 /alpha 接口的非 2xx 响应。
|
||||
type commandCodeHTTPError struct {
|
||||
status int
|
||||
body string
|
||||
}
|
||||
|
||||
func (e *commandCodeHTTPError) Error() string {
|
||||
return fmt.Sprintf("HTTP %d: %s", e.status, e.body)
|
||||
}
|
||||
|
||||
func (e *commandCodeHTTPError) unauthorized() bool {
|
||||
return e.status == http.StatusUnauthorized || e.status == http.StatusForbidden
|
||||
}
|
||||
|
||||
func isCommandCodeAuthError(err error) bool {
|
||||
var httpErr *commandCodeHTTPError
|
||||
return errors.As(err, &httpErr) && httpErr.unauthorized()
|
||||
}
|
||||
|
||||
type commandCodeUsageClient struct {
|
||||
cfg *config.Config
|
||||
httpUpstream HTTPUpstream
|
||||
account *Account
|
||||
apiKey string
|
||||
proxyURL string
|
||||
}
|
||||
|
||||
// get 请求一个 /alpha 接口并解码 JSON。
|
||||
func (c *commandCodeUsageClient) get(ctx context.Context, path string, query url.Values, out any) error {
|
||||
target := commandCodeAPIBase + path
|
||||
if len(query) > 0 {
|
||||
target += "?" + query.Encode()
|
||||
}
|
||||
validatedURL, err := cnValidateProbeURL(c.cfg, target)
|
||||
if err != nil {
|
||||
return infraerrors.New(http.StatusForbidden, "CN_QUOTA_URL_REJECTED", err.Error())
|
||||
}
|
||||
callCtx, cancel := context.WithTimeout(ctx, cnQuotaUpstreamTimeout)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(callCtx, http.MethodGet, validatedURL, nil)
|
||||
if err != nil {
|
||||
return infraerrors.Newf(http.StatusInternalServerError, "CN_QUOTA_REQUEST_BUILD_FAILED", "build request: %v", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
c.account.ApplyHeaderOverrides(req.Header)
|
||||
|
||||
resp, err := c.httpUpstream.Do(req, c.proxyURL, c.account.ID, maxInt(c.account.Concurrency, 1))
|
||||
if err != nil {
|
||||
return fmt.Errorf("upstream request failed: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, cnQuotaMaxBodyBytes))
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return &commandCodeHTTPError{status: resp.StatusCode, body: truncate(strings.TrimSpace(string(body)), 240)}
|
||||
}
|
||||
if err := json.Unmarshal(body, out); err != nil {
|
||||
return fmt.Errorf("%w: %s: %v", errCommandCodeInvalidResponse, path, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// fetch 取积分与窗口;withPeriod 时再取订阅周期与本周期用量以还原月度窗口。
|
||||
// whoami 用来取 orgId,任何失败都整体返回:无从判断是否组织账户时,按个人账户查会把
|
||||
// 组织账号查成成员的个人钱包,并据此按余额停调。
|
||||
func (c *commandCodeUsageClient) fetch(ctx context.Context, withPeriod bool) (*commandCodeUsage, error) {
|
||||
var whoami struct {
|
||||
Org *struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"org"`
|
||||
}
|
||||
query := url.Values{}
|
||||
if err := c.get(ctx, "/alpha/whoami", nil, &whoami); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if whoami.Org != nil && strings.TrimSpace(whoami.Org.ID) != "" {
|
||||
query.Set("orgId", strings.TrimSpace(whoami.Org.ID))
|
||||
}
|
||||
|
||||
usage := &commandCodeUsage{}
|
||||
if err := c.get(ctx, "/alpha/billing/credits", query, &usage.credits); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if usage.credits.Credits == nil {
|
||||
return nil, fmt.Errorf("%w: credits missing", errCommandCodeInvalidResponse)
|
||||
}
|
||||
if !withPeriod {
|
||||
return usage, nil
|
||||
}
|
||||
|
||||
var subscription commandCodeSubscriptionPayload
|
||||
if err := c.get(ctx, "/alpha/billing/subscriptions", query, &subscription); err != nil {
|
||||
if isCommandCodeAuthError(err) {
|
||||
return nil, err
|
||||
}
|
||||
slog.Warn("command_code_subscription_failed", "account_id", c.account.ID, "error", err)
|
||||
usage.periodFetchFailed = true
|
||||
return usage, nil
|
||||
}
|
||||
usage.subscription = &subscription
|
||||
periodStart := strings.TrimSpace(subscription.Data.CurrentPeriodStart)
|
||||
if periodStart == "" {
|
||||
return usage, nil
|
||||
}
|
||||
summaryQuery := url.Values{}
|
||||
for key, values := range query {
|
||||
summaryQuery[key] = values
|
||||
}
|
||||
summaryQuery.Set("since", periodStart)
|
||||
var summary commandCodeSummaryPayload
|
||||
if err := c.get(ctx, "/alpha/usage/summary", summaryQuery, &summary); err != nil {
|
||||
if isCommandCodeAuthError(err) {
|
||||
return nil, err
|
||||
}
|
||||
slog.Warn("command_code_usage_summary_failed", "account_id", c.account.ID, "error", err)
|
||||
usage.periodFetchFailed = true
|
||||
return usage, nil
|
||||
}
|
||||
usage.summary = &summary
|
||||
return usage, nil
|
||||
}
|
||||
|
||||
func isApplicationError(err error) bool {
|
||||
var appErr *infraerrors.ApplicationError
|
||||
return errors.As(err, &appErr)
|
||||
}
|
||||
|
||||
func commandCodePercent(used, limit float64) float64 {
|
||||
percent := used / limit * 100
|
||||
if percent < 0 {
|
||||
percent = 0
|
||||
}
|
||||
return math.Round(percent*100) / 100
|
||||
}
|
||||
|
||||
// tiers 返回 5h / weekly(限额套餐)与月度窗口。
|
||||
func (u *commandCodeUsage) tiers() []CNQuotaTier {
|
||||
var tiers []CNQuotaTier
|
||||
limits := u.credits.WindowLimits
|
||||
if limits.Limited {
|
||||
for _, item := range []struct {
|
||||
window string
|
||||
limit *commandCodeWindow
|
||||
}{
|
||||
{window: "5h", limit: limits.FiveHour},
|
||||
{window: "weekly", limit: limits.Weekly},
|
||||
} {
|
||||
if item.limit == nil || item.limit.Cap <= 0 {
|
||||
continue
|
||||
}
|
||||
tiers = append(tiers, CNQuotaTier{
|
||||
Window: item.window,
|
||||
UsedPercent: commandCodePercent(float64(item.limit.Used), float64(item.limit.Cap)),
|
||||
ResetAt: cnMillisToRFC3339(int64(item.limit.ResetAt)),
|
||||
})
|
||||
}
|
||||
}
|
||||
if u.summary != nil && u.subscription != nil {
|
||||
used := 0.0
|
||||
switch {
|
||||
case u.summary.TotalMonthlyCredits != nil:
|
||||
used = float64(*u.summary.TotalMonthlyCredits)
|
||||
case u.summary.TotalCost != nil:
|
||||
used = float64(*u.summary.TotalCost)
|
||||
}
|
||||
if limit := used + float64(u.credits.Credits.MonthlyCredits); limit > 0 {
|
||||
tiers = append(tiers, CNQuotaTier{
|
||||
Window: "monthly",
|
||||
UsedPercent: commandCodePercent(used, limit),
|
||||
ResetAt: cnNormalizeResetTime(u.subscription.Data.CurrentPeriodEnd),
|
||||
})
|
||||
}
|
||||
}
|
||||
return tiers
|
||||
}
|
||||
|
||||
func (u *commandCodeUsage) planName() string {
|
||||
if u.subscription == nil {
|
||||
return ""
|
||||
}
|
||||
planID := strings.TrimSpace(u.subscription.Data.PlanID)
|
||||
if name, ok := commandCodePlanNames[planID]; ok {
|
||||
return name
|
||||
}
|
||||
return planID
|
||||
}
|
||||
|
||||
// balance 返回可用积分合计(套餐剩余 + 充值 + 赠送,美元)。
|
||||
func (u *commandCodeUsage) balance(now time.Time) *CNProviderBalanceResult {
|
||||
credits := u.credits.Credits
|
||||
total := float64(credits.MonthlyCredits + credits.PurchasedCredits + credits.FreeCredits)
|
||||
total = math.Round(total*10000) / 10000
|
||||
return &CNProviderBalanceResult{
|
||||
Provider: PlatformCommandCode,
|
||||
Success: true,
|
||||
Balance: total,
|
||||
Currency: "USD",
|
||||
Balances: []CNProviderBalanceEntry{{Currency: "USD", Balance: total}},
|
||||
Available: true,
|
||||
StatusCode: http.StatusOK,
|
||||
FetchedAt: now.Unix(),
|
||||
}
|
||||
}
|
||||
|
||||
// clearCommandCodeWindowCooldown 在探测到充值积分时解除窗口用满的临时停调:充值积分
|
||||
// 不受滚动窗口限制,上游已允许继续请求。只清除 commandCodeUsageLimitReason 前缀的停调,
|
||||
// 积分不足、花费上限、管理员规则等其他原因的停调不受影响。
|
||||
func clearCommandCodeWindowCooldown(ctx context.Context, repo AccountRepository, account *Account, usage *commandCodeUsage, now time.Time) bool {
|
||||
if usage == nil || usage.credits.Credits == nil || usage.credits.Credits.PurchasedCredits <= 0 {
|
||||
return false
|
||||
}
|
||||
if account.TempUnschedulableUntil == nil || !now.Before(*account.TempUnschedulableUntil) ||
|
||||
!strings.HasPrefix(account.TempUnschedulableReason, commandCodeUsageLimitReason) {
|
||||
return false
|
||||
}
|
||||
if err := repo.ClearTempUnschedulable(ctx, account.ID); err != nil {
|
||||
slog.Warn("command_code_usage_limit_clear_failed", "account_id", account.ID, "error", err)
|
||||
return false
|
||||
}
|
||||
slog.Info("command_code_usage_limit_cleared", "account_id", account.ID,
|
||||
"purchased_credits", float64(usage.credits.Credits.PurchasedCredits))
|
||||
return true
|
||||
}
|
||||
|
||||
// commandCodeWindowsBinding 报告账号快照中的滚动 / 月度窗口是否限制调度:
|
||||
// 有充值积分时上游跳过窗口检查,窗口用满也不会被拒绝。
|
||||
func commandCodeWindowsBinding(extra map[string]any) bool {
|
||||
purchased, ok := cnParseF64(extra[cnExtraKey(PlatformCommandCode, commandCodeExtraSuffixPurchasedCredits)])
|
||||
return !ok || purchased <= 0
|
||||
}
|
||||
|
||||
// commandCodeUsageExtraUpdates 构造窗口 + 余额快照。缺席的窗口写 nil,避免套餐
|
||||
// 变化后沿用旧快照;窗口尚未开始时没有重置时间,同样清掉上一个窗口的值;月度窗口
|
||||
// 取数失败时保留旧快照;传 nil tiers 时只写余额。
|
||||
func commandCodeUsageExtraUpdates(usage *commandCodeUsage, tiers []CNQuotaTier, balance *CNProviderBalanceResult, now time.Time, withTiers bool) map[string]any {
|
||||
updates := cnBalanceExtraUpdates(PlatformCommandCode, balance, now)
|
||||
updates[cnExtraKey(PlatformCommandCode, commandCodeExtraSuffixPurchasedCredits)] = float64(usage.credits.Credits.PurchasedCredits)
|
||||
if !withTiers {
|
||||
return updates
|
||||
}
|
||||
for key, value := range cnQuotaExtraUpdates(PlatformCommandCode, tiers, now) {
|
||||
updates[key] = value
|
||||
}
|
||||
present := make(map[string]CNQuotaTier, len(tiers))
|
||||
for _, tier := range tiers {
|
||||
present[tier.Window] = tier
|
||||
}
|
||||
for _, window := range []struct {
|
||||
name, used, reset string
|
||||
}{
|
||||
{"5h", cnExtraSuffix5hUsed, cnExtraSuffix5hReset},
|
||||
{"weekly", cnExtraSuffixWeeklyUsed, cnExtraSuffixWeeklyReset},
|
||||
{"monthly", cnExtraSuffixMonthlyUsed, cnExtraSuffixMonthlyReset},
|
||||
} {
|
||||
tier, ok := present[window.name]
|
||||
if !ok && window.name == "monthly" && usage.periodFetchFailed {
|
||||
continue
|
||||
}
|
||||
if !ok {
|
||||
updates[cnExtraKey(PlatformCommandCode, window.used)] = nil
|
||||
}
|
||||
if tier.ResetAt == "" {
|
||||
updates[cnExtraKey(PlatformCommandCode, window.reset)] = nil
|
||||
}
|
||||
}
|
||||
return updates
|
||||
}
|
||||
|
||||
func newCommandCodeUsageClient(cfg *config.Config, httpUpstream HTTPUpstream, account *Account, proxyURL string) *commandCodeUsageClient {
|
||||
return &commandCodeUsageClient{
|
||||
cfg: cfg,
|
||||
httpUpstream: httpUpstream,
|
||||
account: account,
|
||||
apiKey: strings.TrimSpace(account.GetCNAPIKey()),
|
||||
proxyURL: proxyURL,
|
||||
}
|
||||
}
|
||||
|
||||
// commandCodeFetchFailure 把取数错误转换为探测失败文案;ok=false 表示应作为错误返回。
|
||||
func commandCodeFetchFailure(err error) (status int, message string, ok bool) {
|
||||
var httpErr *commandCodeHTTPError
|
||||
switch {
|
||||
case errors.As(err, &httpErr) && httpErr.unauthorized():
|
||||
return httpErr.status, fmt.Sprintf("Authentication failed (HTTP %d)", httpErr.status), true
|
||||
case errors.As(err, &httpErr):
|
||||
return httpErr.status, fmt.Sprintf("API error (HTTP %d): %s", httpErr.status, httpErr.body), true
|
||||
case errors.Is(err, errCommandCodeInvalidResponse):
|
||||
return http.StatusOK, "Invalid usage response", true
|
||||
default:
|
||||
return 0, "", false
|
||||
}
|
||||
}
|
||||
|
||||
// queryCommandCodeUsage 查询 Command Code 积分与窗口,落窗口与余额快照。
|
||||
func (s *CNProviderQuotaService) queryCommandCodeUsage(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
|
||||
client := newCommandCodeUsageClient(s.cfg, s.httpUpstream, account, s.resolveProxyURL(ctx, account))
|
||||
if client.apiKey == "" {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NO_APIKEY", "account api_key is empty")
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
result := &CNProviderQuotaProbeResult{
|
||||
Provider: PlatformCommandCode,
|
||||
Source: "command_code_usage",
|
||||
FetchedAt: now.Unix(),
|
||||
}
|
||||
usage, err := client.fetch(ctx, true)
|
||||
if err != nil {
|
||||
status, message, ok := commandCodeFetchFailure(err)
|
||||
if !ok {
|
||||
if isApplicationError(err) {
|
||||
return nil, err
|
||||
}
|
||||
return nil, infraerrors.Newf(http.StatusBadGateway, "CN_QUOTA_REQUEST_FAILED", "%v", err)
|
||||
}
|
||||
result.StatusCode = status
|
||||
result.Error = message
|
||||
return result, nil
|
||||
}
|
||||
|
||||
tiers := usage.tiers()
|
||||
balance := usage.balance(now)
|
||||
result.StatusCode = http.StatusOK
|
||||
result.Tiers = tiers
|
||||
result.PlanLevel = usage.planName()
|
||||
result.Balance = balance
|
||||
result.Success = true
|
||||
result.CredentialValid = true
|
||||
|
||||
updates := commandCodeUsageExtraUpdates(usage, tiers, balance, now, true)
|
||||
if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil {
|
||||
slog.Warn("cn_quota_persist_failed", "account_id", account.ID, "provider", PlatformCommandCode, "error", err)
|
||||
} else {
|
||||
result.Persisted = true
|
||||
balance.Persisted = true
|
||||
}
|
||||
clearCommandCodeWindowCooldown(ctx, s.accountRepo, account, usage, now)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// queryCommandCodeBalance 只查积分余额(whoami + credits),落余额快照。
|
||||
func (s *CNProviderBalanceService) queryCommandCodeBalance(ctx context.Context, account *Account) (*CNProviderBalanceResult, error) {
|
||||
client := newCommandCodeUsageClient(s.cfg, s.httpUpstream, account, s.resolveProxyURL(ctx, account))
|
||||
if client.apiKey == "" {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_BALANCE_NO_APIKEY", "account api_key is empty")
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
usage, err := client.fetch(ctx, false)
|
||||
if err != nil {
|
||||
status, message, ok := commandCodeFetchFailure(err)
|
||||
if !ok {
|
||||
if isApplicationError(err) {
|
||||
return nil, err
|
||||
}
|
||||
return nil, infraerrors.Newf(http.StatusBadGateway, "CN_BALANCE_REQUEST_FAILED", "%v", err)
|
||||
}
|
||||
return &CNProviderBalanceResult{
|
||||
Provider: PlatformCommandCode,
|
||||
StatusCode: status,
|
||||
FetchedAt: now.Unix(),
|
||||
Available: true,
|
||||
Error: message,
|
||||
}, nil
|
||||
}
|
||||
result := usage.balance(now)
|
||||
if err := s.accountRepo.UpdateExtra(ctx, account.ID, commandCodeUsageExtraUpdates(usage, nil, result, now, false)); err != nil {
|
||||
slog.Warn("cn_balance_persist_failed", "account_id", account.ID, "provider", PlatformCommandCode, "error", err)
|
||||
} else {
|
||||
result.Persisted = true
|
||||
}
|
||||
clearCommandCodeWindowCooldown(ctx, s.accountRepo, account, usage, now)
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,694 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// commandCodeAlphaUpstream 按路径返回 /alpha 接口响应,并记录请求。
|
||||
type commandCodeAlphaUpstream struct {
|
||||
mu sync.Mutex
|
||||
responses map[string]commandCodeAlphaResponse
|
||||
requests []*http.Request
|
||||
}
|
||||
|
||||
type commandCodeAlphaResponse struct {
|
||||
status int
|
||||
body string
|
||||
}
|
||||
|
||||
func (u *commandCodeAlphaUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
u.mu.Lock()
|
||||
u.requests = append(u.requests, req)
|
||||
u.mu.Unlock()
|
||||
resp, ok := u.responses[req.URL.Path]
|
||||
if !ok {
|
||||
resp = commandCodeAlphaResponse{status: http.StatusNotFound, body: `{"error":"not found"}`}
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: resp.status,
|
||||
Body: io.NopCloser(strings.NewReader(resp.body)),
|
||||
Header: make(http.Header),
|
||||
Request: req,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (u *commandCodeAlphaUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, concurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
|
||||
return u.Do(req, proxyURL, accountID, concurrency)
|
||||
}
|
||||
|
||||
func (u *commandCodeAlphaUpstream) request(path string) *http.Request {
|
||||
for _, req := range u.requests {
|
||||
if req.URL.Path == path {
|
||||
return req
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
const (
|
||||
commandCodeTestCredits = `{
|
||||
"credits": {"belowThreshold": false, "monthlyCredits": "34.9", "purchasedCredits": 0, "freeCredits": 0.1},
|
||||
"windowLimits": {
|
||||
"limited": true,
|
||||
"fiveHour": {"used": 1.5, "cap": 14, "resetAt": 1790000000000},
|
||||
"weekly": {"used": 35, "cap": 35, "resetAt": 1790400000000}
|
||||
}
|
||||
}`
|
||||
commandCodeTestSubscription = `{"data": {"planId": "individual-goat", "status": "active",
|
||||
"currentPeriodStart": "2026-09-01T00:00:00Z", "currentPeriodEnd": "2026-10-01T00:00:00Z"}}`
|
||||
commandCodeTestSummary = `{"totalCost": 40, "totalMonthlyCredits": 35.1}`
|
||||
)
|
||||
|
||||
func commandCodeUsageAccount() *Account {
|
||||
account := commandCodeTestAccount(77)
|
||||
account.Concurrency = 1
|
||||
return account
|
||||
}
|
||||
|
||||
func newCommandCodeAlphaUpstream(whoami string) *commandCodeAlphaUpstream {
|
||||
return &commandCodeAlphaUpstream{responses: map[string]commandCodeAlphaResponse{
|
||||
"/alpha/whoami": {status: http.StatusOK, body: whoami},
|
||||
"/alpha/billing/credits": {status: http.StatusOK, body: commandCodeTestCredits},
|
||||
"/alpha/billing/subscriptions": {status: http.StatusOK, body: commandCodeTestSubscription},
|
||||
"/alpha/usage/summary": {status: http.StatusOK, body: commandCodeTestSummary},
|
||||
}}
|
||||
}
|
||||
|
||||
func TestCommandCodeQuotaProbeParsesWindowsAndBalance(t *testing.T) {
|
||||
account := commandCodeUsageAccount()
|
||||
upstream := newCommandCodeAlphaUpstream(`{"user": {"userName": "alice"}, "org": null}`)
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
svc := NewCNProviderQuotaService(repo, nil, upstream, &config.Config{})
|
||||
|
||||
result, err := svc.QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.Success)
|
||||
require.True(t, result.Persisted)
|
||||
require.Equal(t, PlatformCommandCode, result.Provider)
|
||||
require.Equal(t, "GOAT", result.PlanLevel)
|
||||
require.Equal(t, []CNQuotaTier{
|
||||
{Window: "5h", UsedPercent: 10.71, ResetAt: time.UnixMilli(1790000000000).UTC().Format(time.RFC3339)},
|
||||
{Window: "weekly", UsedPercent: 100, ResetAt: time.UnixMilli(1790400000000).UTC().Format(time.RFC3339)},
|
||||
{Window: "monthly", UsedPercent: 50.14, ResetAt: "2026-10-01T00:00:00Z"},
|
||||
}, result.Tiers)
|
||||
require.NotNil(t, result.Balance)
|
||||
require.InDelta(t, 35.0, result.Balance.Balance, 1e-9)
|
||||
require.Equal(t, "USD", result.Balance.Currency)
|
||||
|
||||
// 个人账户不带 orgId;本周期用量从订阅周期起点统计。
|
||||
for _, req := range upstream.requests {
|
||||
require.Equal(t, "Bearer user_test_key", req.Header.Get("Authorization"))
|
||||
require.Equal(t, "api.commandcode.ai", req.URL.Host)
|
||||
require.False(t, req.URL.Query().Has("orgId"), req.URL.String())
|
||||
}
|
||||
require.Equal(t, "2026-09-01T00:00:00Z", upstream.request("/alpha/usage/summary").URL.Query().Get("since"))
|
||||
|
||||
require.Len(t, repo.extraWrites, 1)
|
||||
extra := repo.extraWrites[0]
|
||||
require.Equal(t, 10.71, extra["command_code_5h_used_percent"])
|
||||
require.Equal(t, 100.0, extra["command_code_weekly_used_percent"])
|
||||
require.Equal(t, 50.14, extra["command_code_monthly_used_percent"])
|
||||
require.Equal(t, "2026-10-01T00:00:00Z", extra["command_code_monthly_reset_at"])
|
||||
require.InDelta(t, 35.0, extra["command_code_balance"], 1e-9)
|
||||
require.Equal(t, "USD", extra["command_code_balance_currency"])
|
||||
require.Equal(t, 0.0, extra["command_code_purchased_credits"])
|
||||
require.Equal(t, false, extra["command_code_balance_low"])
|
||||
}
|
||||
|
||||
func TestCommandCodeQuotaProbePassesOrgID(t *testing.T) {
|
||||
account := commandCodeUsageAccount()
|
||||
upstream := newCommandCodeAlphaUpstream(`{"user": {"userName": "alice"}, "org": {"id": "9a0f5d0e-6d2c-4f36-9c55-0b8b0c0d1e2f", "login": "team"}}`)
|
||||
svc := NewCNProviderQuotaService(&cnBalanceProbeRepo{account: account}, nil, upstream, &config.Config{})
|
||||
|
||||
_, err := svc.QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
for _, path := range []string{"/alpha/billing/credits", "/alpha/billing/subscriptions", "/alpha/usage/summary"} {
|
||||
require.Equal(t, "9a0f5d0e-6d2c-4f36-9c55-0b8b0c0d1e2f", upstream.request(path).URL.Query().Get("orgId"), path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCommandCodeQuotaProbeUnlimitedPlanClearsRollingWindows(t *testing.T) {
|
||||
account := commandCodeUsageAccount()
|
||||
upstream := newCommandCodeAlphaUpstream(`{"org": null}`)
|
||||
upstream.responses["/alpha/billing/credits"] = commandCodeAlphaResponse{status: http.StatusOK, body: `{
|
||||
"credits": {"monthlyCredits": 10, "purchasedCredits": 25, "freeCredits": 0},
|
||||
"windowLimits": {"limited": false}
|
||||
}`}
|
||||
upstream.responses["/alpha/usage/summary"] = commandCodeAlphaResponse{status: http.StatusOK, body: `{"totalCost": 5}`}
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
svc := NewCNProviderQuotaService(repo, nil, upstream, &config.Config{})
|
||||
|
||||
result, err := svc.QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []CNQuotaTier{{Window: "monthly", UsedPercent: 33.33, ResetAt: "2026-10-01T00:00:00Z"}}, result.Tiers)
|
||||
extra := repo.extraWrites[0]
|
||||
for _, key := range []string{"command_code_5h_used_percent", "command_code_5h_reset_at", "command_code_weekly_used_percent", "command_code_weekly_reset_at"} {
|
||||
value, ok := extra[key]
|
||||
require.True(t, ok, key)
|
||||
require.Nil(t, value, key)
|
||||
}
|
||||
require.Equal(t, 25.0, extra["command_code_purchased_credits"])
|
||||
require.InDelta(t, 35.0, extra["command_code_balance"], 1e-9)
|
||||
}
|
||||
|
||||
// 窗口尚未开始(resetAt 为空)时清掉上一个窗口的重置时间,不与新的用量配对。
|
||||
func TestCommandCodeQuotaProbeClearsResetOfUnstartedWindow(t *testing.T) {
|
||||
account := commandCodeUsageAccount()
|
||||
upstream := newCommandCodeAlphaUpstream(`{"org": null}`)
|
||||
upstream.responses["/alpha/billing/credits"] = commandCodeAlphaResponse{status: http.StatusOK, body: `{
|
||||
"credits": {"monthlyCredits": 34.9, "purchasedCredits": 0, "freeCredits": 0},
|
||||
"windowLimits": {
|
||||
"limited": true,
|
||||
"fiveHour": {"used": 0, "cap": 14, "resetAt": null},
|
||||
"weekly": {"used": 35, "cap": 35, "resetAt": 1790400000000}
|
||||
}
|
||||
}`}
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
|
||||
_, err := NewCNProviderQuotaService(repo, nil, upstream, &config.Config{}).QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
extra := repo.extraWrites[0]
|
||||
require.Equal(t, 0.0, extra["command_code_5h_used_percent"])
|
||||
value, written := extra["command_code_5h_reset_at"]
|
||||
require.True(t, written)
|
||||
require.Nil(t, value)
|
||||
require.Equal(t, time.UnixMilli(1790400000000).UTC().Format(time.RFC3339), extra["command_code_weekly_reset_at"])
|
||||
require.Equal(t, "2026-10-01T00:00:00Z", extra["command_code_monthly_reset_at"])
|
||||
}
|
||||
|
||||
func TestCommandCodeQuotaProbeAuthFailureDoesNotPersist(t *testing.T) {
|
||||
account := commandCodeUsageAccount()
|
||||
upstream := newCommandCodeAlphaUpstream(`{}`)
|
||||
upstream.responses["/alpha/whoami"] = commandCodeAlphaResponse{status: http.StatusUnauthorized, body: `{"error":"unauthorized"}`}
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
svc := NewCNProviderQuotaService(repo, nil, upstream, &config.Config{})
|
||||
|
||||
result, err := svc.QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.False(t, result.Success)
|
||||
require.Equal(t, http.StatusUnauthorized, result.StatusCode)
|
||||
require.Equal(t, "Authentication failed (HTTP 401)", result.Error)
|
||||
require.Empty(t, repo.extraWrites)
|
||||
require.Len(t, upstream.requests, 1)
|
||||
}
|
||||
|
||||
// 订阅周期或本周期用量取数失败:滚动窗口照常更新,月度快照保留上一次的值(无从判断
|
||||
// 月度窗口是否还在);订阅确实没有周期时才清掉月度快照。
|
||||
// whoami 失败(非鉴权错误)时无从判断是否组织账户:探测整体失败,不按个人账户去查
|
||||
// credits,也不写快照。
|
||||
func TestCommandCodeProbeFailsWhenWhoamiFails(t *testing.T) {
|
||||
for _, whoami := range []commandCodeAlphaResponse{
|
||||
{status: http.StatusBadGateway, body: `bad gateway`},
|
||||
{status: http.StatusOK, body: `not json`},
|
||||
} {
|
||||
account := commandCodeUsageAccount()
|
||||
upstream := newCommandCodeAlphaUpstream(`{}`)
|
||||
upstream.responses["/alpha/whoami"] = whoami
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
|
||||
result, err := NewCNProviderQuotaService(repo, nil, upstream, &config.Config{}).QueryUsage(context.Background(), account.ID)
|
||||
require.NoError(t, err, whoami.body)
|
||||
require.False(t, result.Success, whoami.body)
|
||||
require.NotEmpty(t, result.Error, whoami.body)
|
||||
|
||||
balance, err := NewCNProviderBalanceService(repo, nil, upstream, &config.Config{}).QueryBalance(context.Background(), account.ID)
|
||||
require.NoError(t, err, whoami.body)
|
||||
require.False(t, balance.Success, whoami.body)
|
||||
|
||||
require.Nil(t, upstream.request("/alpha/billing/credits"), whoami.body)
|
||||
require.Empty(t, repo.extraWrites, whoami.body)
|
||||
}
|
||||
|
||||
// 请求本身失败(网络错误)同样整体失败。
|
||||
account := commandCodeUsageAccount()
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
_, err := NewCNProviderQuotaService(repo, nil, failingCommandCodeUpstream{}, &config.Config{}).QueryUsage(context.Background(), account.ID)
|
||||
require.Error(t, err)
|
||||
require.Empty(t, repo.extraWrites)
|
||||
}
|
||||
|
||||
// failingCommandCodeUpstream 让每个请求都在网络层失败。
|
||||
type failingCommandCodeUpstream struct{}
|
||||
|
||||
func (failingCommandCodeUpstream) Do(*http.Request, string, int64, int) (*http.Response, error) {
|
||||
return nil, errors.New("dial tcp: i/o timeout")
|
||||
}
|
||||
|
||||
func (u failingCommandCodeUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, concurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
|
||||
return u.Do(req, proxyURL, accountID, concurrency)
|
||||
}
|
||||
|
||||
func TestCommandCodeQuotaProbeKeepsWindowsWhenPeriodEndpointsFail(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
path string
|
||||
response commandCodeAlphaResponse
|
||||
planLevel string
|
||||
keepMonthly bool
|
||||
}{
|
||||
{"subscriptions failed", "/alpha/billing/subscriptions", commandCodeAlphaResponse{status: http.StatusInternalServerError, body: `oops`}, "", true},
|
||||
{"summary failed", "/alpha/usage/summary", commandCodeAlphaResponse{status: http.StatusBadGateway, body: `oops`}, "GOAT", true},
|
||||
{"no subscription period", "/alpha/billing/subscriptions", commandCodeAlphaResponse{status: http.StatusOK, body: `{"data":{}}`}, "", false},
|
||||
} {
|
||||
account := commandCodeUsageAccount()
|
||||
upstream := newCommandCodeAlphaUpstream(`{"org": null}`)
|
||||
upstream.responses[tc.path] = tc.response
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
|
||||
result, err := NewCNProviderQuotaService(repo, nil, upstream, &config.Config{}).QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err, tc.name)
|
||||
require.True(t, result.Success, tc.name)
|
||||
require.Len(t, result.Tiers, 2, tc.name)
|
||||
require.Equal(t, tc.planLevel, result.PlanLevel, tc.name)
|
||||
if tc.path == "/alpha/billing/subscriptions" {
|
||||
require.Nil(t, upstream.request("/alpha/usage/summary"), tc.name)
|
||||
}
|
||||
extra := repo.extraWrites[0]
|
||||
require.Equal(t, 10.71, extra["command_code_5h_used_percent"], tc.name)
|
||||
for _, key := range []string{"command_code_monthly_used_percent", "command_code_monthly_reset_at"} {
|
||||
value, written := extra[key]
|
||||
require.Equal(t, !tc.keepMonthly, written, tc.name+": "+key)
|
||||
require.Nil(t, value, tc.name+": "+key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCommandCodeUsageRequiresOfficialHost(t *testing.T) {
|
||||
account := commandCodeUsageAccount()
|
||||
require.Equal(t, PlatformCommandCode, account.GetCodingPlanProvider())
|
||||
require.NoError(t, validateCodingPlanAccount(account))
|
||||
require.NoError(t, validatePayGAccount(account))
|
||||
|
||||
account.Credentials["base_url"] = "https://relay.example.com/v1"
|
||||
account.Credentials["api_protocol"] = APIProtocolChatCompletions
|
||||
require.Empty(t, account.GetCodingPlanProvider())
|
||||
require.Error(t, validateCodingPlanAccount(account))
|
||||
require.Error(t, validatePayGAccount(account))
|
||||
}
|
||||
|
||||
func TestCommandCodeBalanceProbeFetchesCreditsOnly(t *testing.T) {
|
||||
account := commandCodeUsageAccount()
|
||||
upstream := newCommandCodeAlphaUpstream(`{"org": null}`)
|
||||
repo := &cnBalanceProbeRepo{account: account}
|
||||
svc := NewCNProviderBalanceService(repo, nil, upstream, &config.Config{})
|
||||
|
||||
result, err := svc.QueryBalance(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.Success)
|
||||
require.True(t, result.Persisted)
|
||||
require.InDelta(t, 35.0, result.Balance, 1e-9)
|
||||
require.Len(t, upstream.requests, 2)
|
||||
require.Len(t, repo.extraWrites, 1)
|
||||
_, hasWindow := repo.extraWrites[0]["command_code_5h_used_percent"]
|
||||
require.False(t, hasWindow)
|
||||
require.InDelta(t, 35.0, repo.extraWrites[0]["command_code_balance"], 1e-9)
|
||||
}
|
||||
|
||||
func TestCommandCodeWindowsIgnoredWithPurchasedCredits(t *testing.T) {
|
||||
future := time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
|
||||
account := commandCodeUsageAccount()
|
||||
account.Extra = map[string]any{
|
||||
"command_code_5h_used_percent": 100.0,
|
||||
"command_code_5h_reset_at": future,
|
||||
"command_code_purchased_credits": 0.0,
|
||||
}
|
||||
require.NotEmpty(t, cnProviderThresholdCandidates(account, PlatformCommandCode))
|
||||
require.NotNil(t, cnProviderQuotaSnapshotReset(account, time.Now()))
|
||||
|
||||
account.Extra["command_code_purchased_credits"] = 12.5
|
||||
require.Empty(t, cnProviderThresholdCandidates(account, PlatformCommandCode))
|
||||
require.Nil(t, cnProviderQuotaSnapshotReset(account, time.Now()))
|
||||
}
|
||||
|
||||
// commandCodeCheckRepo 按平台列出账号,并记录周期检测对账号的停调。
|
||||
type commandCodeCheckRepo struct {
|
||||
AccountRepository
|
||||
accounts []Account
|
||||
mu sync.Mutex
|
||||
paused []int64
|
||||
}
|
||||
|
||||
func (r *commandCodeCheckRepo) ListByPlatform(_ context.Context, platform string) ([]Account, error) {
|
||||
var accounts []Account
|
||||
for _, account := range r.accounts {
|
||||
if account.Platform == platform {
|
||||
accounts = append(accounts, account)
|
||||
}
|
||||
}
|
||||
return accounts, nil
|
||||
}
|
||||
|
||||
func (r *commandCodeCheckRepo) SetTempUnschedulable(_ context.Context, id int64, _ time.Time, _ string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.paused = append(r.paused, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
type commandCodeBalanceProber struct {
|
||||
mu sync.Mutex
|
||||
probed []int64
|
||||
balance float64
|
||||
}
|
||||
|
||||
func (p *commandCodeBalanceProber) QueryUsage(_ context.Context, accountID int64) (*CNProviderQuotaProbeResult, error) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
p.probed = append(p.probed, accountID)
|
||||
return &CNProviderQuotaProbeResult{
|
||||
Success: true,
|
||||
Balance: &CNProviderBalanceResult{Success: true, Balance: p.balance, Currency: "USD", Available: true},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func TestCNProviderBalanceCheckRunOnceProbesCommandCode(t *testing.T) {
|
||||
official := *commandCodeUsageAccount()
|
||||
official.ID, official.Schedulable = 1, true
|
||||
manualOff := *commandCodeUsageAccount()
|
||||
manualOff.ID, manualOff.Schedulable = 2, false
|
||||
relay := *commandCodeTestAccount(3)
|
||||
relay.Schedulable = true
|
||||
relay.Credentials = map[string]any{"api_key": "k", "api_protocol": APIProtocolChatCompletions, "base_url": "https://relay.example.com/v1"}
|
||||
inactive := *commandCodeUsageAccount()
|
||||
inactive.ID, inactive.Status = 4, StatusDisabled
|
||||
|
||||
repo := &commandCodeCheckRepo{accounts: []Account{official, manualOff, relay, inactive}}
|
||||
prober := &commandCodeBalanceProber{balance: 0.2}
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.CNProviders.BalanceThreshold = 0.5
|
||||
svc := &CNProviderBalanceCheckService{accountRepo: repo, quotaService: prober, cfg: cfg}
|
||||
|
||||
svc.runOnce()
|
||||
|
||||
require.ElementsMatch(t, []int64{1, 2}, prober.probed)
|
||||
require.Equal(t, []int64{1}, repo.paused)
|
||||
}
|
||||
|
||||
// 积分不足(402)是可恢复状态:临时停调并标记 balance_low,不能永久置 error。
|
||||
func TestHandleUpstreamError_CommandCode402PausesInsteadOfError(t *testing.T) {
|
||||
repo := &rateLimitAccountRepoStub{}
|
||||
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
account := commandCodeUsageAccount()
|
||||
|
||||
shouldDisable := service.HandleUpstreamError(context.Background(), account, http.StatusPaymentRequired,
|
||||
http.Header{}, []byte(`{"error":{"message":"Insufficient credits"}}`))
|
||||
|
||||
require.True(t, shouldDisable)
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.Equal(t, 1, repo.tempCalls)
|
||||
require.Contains(t, repo.lastTempReason, cnBalanceLowReasonPrefix)
|
||||
require.Equal(t, true, repo.lastExtraUpdates["command_code_balance_low"])
|
||||
}
|
||||
|
||||
func TestHandleUpstreamError_CommandCode429InsufficientBalancePauses(t *testing.T) {
|
||||
repo := &rateLimitAccountRepoStub{}
|
||||
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
account := commandCodeUsageAccount()
|
||||
|
||||
service.HandleUpstreamError(context.Background(), account, http.StatusTooManyRequests,
|
||||
http.Header{}, []byte(`{"error":{"type":"insufficient_credits","message":"no credits left"}}`))
|
||||
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.Equal(t, 1, repo.tempCalls)
|
||||
require.Contains(t, repo.lastTempReason, cnBalanceLowReasonPrefix)
|
||||
require.Equal(t, 0, repo.rateLimitedCalls)
|
||||
}
|
||||
|
||||
// commandCodeRateLimitRepo 记录 Command Code 响应式处理写入的临时停调与模型级限流。
|
||||
type commandCodeRateLimitRepo struct {
|
||||
rateLimitAccountRepoStub
|
||||
tempUntil time.Time
|
||||
modelLimits map[string]time.Time
|
||||
reasons map[string]string
|
||||
}
|
||||
|
||||
func (r *commandCodeRateLimitRepo) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
|
||||
r.tempUntil = until
|
||||
return r.rateLimitAccountRepoStub.SetTempUnschedulable(ctx, id, until, reason)
|
||||
}
|
||||
|
||||
func (r *commandCodeRateLimitRepo) SetModelRateLimit(_ context.Context, _ int64, scope string, resetAt time.Time, reason ...string) error {
|
||||
if r.modelLimits == nil {
|
||||
r.modelLimits = map[string]time.Time{}
|
||||
r.reasons = map[string]string{}
|
||||
}
|
||||
r.modelLimits[scope] = resetAt
|
||||
if len(reason) > 0 {
|
||||
r.reasons[scope] = reason[0]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 窗口用满:按文案中的重置时间写带原因的临时停调(充值后可据原因解除),不依赖额度快照。
|
||||
func TestHandleUpstreamError_CommandCodeWindowLimitCoolsDownToMessageReset(t *testing.T) {
|
||||
for _, status := range []int{http.StatusTooManyRequests, http.StatusForbidden, http.StatusPaymentRequired} {
|
||||
repo := &commandCodeRateLimitRepo{}
|
||||
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
account := commandCodeUsageAccount()
|
||||
|
||||
shouldDisable := service.HandleUpstreamError(context.Background(), account, status, http.Header{},
|
||||
[]byte(`{"error":{"type":"usage_exceeded","message":"You've reached your 5-hour usage limit. Resets in 2h 41m (3:00 PM)."}}`))
|
||||
|
||||
require.Equal(t, status != http.StatusTooManyRequests, shouldDisable, status)
|
||||
require.Equal(t, 0, repo.setErrorCalls, status)
|
||||
require.Equal(t, 0, repo.rateLimitedCalls, status)
|
||||
require.Equal(t, 1, repo.tempCalls, status)
|
||||
require.True(t, strings.HasPrefix(repo.lastTempReason, commandCodeUsageLimitReason+": "), repo.lastTempReason)
|
||||
expected := time.Now().Add(2*time.Hour + 41*time.Minute)
|
||||
require.WithinDuration(t, expected, repo.tempUntil, time.Minute, status)
|
||||
}
|
||||
}
|
||||
|
||||
// 文案没有重置时间时回落额度快照中最早的窗口重置点。
|
||||
func TestHandleUpstreamError_CommandCodeWindowLimitFallsBackToSnapshot(t *testing.T) {
|
||||
repo := &commandCodeRateLimitRepo{}
|
||||
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
account := commandCodeUsageAccount()
|
||||
account.Extra = map[string]any{
|
||||
"command_code_5h_used_percent": 100.0,
|
||||
"command_code_5h_reset_at": time.Now().Add(90 * time.Minute).UTC().Format(time.RFC3339),
|
||||
"command_code_purchased_credits": 0.0,
|
||||
}
|
||||
|
||||
service.HandleUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{},
|
||||
[]byte(`{"error":{"message":"You've reached your weekly usage limit."}}`))
|
||||
|
||||
require.Equal(t, 1, repo.tempCalls)
|
||||
require.WithinDuration(t, time.Now().Add(90*time.Minute), repo.tempUntil, time.Minute)
|
||||
}
|
||||
|
||||
// 普通频率限制很快恢复:即使快照里有窗口重置点,也不冷却到窗口重置。
|
||||
func TestHandleUpstreamError_CommandCodeRateLimitUsesDefaultCooldown(t *testing.T) {
|
||||
repo := &commandCodeRateLimitRepo{}
|
||||
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
account := commandCodeUsageAccount()
|
||||
account.Extra = map[string]any{
|
||||
"command_code_5h_used_percent": 97.0,
|
||||
"command_code_5h_reset_at": time.Now().Add(4 * time.Hour).UTC().Format(time.RFC3339),
|
||||
"command_code_purchased_credits": 0.0,
|
||||
}
|
||||
|
||||
service.HandleUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{},
|
||||
[]byte(`{"error":{"type":"rate_limit","message":"Rate limit exceeded. Please wait a moment and try again."}}`))
|
||||
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.Equal(t, 0, repo.tempCalls)
|
||||
if repo.rateLimitedCalls > 0 {
|
||||
require.True(t, repo.lastRateLimitedAt.Before(time.Now().Add(30*time.Minute)), "must not cool down to the 5h window reset")
|
||||
}
|
||||
}
|
||||
|
||||
const (
|
||||
commandCodeSpendLimitBody = `{"error":{"type":"usage_exceeded","message":"You've reached the $10.00/mo spending limit your organization set for GLM-5. It resets in 3 days."}}`
|
||||
commandCodeMemberSpendLimitBody = `{"error":{"type":"usage_exceeded","message":"You've reached the $50.00/mo spending limit your organization set. It resets in 12 days."}}`
|
||||
)
|
||||
|
||||
// 组织为某个模型设的花费上限只限当前模型:写模型级限流到文案中的重置时间,不停整个账号。
|
||||
func TestHandleUpstreamError_CommandCodeSpendLimitIsModelScoped(t *testing.T) {
|
||||
repo := &commandCodeRateLimitRepo{}
|
||||
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
account := commandCodeUsageAccount()
|
||||
ctx := withTempUnschedulableModel(context.Background(), []string{"glm-5"})
|
||||
|
||||
shouldDisable := service.HandleUpstreamError(ctx, account, http.StatusForbidden, http.Header{}, []byte(commandCodeSpendLimitBody))
|
||||
|
||||
require.False(t, shouldDisable)
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.Equal(t, 0, repo.tempCalls)
|
||||
require.Equal(t, 0, repo.rateLimitedCalls)
|
||||
require.Contains(t, repo.modelLimits, "glm-5")
|
||||
require.WithinDuration(t, time.Now().Add(72*time.Hour), repo.modelLimits["glm-5"], time.Minute)
|
||||
require.Equal(t, commandCodeSpendLimitReason, repo.reasons["glm-5"])
|
||||
}
|
||||
|
||||
// 成员总额度耗尽:即使请求模型已知也按账号冷却,但最长一小时(管理员可能调高上限)。
|
||||
func TestHandleUpstreamError_CommandCodeMemberSpendLimitIsAccountScoped(t *testing.T) {
|
||||
repo := &commandCodeRateLimitRepo{}
|
||||
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
account := commandCodeUsageAccount()
|
||||
ctx := withTempUnschedulableModel(context.Background(), []string{"glm-5"})
|
||||
|
||||
shouldDisable := service.HandleUpstreamError(ctx, account, http.StatusForbidden, http.Header{}, []byte(commandCodeMemberSpendLimitBody))
|
||||
|
||||
require.True(t, shouldDisable)
|
||||
require.Empty(t, repo.modelLimits)
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.Equal(t, 1, repo.tempCalls)
|
||||
require.Equal(t, commandCodeSpendLimitReason, repo.lastTempReason)
|
||||
require.WithinDuration(t, time.Now().Add(commandCodeAccountSpendLimitRecheck), repo.tempUntil, time.Minute)
|
||||
|
||||
// 文案中的重置时间早于一小时时以文案为准。
|
||||
soon := &commandCodeRateLimitRepo{}
|
||||
NewRateLimitService(soon, nil, &config.Config{}, nil, nil).HandleUpstreamError(ctx, account, http.StatusTooManyRequests, http.Header{},
|
||||
[]byte(`{"error":{"message":"You've reached the $50.00/mo spending limit your organization set. It resets in 20 minutes."}}`))
|
||||
require.WithinDuration(t, time.Now().Add(20*time.Minute), soon.tempUntil, time.Minute)
|
||||
}
|
||||
|
||||
// 不知道模型时按账号冷却(最长一小时),不停账号数天,也不进入 403 累计禁用。
|
||||
func TestHandleUpstreamError_CommandCodeSpendLimitWithoutModelShortCooldown(t *testing.T) {
|
||||
repo := &commandCodeRateLimitRepo{}
|
||||
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
account := commandCodeUsageAccount()
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
service.HandleUpstreamError(context.Background(), account, http.StatusForbidden, http.Header{}, []byte(commandCodeSpendLimitBody))
|
||||
}
|
||||
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.Empty(t, repo.modelLimits)
|
||||
require.Equal(t, 5, repo.tempCalls)
|
||||
require.False(t, repo.tempUntil.After(time.Now().Add(commandCodeAccountSpendLimitRecheck)))
|
||||
}
|
||||
|
||||
func TestCommandCodeSpendLimitScope(t *testing.T) {
|
||||
require.True(t, commandCodeSpendLimitIsModelScoped([]byte(commandCodeSpendLimitBody)))
|
||||
require.True(t, commandCodeSpendLimitIsModelScoped([]byte(`spending limit your organization set for gpt-5.5. It resets in 2 days.`)))
|
||||
require.False(t, commandCodeSpendLimitIsModelScoped([]byte(commandCodeMemberSpendLimitBody)))
|
||||
require.False(t, commandCodeSpendLimitIsModelScoped([]byte(`You've reached the $50.00/mo spending limit your organization set for you.`)))
|
||||
require.False(t, commandCodeSpendLimitIsModelScoped([]byte(`spending limit reached`)))
|
||||
}
|
||||
|
||||
// commandCodeStatefulRepo 把停调写回同一个账号对象,用于验证完整的状态转换。
|
||||
type commandCodeStatefulRepo struct {
|
||||
AccountRepository
|
||||
account *Account
|
||||
clearCalls int
|
||||
}
|
||||
|
||||
func (r *commandCodeStatefulRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
|
||||
clone := *r.account
|
||||
return &clone, nil
|
||||
}
|
||||
|
||||
func (r *commandCodeStatefulRepo) UpdateExtra(_ context.Context, _ int64, _ map[string]any) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *commandCodeStatefulRepo) SetTempUnschedulable(_ context.Context, _ int64, until time.Time, reason string) error {
|
||||
r.account.TempUnschedulableUntil = &until
|
||||
r.account.TempUnschedulableReason = reason
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *commandCodeStatefulRepo) ClearTempUnschedulable(_ context.Context, _ int64) error {
|
||||
r.clearCalls++
|
||||
r.account.TempUnschedulableUntil = nil
|
||||
r.account.TempUnschedulableReason = ""
|
||||
return nil
|
||||
}
|
||||
|
||||
func commandCodeCreditsBody(purchased float64) string {
|
||||
return fmt.Sprintf(`{"credits":{"monthlyCredits":0,"purchasedCredits":%g,"freeCredits":0},
|
||||
"windowLimits":{"limited":true,"fiveHour":{"used":14,"cap":14,"resetAt":1790000000000}}}`, purchased)
|
||||
}
|
||||
|
||||
// 窗口用满进入冷却 → 充值 → 下一次额度探测发现充值积分 → 恢复可调度。
|
||||
func TestCommandCodeWindowCooldownClearedAfterTopUp(t *testing.T) {
|
||||
account := commandCodeUsageAccount()
|
||||
account.Schedulable = true
|
||||
repo := &commandCodeStatefulRepo{account: account}
|
||||
rateLimit := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
|
||||
rateLimit.HandleUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{},
|
||||
[]byte(`{"error":{"message":"You've reached your 5-hour usage limit. Resets in 3h 10m."}}`))
|
||||
require.False(t, account.IsSchedulable())
|
||||
|
||||
// 未充值:探测不解除冷却。
|
||||
upstream := newCommandCodeAlphaUpstream(`{"org": null}`)
|
||||
upstream.responses["/alpha/billing/credits"] = commandCodeAlphaResponse{status: http.StatusOK, body: commandCodeCreditsBody(0)}
|
||||
quota := NewCNProviderQuotaService(repo, nil, upstream, &config.Config{})
|
||||
_, err := quota.QueryUsage(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, repo.clearCalls)
|
||||
require.False(t, account.IsSchedulable())
|
||||
|
||||
// 充值后:额度探测与余额探测都能解除窗口冷却。
|
||||
upstream.responses["/alpha/billing/credits"] = commandCodeAlphaResponse{status: http.StatusOK, body: commandCodeCreditsBody(20)}
|
||||
_, err = quota.QueryUsage(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, repo.clearCalls)
|
||||
require.True(t, account.IsSchedulable())
|
||||
|
||||
rateLimit.HandleUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{},
|
||||
[]byte(`{"error":{"message":"You've reached your weekly usage limit. Resets in 2 days."}}`))
|
||||
require.False(t, account.IsSchedulable())
|
||||
_, err = NewCNProviderBalanceService(repo, nil, upstream, &config.Config{}).QueryBalance(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, account.IsSchedulable())
|
||||
}
|
||||
|
||||
// 充值积分只解除窗口冷却:积分不足、花费上限等其他原因的停调保持不变。
|
||||
func TestCommandCodeTopUpKeepsOtherCooldowns(t *testing.T) {
|
||||
for _, reason := range []string{cnBalanceLowReason("no credits"), commandCodeSpendLimitReason, "manual rule"} {
|
||||
until := time.Now().Add(time.Hour)
|
||||
account := commandCodeUsageAccount()
|
||||
account.Schedulable = true
|
||||
account.TempUnschedulableUntil = &until
|
||||
account.TempUnschedulableReason = reason
|
||||
repo := &commandCodeStatefulRepo{account: account}
|
||||
upstream := newCommandCodeAlphaUpstream(`{"org": null}`)
|
||||
upstream.responses["/alpha/billing/credits"] = commandCodeAlphaResponse{status: http.StatusOK, body: commandCodeCreditsBody(20)}
|
||||
|
||||
_, err := NewCNProviderQuotaService(repo, nil, upstream, &config.Config{}).QueryUsage(context.Background(), account.ID)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, repo.clearCalls, reason)
|
||||
require.False(t, account.IsSchedulable(), reason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassifyCommandCodeUsageError(t *testing.T) {
|
||||
cases := map[string]commandCodeUsageErrorKind{
|
||||
`{"error":{"message":"You've reached your 5-hour usage limit. Resets in 2h 41m (3:00 PM)."}}`: commandCodeUsageWindowLimit,
|
||||
`{"error":{"message":"You've reached your weekly usage limit."}}`: commandCodeUsageWindowLimit,
|
||||
commandCodeSpendLimitBody: commandCodeUsageSpendLimit,
|
||||
`{"error":{"type":"usage_exceeded","message":"Monthly credits spent."}}`: commandCodeUsageCreditsExhausted,
|
||||
`{"error":{"message":"Insufficient credits"}}`: commandCodeUsageCreditsExhausted,
|
||||
`{"error":{"type":"rate_limit","message":"Rate limit exceeded. Please wait a moment."}}`: commandCodeUsageErrorNone,
|
||||
`{"error":{"message":"forbidden"}}`: commandCodeUsageErrorNone,
|
||||
}
|
||||
for body, want := range cases {
|
||||
require.Equal(t, want, classifyCommandCodeUsageError([]byte(body)), body)
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
)
|
||||
|
||||
@@ -204,11 +205,5 @@ func (s *GatewayService) resolveCompositeRouteDecision(ctx context.Context, grou
|
||||
}
|
||||
|
||||
func isConcreteRequestPlatform(platform string) bool {
|
||||
switch platform {
|
||||
case PlatformAnthropic, PlatformOpenAI, PlatformGemini, PlatformAntigravity, PlatformGrok,
|
||||
PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return domain.IsConcretePlatform(platform)
|
||||
}
|
||||
|
||||
@@ -218,13 +218,13 @@ func TestCompositeGroupSchedulerHasAllCanonicalPlatformBuckets(t *testing.T) {
|
||||
platforms = append(platforms, platform)
|
||||
}
|
||||
require.ElementsMatch(t,
|
||||
[]string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe},
|
||||
[]string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe, PlatformCommandCode, PlatformCline},
|
||||
platforms,
|
||||
)
|
||||
}
|
||||
|
||||
func TestCompositeConcretePlatformsIncludeCNProviders(t *testing.T) {
|
||||
for _, platform := range []string{PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe} {
|
||||
for _, platform := range []string{PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo, PlatformTypeSafe, PlatformCommandCode, PlatformCline} {
|
||||
require.True(t, isConcreteRequestPlatform(platform))
|
||||
require.True(t, canCopyAccountsFromGroupPlatform(PlatformComposite, platform))
|
||||
}
|
||||
|
||||
@@ -22,6 +22,52 @@ func (r *CompositeRouteResolver) SetModelOwnershipResolver(resolver CompositeMod
|
||||
}
|
||||
}
|
||||
|
||||
// ListExactPublicModels returns enabled, concrete route IDs suitable for a model catalog.
|
||||
func (r *CompositeRouteResolver) ListExactPublicModels(ctx context.Context, groupID int64, endpoint string, includeSystemOne bool) ([]string, error) {
|
||||
if r == nil || r.repo == nil || groupID <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
routes, err := r.repo.ListByGroup(ctx, groupID, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
models := make([]string, 0, len(routes))
|
||||
for _, route := range routes {
|
||||
if route.Enabled && route.MatchType == CompositeRouteMatchExact &&
|
||||
(endpoint == "" || normalizeCompositeRouteEndpoint(route.Endpoint) == CompositeRouteEndpointAny || normalizeCompositeRouteEndpoint(route.Endpoint) == endpoint) {
|
||||
if model := strings.TrimSpace(route.PublicModel); model != "" {
|
||||
models = append(models, model)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !includeSystemOne {
|
||||
models = filterSystemOneRouteModels(routes, models, endpoint)
|
||||
}
|
||||
return models, nil
|
||||
}
|
||||
|
||||
func (r *CompositeRouteResolver) FilterCodexModels(ctx context.Context, groupID int64, models []string) ([]string, error) {
|
||||
if r == nil || r.repo == nil || groupID <= 0 {
|
||||
return models, nil
|
||||
}
|
||||
routes, err := r.repo.ListByGroup(ctx, groupID, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return filterSystemOneRouteModels(routes, models, CompositeRouteEndpointResponses), nil
|
||||
}
|
||||
|
||||
func filterSystemOneRouteModels(routes []CompositeModelRoute, models []string, endpoint string) []string {
|
||||
filtered := make([]string, 0, len(models))
|
||||
for _, model := range models {
|
||||
route, matched := matchCompositeRoute(routes, model, normalizeCompositeRouteEndpoint(endpoint))
|
||||
if !matched || route.TargetPlatform != PlatformTypeSafe {
|
||||
filtered = append(filtered, model)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func (r *CompositeRouteResolver) Resolve(ctx context.Context, groupID int64, model, endpoint string) (CompositeRouteDecision, error) {
|
||||
model = strings.TrimSpace(model)
|
||||
endpoint = normalizeCompositeRouteEndpoint(endpoint)
|
||||
|
||||
@@ -12,6 +12,85 @@ type compositeRouteRepoStub struct {
|
||||
routes []CompositeModelRoute
|
||||
}
|
||||
|
||||
func TestCompositeRouteResolverCodexFilteringMatchesProductionPriority(t *testing.T) {
|
||||
for _, scenario := range []struct {
|
||||
name string
|
||||
routes []CompositeModelRoute
|
||||
want []string
|
||||
}{
|
||||
{
|
||||
name: "endpoint beats priority",
|
||||
routes: []CompositeModelRoute{
|
||||
{ID: 1, PublicModel: "alias", MatchType: CompositeRouteMatchExact, Endpoint: CompositeRouteEndpointAny, TargetPlatform: PlatformOpenAI, Priority: 1},
|
||||
{ID: 2, PublicModel: "alias", MatchType: CompositeRouteMatchExact, Endpoint: CompositeRouteEndpointResponses, TargetPlatform: PlatformTypeSafe, Priority: 100},
|
||||
},
|
||||
want: []string{"unrouted"},
|
||||
},
|
||||
{
|
||||
name: "Responses OpenAI overrides any TypeSafe",
|
||||
routes: []CompositeModelRoute{
|
||||
{ID: 1, PublicModel: "alias", MatchType: CompositeRouteMatchExact, Endpoint: CompositeRouteEndpointAny, TargetPlatform: PlatformTypeSafe, Priority: 1},
|
||||
{ID: 2, PublicModel: "alias", MatchType: CompositeRouteMatchExact, Endpoint: CompositeRouteEndpointResponses, TargetPlatform: PlatformOpenAI, Priority: 100},
|
||||
},
|
||||
want: []string{"alias", "unrouted"},
|
||||
},
|
||||
{
|
||||
name: "longer prefix beats priority",
|
||||
routes: []CompositeModelRoute{
|
||||
{ID: 1, PublicModel: "a", MatchType: CompositeRouteMatchPrefix, TargetPlatform: PlatformOpenAI, Priority: 1},
|
||||
{ID: 2, PublicModel: "ali", MatchType: CompositeRouteMatchPrefix, TargetPlatform: PlatformTypeSafe, Priority: 100},
|
||||
},
|
||||
want: []string{"unrouted"},
|
||||
},
|
||||
{
|
||||
name: "exact beats endpoint prefix",
|
||||
routes: []CompositeModelRoute{
|
||||
{ID: 1, PublicModel: "alias", MatchType: CompositeRouteMatchPrefix, Endpoint: CompositeRouteEndpointResponses, TargetPlatform: PlatformTypeSafe},
|
||||
{ID: 2, PublicModel: "alias", MatchType: CompositeRouteMatchExact, Endpoint: CompositeRouteEndpointAny, TargetPlatform: PlatformOpenAI},
|
||||
},
|
||||
want: []string{"alias", "unrouted"},
|
||||
},
|
||||
{
|
||||
name: "lower priority wins",
|
||||
routes: []CompositeModelRoute{
|
||||
{ID: 1, PublicModel: "alias", MatchType: CompositeRouteMatchExact, TargetPlatform: PlatformTypeSafe, Priority: 100},
|
||||
{ID: 2, PublicModel: "alias", MatchType: CompositeRouteMatchExact, TargetPlatform: PlatformOpenAI, Priority: 1},
|
||||
},
|
||||
want: []string{"alias", "unrouted"},
|
||||
},
|
||||
{
|
||||
name: "lower ID breaks tie",
|
||||
routes: []CompositeModelRoute{
|
||||
{ID: 2, PublicModel: "alias", MatchType: CompositeRouteMatchExact, TargetPlatform: PlatformOpenAI},
|
||||
{ID: 1, PublicModel: "alias", MatchType: CompositeRouteMatchExact, TargetPlatform: PlatformTypeSafe},
|
||||
},
|
||||
want: []string{"unrouted"},
|
||||
},
|
||||
} {
|
||||
t.Run(scenario.name, func(t *testing.T) {
|
||||
for index := range scenario.routes {
|
||||
scenario.routes[index].Enabled = true
|
||||
scenario.routes[index].GroupID = 7
|
||||
}
|
||||
resolver := NewCompositeRouteResolver(compositeRouteRepoStub{routes: scenario.routes})
|
||||
models, err := resolver.FilterCodexModels(context.Background(), 7, []string{"alias", "unrouted"})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, scenario.want, models)
|
||||
decision, err := resolver.Resolve(context.Background(), 7, "alias", CompositeRouteEndpointResponses)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Matched)
|
||||
require.Equal(t, decision.TargetPlatform != PlatformTypeSafe, len(models) == 2)
|
||||
exactModels, err := resolver.ListExactPublicModels(context.Background(), 7, CompositeRouteEndpointResponses, false)
|
||||
require.NoError(t, err)
|
||||
if decision.TargetPlatform == PlatformTypeSafe {
|
||||
require.Empty(t, exactModels)
|
||||
} else {
|
||||
require.Contains(t, exactModels, "alias")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s compositeRouteRepoStub) ListByGroup(ctx context.Context, groupID int64, includeDisabled bool) ([]CompositeModelRoute, error) {
|
||||
routes := make([]CompositeModelRoute, 0, len(s.routes))
|
||||
for _, route := range s.routes {
|
||||
|
||||
@@ -44,13 +44,15 @@ const (
|
||||
PlatformAntigravity = domain.PlatformAntigravity
|
||||
PlatformGrok = domain.PlatformGrok
|
||||
// 国产 OpenAI 兼容供应商(与 grok 一样经 OpenAI 网关转发)。
|
||||
PlatformKimi = domain.PlatformKimi
|
||||
PlatformZhipu = domain.PlatformZhipu
|
||||
PlatformDeepseek = domain.PlatformDeepseek
|
||||
PlatformMiniMax = domain.PlatformMiniMax
|
||||
PlatformTypeSafe = domain.PlatformTypeSafe
|
||||
PlatformOpenCodeGo = domain.PlatformOpenCodeGo
|
||||
PlatformComposite = domain.PlatformComposite
|
||||
PlatformKimi = domain.PlatformKimi
|
||||
PlatformZhipu = domain.PlatformZhipu
|
||||
PlatformDeepseek = domain.PlatformDeepseek
|
||||
PlatformMiniMax = domain.PlatformMiniMax
|
||||
PlatformOpenCodeGo = domain.PlatformOpenCodeGo
|
||||
PlatformTypeSafe = domain.PlatformTypeSafe
|
||||
PlatformCommandCode = domain.PlatformCommandCode
|
||||
PlatformCline = domain.PlatformCline
|
||||
PlatformComposite = domain.PlatformComposite
|
||||
// PlatformKiro is retained for unsupported-platform threshold tests and legacy
|
||||
// account rows. Scheduling-threshold evaluation never pauses kiro accounts.
|
||||
PlatformKiro = "kiro"
|
||||
@@ -86,6 +88,10 @@ const (
|
||||
DefaultOpenCodeGoBaseURL = "https://opencode.ai/zen/go/v1"
|
||||
// OpenCode Zen:按量付费网关,模型列表为 /zen/v1/models。
|
||||
DefaultOpenCodeZenBaseURL = "https://opencode.ai/zen/v1"
|
||||
// Command Code Provider API:Chat Completions / Responses / models 共用 /provider/v1 基址。
|
||||
DefaultCommandCodeBaseURL = "https://api.commandcode.ai/provider/v1"
|
||||
// Cline API:只提供 Chat Completions({base}/chat/completions)与模型列表。
|
||||
DefaultClineBaseURL = "https://api.cline.bot/api/v1"
|
||||
)
|
||||
|
||||
// 国产供应商 Anthropic 协议端点的默认 base_url(上游路径为 {base}/v1/messages)。
|
||||
@@ -99,16 +105,14 @@ const (
|
||||
// OpenCode Go Anthropic 基址不含 /v1:nativeAnthropicTargetURL 会再拼 /v1/messages。
|
||||
DefaultOpenCodeGoAnthropicBaseURL = "https://opencode.ai/zen/go"
|
||||
DefaultOpenCodeZenAnthropicBaseURL = "https://opencode.ai/zen"
|
||||
// Command Code 的 Anthropic 端点为 /provider/v1/messages(Claude 系模型只在此端点提供)。
|
||||
DefaultCommandCodeAnthropicBaseURL = "https://api.commandcode.ai/provider"
|
||||
)
|
||||
|
||||
// IsCNProvider 报告 platform 是否为国产 OpenAI 兼容供应商(kimi/zhipu/deepseek/minimax)。
|
||||
// IsCNProvider 报告 platform 是否为国产 OpenAI 兼容供应商(kimi/zhipu/deepseek/minimax),
|
||||
// 以平台清单(domain/platforms.go)为准。
|
||||
func IsCNProvider(platform string) bool {
|
||||
switch platform {
|
||||
case PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return domain.IsCNProviderPlatform(platform)
|
||||
}
|
||||
|
||||
// IsOpenCodeGo 报告 platform 是否为 OpenCode Go 订阅网关。
|
||||
@@ -118,30 +122,19 @@ func IsOpenCodeGo(platform string) bool {
|
||||
|
||||
// IsMultiProtocolAPIKeyProvider 报告 platform 是否为多协议 API Key 网关
|
||||
// (国产供应商 + OpenCode):走 OpenAI 网关、支持 adaptive 协议分流。
|
||||
// 归属以 provider profile 登记为准(见 provider_profile.go)。
|
||||
func IsMultiProtocolAPIKeyProvider(platform string) bool {
|
||||
return IsCNProvider(platform) || platform == PlatformOpenCodeGo
|
||||
return LookupProviderProfile(platform) != nil
|
||||
}
|
||||
|
||||
// AllowedQuotaPlatforms 是允许设置 user × platform quota 的平台列表(单一权威来源)。
|
||||
// ent/schema/user_platform_quota.go 的 Validate 函数独立维护(构建期约束),
|
||||
// 若新增平台需同步修改该 schema。
|
||||
var AllowedQuotaPlatforms = []string{
|
||||
PlatformAnthropic,
|
||||
PlatformOpenAI,
|
||||
PlatformGemini,
|
||||
PlatformAntigravity,
|
||||
PlatformGrok,
|
||||
PlatformKimi,
|
||||
PlatformZhipu,
|
||||
PlatformDeepseek,
|
||||
PlatformMiniMax,
|
||||
PlatformOpenCodeGo,
|
||||
PlatformTypeSafe,
|
||||
}
|
||||
// AllowedQuotaPlatforms 是允许设置 user × platform quota 的平台列表:全部已登记的
|
||||
// 具体平台(domain/platforms.go),ent/schema/user_platform_quota.go 的校验同源。
|
||||
var AllowedQuotaPlatforms = domain.ConcretePlatformIDs()
|
||||
|
||||
// AllowedSchedulingThresholdPlatforms 是允许设置账号自动停调阈值的平台列表。
|
||||
// openai/anthropic/grok 有原生用量窗口;kimi/zhipu/minimax 的 Coding Plan 同样暴露
|
||||
// 5h/weekly 滚动窗口,纳入阈值评估。deepseek 为余额型,走余额检测而非阈值。
|
||||
// OpenCode Go 与 Command Code 的订阅套餐另有月度窗口。
|
||||
var AllowedSchedulingThresholdPlatforms = []string{
|
||||
PlatformOpenAI,
|
||||
PlatformAnthropic,
|
||||
@@ -150,6 +143,7 @@ var AllowedSchedulingThresholdPlatforms = []string{
|
||||
PlatformZhipu,
|
||||
PlatformMiniMax,
|
||||
PlatformOpenCodeGo,
|
||||
PlatformCommandCode,
|
||||
}
|
||||
|
||||
// IsAllowedQuotaPlatform 报告 s 是否为合法的 quota platform 标识。
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// requireCacheTTLOrderAccepted fails when a ttl=1h breakpoint follows a 5m one
|
||||
// in Anthropic's processing order (tools, system, messages). Anthropic rejects
|
||||
// such requests with HTTP 400.
|
||||
func requireCacheTTLOrderAccepted(t *testing.T, body []byte) {
|
||||
t.Helper()
|
||||
_, messagePaths, toolPaths, systemPaths := collectCacheControlPaths(body)
|
||||
var paths []string
|
||||
paths = append(paths, toolPaths...)
|
||||
paths = append(paths, systemPaths...)
|
||||
paths = append(paths, messagePaths...)
|
||||
shortTTLPath := ""
|
||||
for _, path := range paths {
|
||||
if gjson.GetBytes(body, path+".ttl").String() == cacheTTLTarget1h {
|
||||
require.Empty(t, shortTTLPath, "1h breakpoint %s follows 5m breakpoint %s in %s", path, shortTTLPath, body)
|
||||
continue
|
||||
}
|
||||
if shortTTLPath == "" {
|
||||
shortTTLPath = path
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeCacheControlTTLOrder_RaisesBreakpointsBeforeLast1h(t *testing.T) {
|
||||
body := []byte(`{"tools":[{"name":"t","input_schema":{"type":"object"},"cache_control":{"type":"ephemeral"}}],` +
|
||||
`"system":[{"type":"text","text":"s","cache_control":{"type":"ephemeral","ttl":"5m"}}],` +
|
||||
`"messages":[{"role":"user","content":[` +
|
||||
`{"type":"text","text":"a","cache_control":{"type":"ephemeral","ttl":"1h"}},` +
|
||||
`{"type":"text","text":"b","cache_control":{"type":"ephemeral","ttl":"5m"}}]}]}`)
|
||||
|
||||
out := normalizeCacheControlTTLOrder(body)
|
||||
|
||||
require.Equal(t, "1h", gjson.GetBytes(out, "tools.0.cache_control.ttl").String(), "omitted ttl counts as 5m")
|
||||
require.Equal(t, "1h", gjson.GetBytes(out, "system.0.cache_control.ttl").String())
|
||||
require.Equal(t, "1h", gjson.GetBytes(out, "messages.0.content.0.cache_control.ttl").String())
|
||||
require.Equal(t, "5m", gjson.GetBytes(out, "messages.0.content.1.cache_control.ttl").String(),
|
||||
"a 5m breakpoint after the last 1h one is valid and must keep its ttl")
|
||||
requireCacheTTLOrderAccepted(t, out)
|
||||
}
|
||||
|
||||
func TestNormalizeCacheControlTTLOrder_TopLevelCacheControlIsLast(t *testing.T) {
|
||||
body := []byte(`{"cache_control":{"type":"ephemeral","ttl":"1h"},` +
|
||||
`"system":[{"type":"text","text":"s","cache_control":{"type":"ephemeral","ttl":"5m"}}],` +
|
||||
`"messages":[{"role":"user","content":"hi"}]}`)
|
||||
|
||||
out := normalizeCacheControlTTLOrder(body)
|
||||
|
||||
require.Equal(t, "1h", gjson.GetBytes(out, "system.0.cache_control.ttl").String())
|
||||
}
|
||||
|
||||
func TestNormalizeCacheControlTTLOrder_KeepsValidRequestsUnchanged(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"no breakpoints": `{"system":"s","messages":[{"role":"user","content":"hi"}]}`,
|
||||
"all 5m": `{"tools":[{"name":"t","input_schema":{"type":"object"},"cache_control":{"type":"ephemeral"}}],` +
|
||||
`"system":[{"type":"text","text":"s","cache_control":{"type":"ephemeral","ttl":"5m"}}],` +
|
||||
`"messages":[{"role":"user","content":[{"type":"text","text":"a","cache_control":{"type":"ephemeral","ttl":"5m"}}]}]}`,
|
||||
"1h before 5m": `{"system":[{"type":"text","text":"s","cache_control":{"type":"ephemeral","ttl":"1h"}}],` +
|
||||
`"messages":[{"role":"user","content":[{"type":"text","text":"a","cache_control":{"type":"ephemeral"}}]}]}`,
|
||||
}
|
||||
for name, body := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
require.Equal(t, body, string(normalizeCacheControlTTLOrder([]byte(body))))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A non-Claude Code client sends one 1h breakpoint. The mimic path adds 5m
|
||||
// breakpoints to the injected system block and the last tool, which come first
|
||||
// in Anthropic's processing order. The outbound body must still be accepted.
|
||||
func TestEnforceCacheControlLimit_MimicPathKeepsClient1hBreakpointValid(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"message 1h": `{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":[` +
|
||||
`{"type":"text","text":"Reply OK.","cache_control":{"type":"ephemeral","ttl":"1h"}}]}]}`,
|
||||
"system 1h moved into messages": `{"model":"claude-sonnet-4-6",` +
|
||||
`"system":[{"type":"text","text":"Be terse.","cache_control":{"type":"ephemeral","ttl":"1h"}}],` +
|
||||
`"messages":[{"role":"user","content":"Reply OK."}]}`,
|
||||
"tools and message 1h": `{"model":"claude-sonnet-4-6","tools":[{"name":"probe","input_schema":{"type":"object"}}],` +
|
||||
`"messages":[{"role":"user","content":[` +
|
||||
`{"type":"text","text":"Reply OK.","cache_control":{"type":"ephemeral","ttl":"1h"}}]}]}`,
|
||||
}
|
||||
for name, raw := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
body := []byte(raw)
|
||||
var system any
|
||||
if sys := gjson.GetBytes(body, "system"); sys.Exists() {
|
||||
require.NoError(t, json.Unmarshal([]byte(sys.Raw), &system))
|
||||
}
|
||||
body = rewriteSystemForNonClaudeCode(body, system)
|
||||
body = applyToolsLastCacheBreakpoint(body)
|
||||
|
||||
out := enforceCacheControlLimit(body)
|
||||
|
||||
requireCacheTTLOrderAccepted(t, out)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardCountTokens_MimicPathKeepsClient1hBreakpointValid(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil)
|
||||
c.Request.Header.Set("User-Agent", "third-party-client/1.0")
|
||||
|
||||
// The mimic path adds a 5m breakpoint to tools[-1], ahead of the client's 1h one.
|
||||
body := []byte(`{"model":"claude-sonnet-4-6",` +
|
||||
`"tools":[{"name":"probe","description":"d","input_schema":{"type":"object"}}],` +
|
||||
`"messages":[{"role":"user","content":[{"type":"text","text":"Reply OK.","cache_control":{"type":"ephemeral","ttl":"1h"}}]}]}`)
|
||||
parsed := &ParsedRequest{Body: NewRequestBodyRef(body), Model: "claude-sonnet-4-6"}
|
||||
|
||||
upstream := &anthropicHTTPUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"input_tokens":42}`)),
|
||||
},
|
||||
}
|
||||
svc := &GatewayService{
|
||||
cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
|
||||
httpUpstream: upstream,
|
||||
rateLimitService: &RateLimitService{},
|
||||
}
|
||||
account := &Account{
|
||||
ID: 402,
|
||||
Name: "count-tokens-ttl-order-test",
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeSetupToken,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{"access_token": "test-oauth-token"},
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
}
|
||||
|
||||
require.NoError(t, svc.ForwardCountTokens(context.Background(), c, account, parsed))
|
||||
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, "tools.0.cache_control").Exists(),
|
||||
"the mimic tools breakpoint should still be present")
|
||||
requireCacheTTLOrderAccepted(t, upstream.lastBody)
|
||||
}
|
||||
@@ -1052,9 +1052,15 @@ func collectCacheControlPaths(body []byte) (invalidThinking []cacheControlPath,
|
||||
return invalidThinking, messagePaths, toolPaths, systemPaths
|
||||
}
|
||||
|
||||
// enforceCacheControlLimit 强制执行 cache_control 块数量限制(最多 4 个)
|
||||
// 超限时优先移除工具断点,再移除 messages 断点,最后才移除 system 断点。
|
||||
// enforceCacheControlLimit 让出站 cache_control 满足上游约束:
|
||||
// 先执行块数量限制,再修正 TTL 顺序。
|
||||
func enforceCacheControlLimit(body []byte) []byte {
|
||||
return normalizeCacheControlTTLOrder(enforceCacheControlBlockLimit(body))
|
||||
}
|
||||
|
||||
// enforceCacheControlBlockLimit 强制执行 cache_control 块数量限制(最多 4 个)
|
||||
// 超限时优先移除工具断点,再移除 messages 断点,最后才移除 system 断点。
|
||||
func enforceCacheControlBlockLimit(body []byte) []byte {
|
||||
if len(body) == 0 {
|
||||
return body
|
||||
}
|
||||
@@ -1137,6 +1143,46 @@ func enforceCacheControlLimit(body []byte) []byte {
|
||||
return body
|
||||
}
|
||||
|
||||
// normalizeCacheControlTTLOrder raises earlier ephemeral breakpoints to 1h when
|
||||
// a later breakpoint uses 1h. Anthropic processes tools, system, then messages,
|
||||
// and rejects a ttl=1h breakpoint that follows a 5m one; an omitted ttl means 5m.
|
||||
// The mimic path adds its own 5m breakpoints to tools and system, so a client
|
||||
// 1h breakpoint in messages would otherwise turn a valid request into a 400.
|
||||
//
|
||||
// Raising keeps the client's choice. Tokens before a 1h breakpoint are already
|
||||
// written at the 1h rate, so the earlier breakpoints add no cache write cost.
|
||||
// Requests with a valid order are returned unchanged.
|
||||
func normalizeCacheControlTTLOrder(body []byte) []byte {
|
||||
_, messagePaths, toolPaths, systemPaths := collectCacheControlPaths(body)
|
||||
paths := make([]string, 0, len(toolPaths)+len(systemPaths)+len(messagePaths)+1)
|
||||
paths = append(paths, toolPaths...)
|
||||
paths = append(paths, systemPaths...)
|
||||
paths = append(paths, messagePaths...)
|
||||
// A top-level cache_control applies to the last cacheable block.
|
||||
if gjson.GetBytes(body, "cache_control").Exists() {
|
||||
paths = append(paths, "cache_control")
|
||||
}
|
||||
|
||||
last1h := -1
|
||||
for i, path := range paths {
|
||||
if gjson.GetBytes(body, path+".ttl").String() == cacheTTLTarget1h {
|
||||
last1h = i
|
||||
}
|
||||
}
|
||||
|
||||
out := body
|
||||
for _, path := range paths[:last1h+1] {
|
||||
cc := gjson.GetBytes(out, path)
|
||||
if cc.Get("type").String() != "ephemeral" || cc.Get("ttl").String() == cacheTTLTarget1h {
|
||||
continue
|
||||
}
|
||||
if next, err := sjson.SetBytes(out, path+".ttl", cacheTTLTarget1h); err == nil {
|
||||
out = next
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// injectAnthropicCacheControlTTL1h 将已有 ephemeral cache_control 块的 ttl 强制写为 1h。
|
||||
// 仅修改已经存在的 cache_control,不新增缓存断点。
|
||||
func injectAnthropicCacheControlTTL1h(body []byte) []byte {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user