This commit is contained in:
yangjianbo
2026-02-09 20:26:03 +08:00
103 changed files with 8044 additions and 2470 deletions

323
DEV_GUIDE.md Normal file
View File

@@ -0,0 +1,323 @@
# sub2api 项目开发指南
> 本文档记录项目环境配置、常见坑点和注意事项,供 Claude Code 和团队成员参考。
## 一、项目基本信息
| 项目 | 说明 |
|------|------|
| **上游仓库** | Wei-Shaw/sub2api |
| **Fork 仓库** | bayma888/sub2api-bmai |
| **技术栈** | Go 后端 (Ent ORM + Gin) + Vue3 前端 (pnpm) |
| **数据库** | PostgreSQL 16 + Redis |
| **包管理** | 后端: go modules, 前端: **pnpm**(不是 npm |
## 二、本地环境配置
### PostgreSQL 16 (Windows 服务)
| 配置项 | 值 |
|--------|-----|
| 端口 | 5432 |
| psql 路径 | `C:\Program Files\PostgreSQL\16\bin\psql.exe` |
| pg_hba.conf | `C:\Program Files\PostgreSQL\16\data\pg_hba.conf` |
| 数据库凭据 | user=`sub2api`, password=`sub2api`, dbname=`sub2api` |
| 超级用户 | user=`postgres`, password=`postgres` |
### Redis
| 配置项 | 值 |
|--------|-----|
| 端口 | 6379 |
| 密码 | 无 |
### 开发工具
```bash
# golangci-lint v2.7
go install github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.7
# pnpm (前端包管理)
npm install -g pnpm
```
## 三、CI/CD 流水线
### GitHub Actions Workflows
| Workflow | 触发条件 | 检查内容 |
|----------|----------|----------|
| **backend-ci.yml** | push, pull_request | 单元测试 + 集成测试 + golangci-lint v2.7 |
| **security-scan.yml** | push, pull_request, 每周一 | govulncheck + gosec + pnpm audit |
| **release.yml** | tag `v*` | 构建发布PR 不触发) |
### CI 要求
- Go 版本必须是 **1.25.7**
- 前端使用 `pnpm install --frozen-lockfile`,必须提交 `pnpm-lock.yaml`
### 本地测试命令
```bash
# 后端单元测试
cd backend && go test -tags=unit ./...
# 后端集成测试
cd backend && go test -tags=integration ./...
# 代码质量检查
cd backend && golangci-lint run ./...
# 前端依赖安装(必须用 pnpm
cd frontend && pnpm install
```
## 四、常见坑点 & 解决方案
### 坑 1pnpm-lock.yaml 必须同步提交
**问题**`package.json` 新增依赖后CI 的 `pnpm install --frozen-lockfile` 失败。
**原因**:上游 CI 使用 pnpmlock 文件不同步会报错。
**解决**
```bash
cd frontend
pnpm install # 更新 pnpm-lock.yaml
git add pnpm-lock.yaml
git commit -m "chore: update pnpm-lock.yaml"
```
---
### 坑 2npm 和 pnpm 的 node_modules 冲突
**问题**:之前用 npm 装过 `node_modules`pnpm install 报 `EPERM` 错误。
**解决**
```bash
cd frontend
rm -rf node_modules # 或 PowerShell: Remove-Item -Recurse -Force node_modules
pnpm install
```
---
### 坑 3PowerShell 中 bcrypt hash 的 `$` 被转义
**问题**bcrypt hash 格式如 `$2a$10$xxx...`PowerShell 把 `$2a` 当变量解析,导致数据丢失。
**解决**:将 SQL 写入文件,用 `psql -f` 执行:
```bash
# 错误示范PowerShell 会吃掉 $
psql -c "INSERT INTO users ... VALUES ('$2a$10$...')"
# 正确做法
echo "INSERT INTO users ... VALUES ('\$2a\$10\$...')" > temp.sql
psql -U sub2api -h 127.0.0.1 -d sub2api -f temp.sql
```
---
### 坑 4psql 不支持中文路径
**问题**`psql -f "D:\中文路径\file.sql"` 报错找不到文件。
**解决**:复制到纯英文路径再执行:
```bash
cp "D:\中文路径\file.sql" "C:\temp.sql"
psql -f "C:\temp.sql"
```
---
### 坑 5PostgreSQL 密码重置流程
**场景**:忘记 PostgreSQL 密码。
**步骤**
1. 修改 `C:\Program Files\PostgreSQL\16\data\pg_hba.conf`
```
# 将 scram-sha-256 改为 trust
host all all 127.0.0.1/32 trust
```
2. 重启 PostgreSQL 服务
```powershell
Restart-Service postgresql-x64-16
```
3. 无密码登录并重置
```bash
psql -U postgres -h 127.0.0.1
ALTER USER sub2api WITH PASSWORD 'sub2api';
ALTER USER postgres WITH PASSWORD 'postgres';
```
4. 改回 `scram-sha-256` 并重启
---
### 坑 6Go interface 新增方法后 test stub 必须补全
**问题**:给 interface 新增方法后,编译报错 `does not implement interface (missing method XXX)`。
**原因**:所有测试文件中实现该 interface 的 stub/mock 都必须补上新方法。
**解决**
```bash
# 搜索所有实现该 interface 的 struct
cd backend
grep -r "type.*Stub.*struct" internal/
grep -r "type.*Mock.*struct" internal/
# 逐一补全新方法
```
---
### 坑 7Windows 上 psql 连 localhost 的 IPv6 问题
**问题**psql 连 `localhost` 先尝试 IPv6 (::1),可能报错后再回退 IPv4。
**建议**:直接用 `127.0.0.1` 代替 `localhost`。
---
### 坑 8Windows 没有 make 命令
**问题**CI 里用 `make test-unit`,本地 Windows 没有 make。
**解决**:直接用 Makefile 里的原始命令:
```bash
# 代替 make test-unit
go test -tags=unit ./...
# 代替 make test-integration
go test -tags=integration ./...
```
---
### 坑 9Ent Schema 修改后必须重新生成
**问题**:修改 `ent/schema/*.go` 后,代码不生效。
**解决**
```bash
cd backend
go generate ./ent # 重新生成 ent 代码
git add ent/ # 生成的文件也要提交
```
---
### 坑 10PR 提交前检查清单
提交 PR 前务必本地验证:
- [ ] `go test -tags=unit ./...` 通过
- [ ] `go test -tags=integration ./...` 通过
- [ ] `golangci-lint run ./...` 无新增问题
- [ ] `pnpm-lock.yaml` 已同步(如果改了 package.json
- [ ] 所有 test stub 补全新接口方法(如果改了 interface
- [ ] Ent 生成的代码已提交(如果改了 schema
## 五、常用命令速查
### 数据库操作
```bash
# 连接数据库
psql -U sub2api -h 127.0.0.1 -d sub2api
# 查看所有用户
psql -U postgres -h 127.0.0.1 -c "\du"
# 查看所有数据库
psql -U postgres -h 127.0.0.1 -c "\l"
# 执行 SQL 文件
psql -U sub2api -h 127.0.0.1 -d sub2api -f migration.sql
```
### Git 操作
```bash
# 同步上游
git fetch upstream
git checkout main
git merge upstream/main
git push origin main
# 创建功能分支
git checkout -b feature/xxx
# Rebase 到最新 main
git fetch upstream
git rebase upstream/main
```
### 前端操作
```bash
# 安装依赖(必须用 pnpm
cd frontend
pnpm install
# 开发服务器
pnpm dev
# 构建
pnpm build
```
### 后端操作
```bash
# 运行服务器
cd backend
go run ./cmd/server/
# 生成 Ent 代码
go generate ./ent
# 运行测试
go test -tags=unit ./...
go test -tags=integration ./...
# Lint 检查
golangci-lint run ./...
```
## 六、项目结构速览
```
sub2api-bmai/
├── backend/
│ ├── cmd/server/ # 主程序入口
│ ├── ent/ # Ent ORM 生成代码
│ │ └── schema/ # 数据库 Schema 定义
│ ├── internal/
│ │ ├── handler/ # HTTP 处理器
│ │ ├── service/ # 业务逻辑
│ │ ├── repository/ # 数据访问层
│ │ └── server/ # 服务器配置
│ ├── migrations/ # 数据库迁移脚本
│ └── config.yaml # 配置文件
├── frontend/
│ ├── src/
│ │ ├── api/ # API 调用
│ │ ├── components/ # Vue 组件
│ │ ├── views/ # 页面视图
│ │ ├── types/ # TypeScript 类型
│ │ └── i18n/ # 国际化
│ ├── package.json # 依赖配置
│ └── pnpm-lock.yaml # pnpm 锁文件(必须提交)
└── .claude/
└── CLAUDE.md # 本文档
```
## 七、参考资源
- [上游仓库](https://github.com/Wei-Shaw/sub2api)
- [Ent 文档](https://entgo.io/docs/getting-started)
- [Vue3 文档](https://vuejs.org/)
- [pnpm 文档](https://pnpm.io/)

View File

@@ -1 +1 @@
0.1.70.2 0.1.74.2

View File

@@ -102,7 +102,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig) proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig)
proxyLatencyCache := repository.NewProxyLatencyCache(redisClient) proxyLatencyCache := repository.NewProxyLatencyCache(redisClient)
adminService := service.NewAdminService(userRepository, groupRepository, accountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator) adminService := service.NewAdminService(userRepository, groupRepository, accountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator)
adminUserHandler := admin.NewUserHandler(adminService) concurrencyCache := repository.ProvideConcurrencyCache(redisClient, configConfig)
concurrencyService := service.ProvideConcurrencyService(concurrencyCache, accountRepository, configConfig)
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService)
groupHandler := admin.NewGroupHandler(adminService) groupHandler := admin.NewGroupHandler(adminService)
claudeOAuthClient := repository.NewClaudeOAuthClient() claudeOAuthClient := repository.NewClaudeOAuthClient()
oAuthService := service.NewOAuthService(proxyRepository, claudeOAuthClient) oAuthService := service.NewOAuthService(proxyRepository, claudeOAuthClient)
@@ -126,13 +128,11 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, usageCache, identityCache) accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, usageCache, identityCache)
geminiTokenProvider := service.NewGeminiTokenProvider(accountRepository, geminiTokenCache, geminiOAuthService) geminiTokenProvider := service.NewGeminiTokenProvider(accountRepository, geminiTokenCache, geminiOAuthService)
gatewayCache := repository.NewGatewayCache(redisClient) gatewayCache := repository.NewGatewayCache(redisClient)
antigravityTokenProvider := service.NewAntigravityTokenProvider(accountRepository, geminiTokenCache, antigravityOAuthService)
schedulerOutboxRepository := repository.NewSchedulerOutboxRepository(db) schedulerOutboxRepository := repository.NewSchedulerOutboxRepository(db)
schedulerSnapshotService := service.ProvideSchedulerSnapshotService(schedulerCache, schedulerOutboxRepository, accountRepository, groupRepository, configConfig) schedulerSnapshotService := service.ProvideSchedulerSnapshotService(schedulerCache, schedulerOutboxRepository, accountRepository, groupRepository, configConfig)
antigravityTokenProvider := service.NewAntigravityTokenProvider(accountRepository, geminiTokenCache, antigravityOAuthService)
antigravityGatewayService := service.NewAntigravityGatewayService(accountRepository, gatewayCache, schedulerSnapshotService, antigravityTokenProvider, rateLimitService, httpUpstream, settingService) antigravityGatewayService := service.NewAntigravityGatewayService(accountRepository, gatewayCache, schedulerSnapshotService, antigravityTokenProvider, rateLimitService, httpUpstream, settingService)
accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, antigravityGatewayService, httpUpstream, configConfig) accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, antigravityGatewayService, httpUpstream, configConfig)
concurrencyCache := repository.ProvideConcurrencyCache(redisClient, configConfig)
concurrencyService := service.ProvideConcurrencyService(concurrencyCache, accountRepository, configConfig)
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig) crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
sessionLimitCache := repository.ProvideSessionLimitCache(redisClient, configConfig) sessionLimitCache := repository.ProvideSessionLimitCache(redisClient, configConfig)
accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, compositeTokenCacheInvalidator) accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, compositeTokenCacheInvalidator)
@@ -154,7 +154,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
identityService := service.NewIdentityService(identityCache) identityService := service.NewIdentityService(identityCache)
deferredService := service.ProvideDeferredService(accountRepository, timingWheelService) deferredService := service.ProvideDeferredService(accountRepository, timingWheelService)
claudeTokenProvider := service.NewClaudeTokenProvider(accountRepository, geminiTokenCache, oAuthService) claudeTokenProvider := service.NewClaudeTokenProvider(accountRepository, geminiTokenCache, oAuthService)
gatewayService := service.NewGatewayService(accountRepository, groupRepository, usageLogRepository, userRepository, userSubscriptionRepository, userGroupRateRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, identityService, httpUpstream, deferredService, claudeTokenProvider, sessionLimitCache) digestSessionStore := service.NewDigestSessionStore()
gatewayService := service.NewGatewayService(accountRepository, groupRepository, usageLogRepository, userRepository, userSubscriptionRepository, userGroupRateRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, identityService, httpUpstream, deferredService, claudeTokenProvider, sessionLimitCache, digestSessionStore)
openAITokenProvider := service.NewOpenAITokenProvider(accountRepository, geminiTokenCache, openAIOAuthService) openAITokenProvider := service.NewOpenAITokenProvider(accountRepository, geminiTokenCache, openAIOAuthService)
openAIGatewayService := service.NewOpenAIGatewayService(accountRepository, usageLogRepository, userRepository, userSubscriptionRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, httpUpstream, deferredService, openAITokenProvider) openAIGatewayService := service.NewOpenAIGatewayService(accountRepository, usageLogRepository, userRepository, userSubscriptionRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, httpUpstream, deferredService, openAITokenProvider)
geminiMessagesCompatService := service.NewGeminiMessagesCompatService(accountRepository, groupRepository, gatewayCache, schedulerSnapshotService, geminiTokenProvider, rateLimitService, httpUpstream, antigravityGatewayService, configConfig) geminiMessagesCompatService := service.NewGeminiMessagesCompatService(accountRepository, groupRepository, gatewayCache, schedulerSnapshotService, geminiTokenProvider, rateLimitService, httpUpstream, antigravityGatewayService, configConfig)

View File

@@ -66,6 +66,8 @@ type Group struct {
McpXMLInject bool `json:"mcp_xml_inject,omitempty"` McpXMLInject bool `json:"mcp_xml_inject,omitempty"`
// 支持的模型系列claude, gemini_text, gemini_image // 支持的模型系列claude, gemini_text, gemini_image
SupportedModelScopes []string `json:"supported_model_scopes,omitempty"` SupportedModelScopes []string `json:"supported_model_scopes,omitempty"`
// 分组显示排序,数值越小越靠前
SortOrder int `json:"sort_order,omitempty"`
// Edges holds the relations/edges for other nodes in the graph. // Edges holds the relations/edges for other nodes in the graph.
// The values are being populated by the GroupQuery when eager-loading is set. // The values are being populated by the GroupQuery when eager-loading is set.
Edges GroupEdges `json:"edges"` Edges GroupEdges `json:"edges"`
@@ -178,7 +180,7 @@ func (*Group) scanValues(columns []string) ([]any, error) {
values[i] = new(sql.NullBool) values[i] = new(sql.NullBool)
case group.FieldRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k: case group.FieldRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k:
values[i] = new(sql.NullFloat64) values[i] = new(sql.NullFloat64)
case group.FieldID, group.FieldDefaultValidityDays, group.FieldFallbackGroupID, group.FieldFallbackGroupIDOnInvalidRequest: case group.FieldID, group.FieldDefaultValidityDays, group.FieldFallbackGroupID, group.FieldFallbackGroupIDOnInvalidRequest, group.FieldSortOrder:
values[i] = new(sql.NullInt64) values[i] = new(sql.NullInt64)
case group.FieldName, group.FieldDescription, group.FieldStatus, group.FieldPlatform, group.FieldSubscriptionType: case group.FieldName, group.FieldDescription, group.FieldStatus, group.FieldPlatform, group.FieldSubscriptionType:
values[i] = new(sql.NullString) values[i] = new(sql.NullString)
@@ -363,6 +365,12 @@ func (_m *Group) assignValues(columns []string, values []any) error {
return fmt.Errorf("unmarshal field supported_model_scopes: %w", err) return fmt.Errorf("unmarshal field supported_model_scopes: %w", err)
} }
} }
case group.FieldSortOrder:
if value, ok := values[i].(*sql.NullInt64); !ok {
return fmt.Errorf("unexpected type %T for field sort_order", values[i])
} else if value.Valid {
_m.SortOrder = int(value.Int64)
}
default: default:
_m.selectValues.Set(columns[i], values[i]) _m.selectValues.Set(columns[i], values[i])
} }
@@ -530,6 +538,9 @@ func (_m *Group) String() string {
builder.WriteString(", ") builder.WriteString(", ")
builder.WriteString("supported_model_scopes=") builder.WriteString("supported_model_scopes=")
builder.WriteString(fmt.Sprintf("%v", _m.SupportedModelScopes)) builder.WriteString(fmt.Sprintf("%v", _m.SupportedModelScopes))
builder.WriteString(", ")
builder.WriteString("sort_order=")
builder.WriteString(fmt.Sprintf("%v", _m.SortOrder))
builder.WriteByte(')') builder.WriteByte(')')
return builder.String() return builder.String()
} }

View File

@@ -63,6 +63,8 @@ const (
FieldMcpXMLInject = "mcp_xml_inject" FieldMcpXMLInject = "mcp_xml_inject"
// FieldSupportedModelScopes holds the string denoting the supported_model_scopes field in the database. // FieldSupportedModelScopes holds the string denoting the supported_model_scopes field in the database.
FieldSupportedModelScopes = "supported_model_scopes" FieldSupportedModelScopes = "supported_model_scopes"
// FieldSortOrder holds the string denoting the sort_order field in the database.
FieldSortOrder = "sort_order"
// EdgeAPIKeys holds the string denoting the api_keys edge name in mutations. // EdgeAPIKeys holds the string denoting the api_keys edge name in mutations.
EdgeAPIKeys = "api_keys" EdgeAPIKeys = "api_keys"
// EdgeRedeemCodes holds the string denoting the redeem_codes edge name in mutations. // EdgeRedeemCodes holds the string denoting the redeem_codes edge name in mutations.
@@ -162,6 +164,7 @@ var Columns = []string{
FieldModelRoutingEnabled, FieldModelRoutingEnabled,
FieldMcpXMLInject, FieldMcpXMLInject,
FieldSupportedModelScopes, FieldSupportedModelScopes,
FieldSortOrder,
} }
var ( var (
@@ -225,6 +228,8 @@ var (
DefaultMcpXMLInject bool DefaultMcpXMLInject bool
// DefaultSupportedModelScopes holds the default value on creation for the "supported_model_scopes" field. // DefaultSupportedModelScopes holds the default value on creation for the "supported_model_scopes" field.
DefaultSupportedModelScopes []string DefaultSupportedModelScopes []string
// DefaultSortOrder holds the default value on creation for the "sort_order" field.
DefaultSortOrder int
) )
// OrderOption defines the ordering options for the Group queries. // OrderOption defines the ordering options for the Group queries.
@@ -345,6 +350,11 @@ func ByMcpXMLInject(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldMcpXMLInject, opts...).ToFunc() return sql.OrderByField(FieldMcpXMLInject, opts...).ToFunc()
} }
// BySortOrder orders the results by the sort_order field.
func BySortOrder(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldSortOrder, opts...).ToFunc()
}
// ByAPIKeysCount orders the results by api_keys count. // ByAPIKeysCount orders the results by api_keys count.
func ByAPIKeysCount(opts ...sql.OrderTermOption) OrderOption { func ByAPIKeysCount(opts ...sql.OrderTermOption) OrderOption {
return func(s *sql.Selector) { return func(s *sql.Selector) {

View File

@@ -165,6 +165,11 @@ func McpXMLInject(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldMcpXMLInject, v)) return predicate.Group(sql.FieldEQ(FieldMcpXMLInject, v))
} }
// SortOrder applies equality check predicate on the "sort_order" field. It's identical to SortOrderEQ.
func SortOrder(v int) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldSortOrder, v))
}
// CreatedAtEQ applies the EQ predicate on the "created_at" field. // CreatedAtEQ applies the EQ predicate on the "created_at" field.
func CreatedAtEQ(v time.Time) predicate.Group { func CreatedAtEQ(v time.Time) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldCreatedAt, v)) return predicate.Group(sql.FieldEQ(FieldCreatedAt, v))
@@ -1160,6 +1165,46 @@ func McpXMLInjectNEQ(v bool) predicate.Group {
return predicate.Group(sql.FieldNEQ(FieldMcpXMLInject, v)) return predicate.Group(sql.FieldNEQ(FieldMcpXMLInject, v))
} }
// SortOrderEQ applies the EQ predicate on the "sort_order" field.
func SortOrderEQ(v int) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldSortOrder, v))
}
// SortOrderNEQ applies the NEQ predicate on the "sort_order" field.
func SortOrderNEQ(v int) predicate.Group {
return predicate.Group(sql.FieldNEQ(FieldSortOrder, v))
}
// SortOrderIn applies the In predicate on the "sort_order" field.
func SortOrderIn(vs ...int) predicate.Group {
return predicate.Group(sql.FieldIn(FieldSortOrder, vs...))
}
// SortOrderNotIn applies the NotIn predicate on the "sort_order" field.
func SortOrderNotIn(vs ...int) predicate.Group {
return predicate.Group(sql.FieldNotIn(FieldSortOrder, vs...))
}
// SortOrderGT applies the GT predicate on the "sort_order" field.
func SortOrderGT(v int) predicate.Group {
return predicate.Group(sql.FieldGT(FieldSortOrder, v))
}
// SortOrderGTE applies the GTE predicate on the "sort_order" field.
func SortOrderGTE(v int) predicate.Group {
return predicate.Group(sql.FieldGTE(FieldSortOrder, v))
}
// SortOrderLT applies the LT predicate on the "sort_order" field.
func SortOrderLT(v int) predicate.Group {
return predicate.Group(sql.FieldLT(FieldSortOrder, v))
}
// SortOrderLTE applies the LTE predicate on the "sort_order" field.
func SortOrderLTE(v int) predicate.Group {
return predicate.Group(sql.FieldLTE(FieldSortOrder, v))
}
// HasAPIKeys applies the HasEdge predicate on the "api_keys" edge. // HasAPIKeys applies the HasEdge predicate on the "api_keys" edge.
func HasAPIKeys() predicate.Group { func HasAPIKeys() predicate.Group {
return predicate.Group(func(s *sql.Selector) { return predicate.Group(func(s *sql.Selector) {

View File

@@ -340,6 +340,20 @@ func (_c *GroupCreate) SetSupportedModelScopes(v []string) *GroupCreate {
return _c return _c
} }
// SetSortOrder sets the "sort_order" field.
func (_c *GroupCreate) SetSortOrder(v int) *GroupCreate {
_c.mutation.SetSortOrder(v)
return _c
}
// SetNillableSortOrder sets the "sort_order" field if the given value is not nil.
func (_c *GroupCreate) SetNillableSortOrder(v *int) *GroupCreate {
if v != nil {
_c.SetSortOrder(*v)
}
return _c
}
// AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by IDs. // AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by IDs.
func (_c *GroupCreate) AddAPIKeyIDs(ids ...int64) *GroupCreate { func (_c *GroupCreate) AddAPIKeyIDs(ids ...int64) *GroupCreate {
_c.mutation.AddAPIKeyIDs(ids...) _c.mutation.AddAPIKeyIDs(ids...)
@@ -521,6 +535,10 @@ func (_c *GroupCreate) defaults() error {
v := group.DefaultSupportedModelScopes v := group.DefaultSupportedModelScopes
_c.mutation.SetSupportedModelScopes(v) _c.mutation.SetSupportedModelScopes(v)
} }
if _, ok := _c.mutation.SortOrder(); !ok {
v := group.DefaultSortOrder
_c.mutation.SetSortOrder(v)
}
return nil return nil
} }
@@ -585,6 +603,9 @@ func (_c *GroupCreate) check() error {
if _, ok := _c.mutation.SupportedModelScopes(); !ok { if _, ok := _c.mutation.SupportedModelScopes(); !ok {
return &ValidationError{Name: "supported_model_scopes", err: errors.New(`ent: missing required field "Group.supported_model_scopes"`)} return &ValidationError{Name: "supported_model_scopes", err: errors.New(`ent: missing required field "Group.supported_model_scopes"`)}
} }
if _, ok := _c.mutation.SortOrder(); !ok {
return &ValidationError{Name: "sort_order", err: errors.New(`ent: missing required field "Group.sort_order"`)}
}
return nil return nil
} }
@@ -708,6 +729,10 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) {
_spec.SetField(group.FieldSupportedModelScopes, field.TypeJSON, value) _spec.SetField(group.FieldSupportedModelScopes, field.TypeJSON, value)
_node.SupportedModelScopes = value _node.SupportedModelScopes = value
} }
if value, ok := _c.mutation.SortOrder(); ok {
_spec.SetField(group.FieldSortOrder, field.TypeInt, value)
_node.SortOrder = value
}
if nodes := _c.mutation.APIKeysIDs(); len(nodes) > 0 { if nodes := _c.mutation.APIKeysIDs(); len(nodes) > 0 {
edge := &sqlgraph.EdgeSpec{ edge := &sqlgraph.EdgeSpec{
Rel: sqlgraph.O2M, Rel: sqlgraph.O2M,
@@ -1266,6 +1291,24 @@ func (u *GroupUpsert) UpdateSupportedModelScopes() *GroupUpsert {
return u return u
} }
// SetSortOrder sets the "sort_order" field.
func (u *GroupUpsert) SetSortOrder(v int) *GroupUpsert {
u.Set(group.FieldSortOrder, v)
return u
}
// UpdateSortOrder sets the "sort_order" field to the value that was provided on create.
func (u *GroupUpsert) UpdateSortOrder() *GroupUpsert {
u.SetExcluded(group.FieldSortOrder)
return u
}
// AddSortOrder adds v to the "sort_order" field.
func (u *GroupUpsert) AddSortOrder(v int) *GroupUpsert {
u.Add(group.FieldSortOrder, v)
return u
}
// UpdateNewValues updates the mutable fields using the new values that were set on create. // UpdateNewValues updates the mutable fields using the new values that were set on create.
// Using this option is equivalent to using: // Using this option is equivalent to using:
// //
@@ -1780,6 +1823,27 @@ func (u *GroupUpsertOne) UpdateSupportedModelScopes() *GroupUpsertOne {
}) })
} }
// SetSortOrder sets the "sort_order" field.
func (u *GroupUpsertOne) SetSortOrder(v int) *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.SetSortOrder(v)
})
}
// AddSortOrder adds v to the "sort_order" field.
func (u *GroupUpsertOne) AddSortOrder(v int) *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.AddSortOrder(v)
})
}
// UpdateSortOrder sets the "sort_order" field to the value that was provided on create.
func (u *GroupUpsertOne) UpdateSortOrder() *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.UpdateSortOrder()
})
}
// Exec executes the query. // Exec executes the query.
func (u *GroupUpsertOne) Exec(ctx context.Context) error { func (u *GroupUpsertOne) Exec(ctx context.Context) error {
if len(u.create.conflict) == 0 { if len(u.create.conflict) == 0 {
@@ -2460,6 +2524,27 @@ func (u *GroupUpsertBulk) UpdateSupportedModelScopes() *GroupUpsertBulk {
}) })
} }
// SetSortOrder sets the "sort_order" field.
func (u *GroupUpsertBulk) SetSortOrder(v int) *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.SetSortOrder(v)
})
}
// AddSortOrder adds v to the "sort_order" field.
func (u *GroupUpsertBulk) AddSortOrder(v int) *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.AddSortOrder(v)
})
}
// UpdateSortOrder sets the "sort_order" field to the value that was provided on create.
func (u *GroupUpsertBulk) UpdateSortOrder() *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.UpdateSortOrder()
})
}
// Exec executes the query. // Exec executes the query.
func (u *GroupUpsertBulk) Exec(ctx context.Context) error { func (u *GroupUpsertBulk) Exec(ctx context.Context) error {
if u.create.err != nil { if u.create.err != nil {

View File

@@ -475,6 +475,27 @@ func (_u *GroupUpdate) AppendSupportedModelScopes(v []string) *GroupUpdate {
return _u return _u
} }
// SetSortOrder sets the "sort_order" field.
func (_u *GroupUpdate) SetSortOrder(v int) *GroupUpdate {
_u.mutation.ResetSortOrder()
_u.mutation.SetSortOrder(v)
return _u
}
// SetNillableSortOrder sets the "sort_order" field if the given value is not nil.
func (_u *GroupUpdate) SetNillableSortOrder(v *int) *GroupUpdate {
if v != nil {
_u.SetSortOrder(*v)
}
return _u
}
// AddSortOrder adds value to the "sort_order" field.
func (_u *GroupUpdate) AddSortOrder(v int) *GroupUpdate {
_u.mutation.AddSortOrder(v)
return _u
}
// AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by IDs. // AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by IDs.
func (_u *GroupUpdate) AddAPIKeyIDs(ids ...int64) *GroupUpdate { func (_u *GroupUpdate) AddAPIKeyIDs(ids ...int64) *GroupUpdate {
_u.mutation.AddAPIKeyIDs(ids...) _u.mutation.AddAPIKeyIDs(ids...)
@@ -912,6 +933,12 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) {
sqljson.Append(u, group.FieldSupportedModelScopes, value) sqljson.Append(u, group.FieldSupportedModelScopes, value)
}) })
} }
if value, ok := _u.mutation.SortOrder(); ok {
_spec.SetField(group.FieldSortOrder, field.TypeInt, value)
}
if value, ok := _u.mutation.AddedSortOrder(); ok {
_spec.AddField(group.FieldSortOrder, field.TypeInt, value)
}
if _u.mutation.APIKeysCleared() { if _u.mutation.APIKeysCleared() {
edge := &sqlgraph.EdgeSpec{ edge := &sqlgraph.EdgeSpec{
Rel: sqlgraph.O2M, Rel: sqlgraph.O2M,
@@ -1666,6 +1693,27 @@ func (_u *GroupUpdateOne) AppendSupportedModelScopes(v []string) *GroupUpdateOne
return _u return _u
} }
// SetSortOrder sets the "sort_order" field.
func (_u *GroupUpdateOne) SetSortOrder(v int) *GroupUpdateOne {
_u.mutation.ResetSortOrder()
_u.mutation.SetSortOrder(v)
return _u
}
// SetNillableSortOrder sets the "sort_order" field if the given value is not nil.
func (_u *GroupUpdateOne) SetNillableSortOrder(v *int) *GroupUpdateOne {
if v != nil {
_u.SetSortOrder(*v)
}
return _u
}
// AddSortOrder adds value to the "sort_order" field.
func (_u *GroupUpdateOne) AddSortOrder(v int) *GroupUpdateOne {
_u.mutation.AddSortOrder(v)
return _u
}
// AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by IDs. // AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by IDs.
func (_u *GroupUpdateOne) AddAPIKeyIDs(ids ...int64) *GroupUpdateOne { func (_u *GroupUpdateOne) AddAPIKeyIDs(ids ...int64) *GroupUpdateOne {
_u.mutation.AddAPIKeyIDs(ids...) _u.mutation.AddAPIKeyIDs(ids...)
@@ -2133,6 +2181,12 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error)
sqljson.Append(u, group.FieldSupportedModelScopes, value) sqljson.Append(u, group.FieldSupportedModelScopes, value)
}) })
} }
if value, ok := _u.mutation.SortOrder(); ok {
_spec.SetField(group.FieldSortOrder, field.TypeInt, value)
}
if value, ok := _u.mutation.AddedSortOrder(); ok {
_spec.AddField(group.FieldSortOrder, field.TypeInt, value)
}
if _u.mutation.APIKeysCleared() { if _u.mutation.APIKeysCleared() {
edge := &sqlgraph.EdgeSpec{ edge := &sqlgraph.EdgeSpec{
Rel: sqlgraph.O2M, Rel: sqlgraph.O2M,

View File

@@ -372,6 +372,7 @@ var (
{Name: "model_routing_enabled", Type: field.TypeBool, Default: false}, {Name: "model_routing_enabled", Type: field.TypeBool, Default: false},
{Name: "mcp_xml_inject", Type: field.TypeBool, Default: true}, {Name: "mcp_xml_inject", Type: field.TypeBool, Default: true},
{Name: "supported_model_scopes", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}}, {Name: "supported_model_scopes", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
{Name: "sort_order", Type: field.TypeInt, Default: 0},
} }
// GroupsTable holds the schema information for the "groups" table. // GroupsTable holds the schema information for the "groups" table.
GroupsTable = &schema.Table{ GroupsTable = &schema.Table{
@@ -404,6 +405,11 @@ var (
Unique: false, Unique: false,
Columns: []*schema.Column{GroupsColumns[3]}, Columns: []*schema.Column{GroupsColumns[3]},
}, },
{
Name: "group_sort_order",
Unique: false,
Columns: []*schema.Column{GroupsColumns[25]},
},
}, },
} }
// PromoCodesColumns holds the columns for the "promo_codes" table. // PromoCodesColumns holds the columns for the "promo_codes" table.

View File

@@ -7059,6 +7059,8 @@ type GroupMutation struct {
mcp_xml_inject *bool mcp_xml_inject *bool
supported_model_scopes *[]string supported_model_scopes *[]string
appendsupported_model_scopes []string appendsupported_model_scopes []string
sort_order *int
addsort_order *int
clearedFields map[string]struct{} clearedFields map[string]struct{}
api_keys map[int64]struct{} api_keys map[int64]struct{}
removedapi_keys map[int64]struct{} removedapi_keys map[int64]struct{}
@@ -8411,6 +8413,62 @@ func (m *GroupMutation) ResetSupportedModelScopes() {
m.appendsupported_model_scopes = nil m.appendsupported_model_scopes = nil
} }
// SetSortOrder sets the "sort_order" field.
func (m *GroupMutation) SetSortOrder(i int) {
m.sort_order = &i
m.addsort_order = nil
}
// SortOrder returns the value of the "sort_order" field in the mutation.
func (m *GroupMutation) SortOrder() (r int, exists bool) {
v := m.sort_order
if v == nil {
return
}
return *v, true
}
// OldSortOrder returns the old "sort_order" field's value of the Group entity.
// If the Group object wasn't provided to the builder, the object is fetched from the database.
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
func (m *GroupMutation) OldSortOrder(ctx context.Context) (v int, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldSortOrder is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldSortOrder requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldSortOrder: %w", err)
}
return oldValue.SortOrder, nil
}
// AddSortOrder adds i to the "sort_order" field.
func (m *GroupMutation) AddSortOrder(i int) {
if m.addsort_order != nil {
*m.addsort_order += i
} else {
m.addsort_order = &i
}
}
// AddedSortOrder returns the value that was added to the "sort_order" field in this mutation.
func (m *GroupMutation) AddedSortOrder() (r int, exists bool) {
v := m.addsort_order
if v == nil {
return
}
return *v, true
}
// ResetSortOrder resets all changes to the "sort_order" field.
func (m *GroupMutation) ResetSortOrder() {
m.sort_order = nil
m.addsort_order = nil
}
// AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by ids. // AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by ids.
func (m *GroupMutation) AddAPIKeyIDs(ids ...int64) { func (m *GroupMutation) AddAPIKeyIDs(ids ...int64) {
if m.api_keys == nil { if m.api_keys == nil {
@@ -8769,7 +8827,7 @@ func (m *GroupMutation) Type() string {
// order to get all numeric fields that were incremented/decremented, call // order to get all numeric fields that were incremented/decremented, call
// AddedFields(). // AddedFields().
func (m *GroupMutation) Fields() []string { func (m *GroupMutation) Fields() []string {
fields := make([]string, 0, 24) fields := make([]string, 0, 25)
if m.created_at != nil { if m.created_at != nil {
fields = append(fields, group.FieldCreatedAt) fields = append(fields, group.FieldCreatedAt)
} }
@@ -8842,6 +8900,9 @@ func (m *GroupMutation) Fields() []string {
if m.supported_model_scopes != nil { if m.supported_model_scopes != nil {
fields = append(fields, group.FieldSupportedModelScopes) fields = append(fields, group.FieldSupportedModelScopes)
} }
if m.sort_order != nil {
fields = append(fields, group.FieldSortOrder)
}
return fields return fields
} }
@@ -8898,6 +8959,8 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) {
return m.McpXMLInject() return m.McpXMLInject()
case group.FieldSupportedModelScopes: case group.FieldSupportedModelScopes:
return m.SupportedModelScopes() return m.SupportedModelScopes()
case group.FieldSortOrder:
return m.SortOrder()
} }
return nil, false return nil, false
} }
@@ -8955,6 +9018,8 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e
return m.OldMcpXMLInject(ctx) return m.OldMcpXMLInject(ctx)
case group.FieldSupportedModelScopes: case group.FieldSupportedModelScopes:
return m.OldSupportedModelScopes(ctx) return m.OldSupportedModelScopes(ctx)
case group.FieldSortOrder:
return m.OldSortOrder(ctx)
} }
return nil, fmt.Errorf("unknown Group field %s", name) return nil, fmt.Errorf("unknown Group field %s", name)
} }
@@ -9132,6 +9197,13 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error {
} }
m.SetSupportedModelScopes(v) m.SetSupportedModelScopes(v)
return nil return nil
case group.FieldSortOrder:
v, ok := value.(int)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetSortOrder(v)
return nil
} }
return fmt.Errorf("unknown Group field %s", name) return fmt.Errorf("unknown Group field %s", name)
} }
@@ -9170,6 +9242,9 @@ func (m *GroupMutation) AddedFields() []string {
if m.addfallback_group_id_on_invalid_request != nil { if m.addfallback_group_id_on_invalid_request != nil {
fields = append(fields, group.FieldFallbackGroupIDOnInvalidRequest) fields = append(fields, group.FieldFallbackGroupIDOnInvalidRequest)
} }
if m.addsort_order != nil {
fields = append(fields, group.FieldSortOrder)
}
return fields return fields
} }
@@ -9198,6 +9273,8 @@ func (m *GroupMutation) AddedField(name string) (ent.Value, bool) {
return m.AddedFallbackGroupID() return m.AddedFallbackGroupID()
case group.FieldFallbackGroupIDOnInvalidRequest: case group.FieldFallbackGroupIDOnInvalidRequest:
return m.AddedFallbackGroupIDOnInvalidRequest() return m.AddedFallbackGroupIDOnInvalidRequest()
case group.FieldSortOrder:
return m.AddedSortOrder()
} }
return nil, false return nil, false
} }
@@ -9277,6 +9354,13 @@ func (m *GroupMutation) AddField(name string, value ent.Value) error {
} }
m.AddFallbackGroupIDOnInvalidRequest(v) m.AddFallbackGroupIDOnInvalidRequest(v)
return nil return nil
case group.FieldSortOrder:
v, ok := value.(int)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.AddSortOrder(v)
return nil
} }
return fmt.Errorf("unknown Group numeric field %s", name) return fmt.Errorf("unknown Group numeric field %s", name)
} }
@@ -9445,6 +9529,9 @@ func (m *GroupMutation) ResetField(name string) error {
case group.FieldSupportedModelScopes: case group.FieldSupportedModelScopes:
m.ResetSupportedModelScopes() m.ResetSupportedModelScopes()
return nil return nil
case group.FieldSortOrder:
m.ResetSortOrder()
return nil
} }
return fmt.Errorf("unknown Group field %s", name) return fmt.Errorf("unknown Group field %s", name)
} }

View File

@@ -409,6 +409,10 @@ func init() {
groupDescSupportedModelScopes := groupFields[20].Descriptor() groupDescSupportedModelScopes := groupFields[20].Descriptor()
// group.DefaultSupportedModelScopes holds the default value on creation for the supported_model_scopes field. // group.DefaultSupportedModelScopes holds the default value on creation for the supported_model_scopes field.
group.DefaultSupportedModelScopes = groupDescSupportedModelScopes.Default.([]string) group.DefaultSupportedModelScopes = groupDescSupportedModelScopes.Default.([]string)
// groupDescSortOrder is the schema descriptor for sort_order field.
groupDescSortOrder := groupFields[21].Descriptor()
// group.DefaultSortOrder holds the default value on creation for the sort_order field.
group.DefaultSortOrder = groupDescSortOrder.Default.(int)
promocodeFields := schema.PromoCode{}.Fields() promocodeFields := schema.PromoCode{}.Fields()
_ = promocodeFields _ = promocodeFields
// promocodeDescCode is the schema descriptor for code field. // promocodeDescCode is the schema descriptor for code field.

View File

@@ -121,6 +121,11 @@ func (Group) Fields() []ent.Field {
Default([]string{"claude", "gemini_text", "gemini_image"}). Default([]string{"claude", "gemini_text", "gemini_image"}).
SchemaType(map[string]string{dialect.Postgres: "jsonb"}). SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
Comment("支持的模型系列claude, gemini_text, gemini_image"), Comment("支持的模型系列claude, gemini_text, gemini_image"),
// 分组排序 (added by migration 052)
field.Int("sort_order").
Default(0).
Comment("分组显示排序,数值越小越靠前"),
} }
} }
@@ -149,5 +154,6 @@ func (Group) Indexes() []ent.Index {
index.Fields("subscription_type"), index.Fields("subscription_type"),
index.Fields("is_exclusive"), index.Fields("is_exclusive"),
index.Fields("deleted_at"), index.Fields("deleted_at"),
index.Fields("sort_order"),
} }
} }

View File

@@ -103,6 +103,7 @@ require (
github.com/ncruces/go-strftime v1.0.0 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect
github.com/opencontainers/go-digest v1.0.0 // indirect github.com/opencontainers/go-digest v1.0.0 // indirect
github.com/opencontainers/image-spec v1.1.1 // indirect github.com/opencontainers/image-spec v1.1.1 // indirect
github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
github.com/pelletier/go-toml/v2 v2.2.2 // indirect github.com/pelletier/go-toml/v2 v2.2.2 // indirect
github.com/pkg/errors v0.9.1 // indirect github.com/pkg/errors v0.9.1 // indirect
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect

View File

@@ -135,6 +135,8 @@ github.com/icholy/digest v1.1.0 h1:HfGg9Irj7i+IX1o1QAmPfIBNu/Q5A5Tu3n/MED9k9H4=
github.com/icholy/digest v1.1.0/go.mod h1:QNrsSGQ5v7v9cReDI0+eyjsXGUoRSUZQHeQ5C4XLa0Y= github.com/icholy/digest v1.1.0/go.mod h1:QNrsSGQ5v7v9cReDI0+eyjsXGUoRSUZQHeQ5C4XLa0Y=
github.com/imroc/req/v3 v3.57.0 h1:LMTUjNRUybUkTPn8oJDq8Kg3JRBOBTcnDhKu7mzupKI= github.com/imroc/req/v3 v3.57.0 h1:LMTUjNRUybUkTPn8oJDq8Kg3JRBOBTcnDhKu7mzupKI=
github.com/imroc/req/v3 v3.57.0/go.mod h1:JL62ey1nvSLq81HORNcosvlf7SxZStONNqOprg0Pz00= github.com/imroc/req/v3 v3.57.0/go.mod h1:JL62ey1nvSLq81HORNcosvlf7SxZStONNqOprg0Pz00=
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
@@ -170,6 +172,8 @@ github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovk
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U=
github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
github.com/mattn/go-sqlite3 v1.14.17 h1:mCRHCLDUBXgpKAqIKsaAaAsrAlbkeomtRFKXh2L6YIM= github.com/mattn/go-sqlite3 v1.14.17 h1:mCRHCLDUBXgpKAqIKsaAaAsrAlbkeomtRFKXh2L6YIM=
github.com/mattn/go-sqlite3 v1.14.17/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg= github.com/mattn/go-sqlite3 v1.14.17/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
github.com/mdelapenya/tlscert v0.2.0 h1:7H81W6Z/4weDvZBNOfQte5GpIMo0lGYEeWbkGp5LJHI= github.com/mdelapenya/tlscert v0.2.0 h1:7H81W6Z/4weDvZBNOfQte5GpIMo0lGYEeWbkGp5LJHI=
@@ -203,10 +207,14 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec=
github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040=
github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M= github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M=
github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM= github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM=
github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs= github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
@@ -230,6 +238,8 @@ github.com/refraction-networking/utls v1.8.1 h1:yNY1kapmQU8JeM1sSw2H2asfTIwWxIkr
github.com/refraction-networking/utls v1.8.1/go.mod h1:jkSOEkLqn+S/jtpEHPOsVv/4V4EVnelwbMQl4vCWXAM= github.com/refraction-networking/utls v1.8.1/go.mod h1:jkSOEkLqn+S/jtpEHPOsVv/4V4EVnelwbMQl4vCWXAM=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
@@ -252,6 +262,8 @@ github.com/spf13/afero v1.11.0 h1:WJQKhtpdm3v2IzqG8VMqrr6Rf3UYpEF239Jy9wNepM8=
github.com/spf13/afero v1.11.0/go.mod h1:GH9Y3pIexgf1MTIWtNGyogA5MwRIDXGUr+hbWNoBjkY= github.com/spf13/afero v1.11.0/go.mod h1:GH9Y3pIexgf1MTIWtNGyogA5MwRIDXGUr+hbWNoBjkY=
github.com/spf13/cast v1.6.0 h1:GEiTHELF+vaR5dhz3VqZfFSzZjYbgeKDpBxQVS4GYJ0= github.com/spf13/cast v1.6.0 h1:GEiTHELF+vaR5dhz3VqZfFSzZjYbgeKDpBxQVS4GYJ0=
github.com/spf13/cast v1.6.0/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= github.com/spf13/cast v1.6.0/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0=
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
github.com/spf13/viper v1.18.2 h1:LUXCnvUvSM6FXAsj6nnfc8Q2tp1dIgUfY9Kc8GsSOiQ= github.com/spf13/viper v1.18.2 h1:LUXCnvUvSM6FXAsj6nnfc8Q2tp1dIgUfY9Kc8GsSOiQ=

View File

@@ -425,10 +425,17 @@ type TestAccountRequest struct {
} }
type SyncFromCRSRequest struct { type SyncFromCRSRequest struct {
BaseURL string `json:"base_url" binding:"required"` BaseURL string `json:"base_url" binding:"required"`
Username string `json:"username" binding:"required"` Username string `json:"username" binding:"required"`
Password string `json:"password" binding:"required"` Password string `json:"password" binding:"required"`
SyncProxies *bool `json:"sync_proxies"` SyncProxies *bool `json:"sync_proxies"`
SelectedAccountIDs []string `json:"selected_account_ids"`
}
type PreviewFromCRSRequest struct {
BaseURL string `json:"base_url" binding:"required"`
Username string `json:"username" binding:"required"`
Password string `json:"password" binding:"required"`
} }
// Test handles testing account connectivity with SSE streaming // Test handles testing account connectivity with SSE streaming
@@ -467,10 +474,11 @@ func (h *AccountHandler) SyncFromCRS(c *gin.Context) {
} }
result, err := h.crsSyncService.SyncFromCRS(c.Request.Context(), service.SyncFromCRSInput{ result, err := h.crsSyncService.SyncFromCRS(c.Request.Context(), service.SyncFromCRSInput{
BaseURL: req.BaseURL, BaseURL: req.BaseURL,
Username: req.Username, Username: req.Username,
Password: req.Password, Password: req.Password,
SyncProxies: syncProxies, SyncProxies: syncProxies,
SelectedAccountIDs: req.SelectedAccountIDs,
}) })
if err != nil { if err != nil {
// Provide detailed error message for CRS sync failures // Provide detailed error message for CRS sync failures
@@ -481,6 +489,28 @@ func (h *AccountHandler) SyncFromCRS(c *gin.Context) {
response.Success(c, result) response.Success(c, result)
} }
// PreviewFromCRS handles previewing accounts from CRS before sync
// POST /api/v1/admin/accounts/sync/crs/preview
func (h *AccountHandler) PreviewFromCRS(c *gin.Context) {
var req PreviewFromCRSRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
result, err := h.crsSyncService.PreviewFromCRS(c.Request.Context(), service.SyncFromCRSInput{
BaseURL: req.BaseURL,
Username: req.Username,
Password: req.Password,
})
if err != nil {
response.InternalError(c, "CRS preview failed: "+err.Error())
return
}
response.Success(c, result)
}
// Refresh handles refreshing account credentials // Refresh handles refreshing account credentials
// POST /api/v1/admin/accounts/:id/refresh // POST /api/v1/admin/accounts/:id/refresh
func (h *AccountHandler) Refresh(c *gin.Context) { func (h *AccountHandler) Refresh(c *gin.Context) {

View File

@@ -16,7 +16,7 @@ func setupAdminRouter() (*gin.Engine, *stubAdminService) {
router := gin.New() router := gin.New()
adminSvc := newStubAdminService() adminSvc := newStubAdminService()
userHandler := NewUserHandler(adminSvc) userHandler := NewUserHandler(adminSvc, nil)
groupHandler := NewGroupHandler(adminSvc) groupHandler := NewGroupHandler(adminSvc)
proxyHandler := NewProxyHandler(adminSvc) proxyHandler := NewProxyHandler(adminSvc)
redeemHandler := NewRedeemHandler(adminSvc) redeemHandler := NewRedeemHandler(adminSvc)

View File

@@ -357,5 +357,9 @@ func (s *stubAdminService) GetUserBalanceHistory(ctx context.Context, userID int
return s.redeems, int64(len(s.redeems)), 100.0, nil return s.redeems, int64(len(s.redeems)), 100.0, nil
} }
func (s *stubAdminService) UpdateGroupSortOrders(ctx context.Context, updates []service.GroupSortOrderUpdate) error {
return nil
}
// Ensure stub implements interface. // Ensure stub implements interface.
var _ service.AdminService = (*stubAdminService)(nil) var _ service.AdminService = (*stubAdminService)(nil)

View File

@@ -302,3 +302,36 @@ func (h *GroupHandler) GetGroupAPIKeys(c *gin.Context) {
} }
response.Paginated(c, outKeys, total, page, pageSize) response.Paginated(c, outKeys, total, page, pageSize)
} }
// UpdateSortOrderRequest represents the request to update group sort orders
type UpdateSortOrderRequest struct {
Updates []struct {
ID int64 `json:"id" binding:"required"`
SortOrder int `json:"sort_order"`
} `json:"updates" binding:"required,min=1"`
}
// UpdateSortOrder handles updating group sort orders
// PUT /api/v1/admin/groups/sort-order
func (h *GroupHandler) UpdateSortOrder(c *gin.Context) {
var req UpdateSortOrderRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
updates := make([]service.GroupSortOrderUpdate, 0, len(req.Updates))
for _, u := range req.Updates {
updates = append(updates, service.GroupSortOrderUpdate{
ID: u.ID,
SortOrder: u.SortOrder,
})
}
if err := h.adminService.UpdateGroupSortOrders(c.Request.Context(), updates); err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, gin.H{"message": "Sort order updated successfully"})
}

View File

@@ -11,15 +11,23 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
) )
// UserWithConcurrency wraps AdminUser with current concurrency info
type UserWithConcurrency struct {
dto.AdminUser
CurrentConcurrency int `json:"current_concurrency"`
}
// UserHandler handles admin user management // UserHandler handles admin user management
type UserHandler struct { type UserHandler struct {
adminService service.AdminService adminService service.AdminService
concurrencyService *service.ConcurrencyService
} }
// NewUserHandler creates a new admin user handler // NewUserHandler creates a new admin user handler
func NewUserHandler(adminService service.AdminService) *UserHandler { func NewUserHandler(adminService service.AdminService, concurrencyService *service.ConcurrencyService) *UserHandler {
return &UserHandler{ return &UserHandler{
adminService: adminService, adminService: adminService,
concurrencyService: concurrencyService,
} }
} }
@@ -87,10 +95,30 @@ func (h *UserHandler) List(c *gin.Context) {
return return
} }
out := make([]dto.AdminUser, 0, len(users)) // Batch get current concurrency (nil map if unavailable)
for i := range users { var loadInfo map[int64]*service.UserLoadInfo
out = append(out, *dto.UserFromServiceAdmin(&users[i])) if len(users) > 0 && h.concurrencyService != nil {
usersConcurrency := make([]service.UserWithConcurrency, len(users))
for i := range users {
usersConcurrency[i] = service.UserWithConcurrency{
ID: users[i].ID,
MaxConcurrency: users[i].Concurrency,
}
}
loadInfo, _ = h.concurrencyService.GetUsersLoadBatch(c.Request.Context(), usersConcurrency)
} }
// Build response with concurrency info
out := make([]UserWithConcurrency, len(users))
for i := range users {
out[i] = UserWithConcurrency{
AdminUser: *dto.UserFromServiceAdmin(&users[i]),
}
if info := loadInfo[users[i].ID]; info != nil {
out[i].CurrentConcurrency = info.CurrentConcurrency
}
}
response.Paginated(c, out, total, page, pageSize) response.Paginated(c, out, total, page, pageSize)
} }

View File

@@ -115,6 +115,7 @@ func GroupFromServiceAdmin(g *service.Group) *AdminGroup {
MCPXMLInject: g.MCPXMLInject, MCPXMLInject: g.MCPXMLInject,
SupportedModelScopes: g.SupportedModelScopes, SupportedModelScopes: g.SupportedModelScopes,
AccountCount: g.AccountCount, AccountCount: g.AccountCount,
SortOrder: g.SortOrder,
} }
if len(g.AccountGroups) > 0 { if len(g.AccountGroups) > 0 {
out.AccountGroups = make([]AccountGroup, 0, len(g.AccountGroups)) out.AccountGroups = make([]AccountGroup, 0, len(g.AccountGroups))

View File

@@ -2,11 +2,6 @@ package dto
import "time" import "time"
type ScopeRateLimitInfo struct {
ResetAt time.Time `json:"reset_at"`
RemainingSec int64 `json:"remaining_sec"`
}
type User struct { type User struct {
ID int64 `json:"id"` ID int64 `json:"id"`
Email string `json:"email"` Email string `json:"email"`
@@ -98,6 +93,9 @@ type AdminGroup struct {
SupportedModelScopes []string `json:"supported_model_scopes"` SupportedModelScopes []string `json:"supported_model_scopes"`
AccountGroups []AccountGroup `json:"account_groups,omitempty"` AccountGroups []AccountGroup `json:"account_groups,omitempty"`
AccountCount int64 `json:"account_count,omitempty"` AccountCount int64 `json:"account_count,omitempty"`
// 分组排序
SortOrder int `json:"sort_order"`
} }
type Account struct { type Account struct {
@@ -126,9 +124,6 @@ type Account struct {
RateLimitResetAt *time.Time `json:"rate_limit_reset_at"` RateLimitResetAt *time.Time `json:"rate_limit_reset_at"`
OverloadUntil *time.Time `json:"overload_until"` OverloadUntil *time.Time `json:"overload_until"`
// Antigravity scope 级限流状态(从 extra 提取)
ScopeRateLimits map[string]ScopeRateLimitInfo `json:"scope_rate_limits,omitempty"`
TempUnschedulableUntil *time.Time `json:"temp_unschedulable_until"` TempUnschedulableUntil *time.Time `json:"temp_unschedulable_until"`
TempUnschedulableReason string `json:"temp_unschedulable_reason"` TempUnschedulableReason string `json:"temp_unschedulable_reason"`

View File

@@ -13,6 +13,7 @@ import (
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/config" "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/antigravity"
"github.com/Wei-Shaw/sub2api/internal/pkg/claude" "github.com/Wei-Shaw/sub2api/internal/pkg/claude"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
@@ -114,7 +115,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
setOpsRequestContext(c, "", false, body) setOpsRequestContext(c, "", false, body)
parsedReq, err := service.ParseGatewayRequest(body) parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic)
if err != nil { if err != nil {
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
return return
@@ -203,6 +204,11 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
} }
// 计算粘性会话hash // 计算粘性会话hash
parsedReq.SessionContext = &service.SessionContext{
ClientIP: ip.GetClientIP(c),
UserAgent: c.GetHeader("User-Agent"),
APIKeyID: apiKey.ID,
}
sessionHash := h.gatewayService.GenerateSessionHash(parsedReq) sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
// 获取平台:优先使用强制平台(/antigravity 路由,中间件已设置 request.Context否则使用分组平台 // 获取平台:优先使用强制平台(/antigravity 路由,中间件已设置 request.Context否则使用分组平台
@@ -334,7 +340,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
if errors.As(err, &failoverErr) { if errors.As(err, &failoverErr) {
failedAccountIDs[account.ID] = struct{}{} failedAccountIDs[account.ID] = struct{}{}
lastFailoverErr = failoverErr lastFailoverErr = failoverErr
if failoverErr.ForceCacheBilling { if needForceCacheBilling(hasBoundSession, failoverErr) {
forceCacheBilling = true forceCacheBilling = true
} }
if switchCount >= maxAccountSwitches { if switchCount >= maxAccountSwitches {
@@ -343,6 +349,11 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
} }
switchCount++ switchCount++
log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches) log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
if account.Platform == service.PlatformAntigravity {
if !sleepFailoverDelay(c.Request.Context(), switchCount) {
return
}
}
continue continue
} }
// 错误响应已在Forward中处理这里只记录日志 // 错误响应已在Forward中处理这里只记录日志
@@ -482,7 +493,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
if switchCount > 0 { if switchCount > 0 {
requestCtx = context.WithValue(requestCtx, ctxkey.AccountSwitchCount, switchCount) requestCtx = context.WithValue(requestCtx, ctxkey.AccountSwitchCount, switchCount)
} }
if account.Platform == service.PlatformAntigravity { if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey {
result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, body, hasBoundSession) result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, body, hasBoundSession)
} else { } else {
result, err = h.gatewayService.Forward(requestCtx, c, account, parsedReq) result, err = h.gatewayService.Forward(requestCtx, c, account, parsedReq)
@@ -530,7 +541,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
if errors.As(err, &failoverErr) { if errors.As(err, &failoverErr) {
failedAccountIDs[account.ID] = struct{}{} failedAccountIDs[account.ID] = struct{}{}
lastFailoverErr = failoverErr lastFailoverErr = failoverErr
if failoverErr.ForceCacheBilling { if needForceCacheBilling(hasBoundSession, failoverErr) {
forceCacheBilling = true forceCacheBilling = true
} }
if switchCount >= maxAccountSwitches { if switchCount >= maxAccountSwitches {
@@ -539,6 +550,11 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
} }
switchCount++ switchCount++
log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches) log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
if account.Platform == service.PlatformAntigravity {
if !sleepFailoverDelay(c.Request.Context(), switchCount) {
return
}
}
continue continue
} }
// 错误响应已在Forward中处理这里只记录日志 // 错误响应已在Forward中处理这里只记录日志
@@ -801,6 +817,27 @@ func (h *GatewayHandler) handleConcurrencyError(c *gin.Context, err error, slotT
fmt.Sprintf("Concurrency limit exceeded for %s, please retry later", slotType), streamStarted) fmt.Sprintf("Concurrency limit exceeded for %s, please retry later", slotType), streamStarted)
} }
// needForceCacheBilling 判断 failover 时是否需要强制缓存计费
// 粘性会话切换账号、或上游明确标记时,将 input_tokens 转为 cache_read 计费
func needForceCacheBilling(hasBoundSession bool, failoverErr *service.UpstreamFailoverError) bool {
return hasBoundSession || (failoverErr != nil && failoverErr.ForceCacheBilling)
}
// sleepFailoverDelay 账号切换线性递增延时第1次0s、第2次1s、第3次2s…
// 返回 false 表示 context 已取消。
func sleepFailoverDelay(ctx context.Context, switchCount int) bool {
delay := time.Duration(switchCount-1) * time.Second
if delay <= 0 {
return true
}
select {
case <-ctx.Done():
return false
case <-time.After(delay):
return true
}
}
func (h *GatewayHandler) handleFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, platform string, streamStarted bool) { func (h *GatewayHandler) handleFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, platform string, streamStarted bool) {
statusCode := failoverErr.StatusCode statusCode := failoverErr.StatusCode
responseBody := failoverErr.ResponseBody responseBody := failoverErr.ResponseBody
@@ -934,7 +971,7 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) {
setOpsRequestContext(c, "", false, body) setOpsRequestContext(c, "", false, body)
parsedReq, err := service.ParseGatewayRequest(body) parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic)
if err != nil { if err != nil {
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
return return
@@ -962,6 +999,11 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) {
} }
// 计算粘性会话 hash // 计算粘性会话 hash
parsedReq.SessionContext = &service.SessionContext{
ClientIP: ip.GetClientIP(c),
UserAgent: c.GetHeader("User-Agent"),
APIKeyID: apiKey.ID,
}
sessionHash := h.gatewayService.GenerateSessionHash(parsedReq) sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
// 选择支持该模型的账号 // 选择支持该模型的账号

View File

@@ -14,6 +14,7 @@ import (
"strings" "strings"
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/gemini" "github.com/Wei-Shaw/sub2api/internal/pkg/gemini"
@@ -30,13 +31,6 @@ import (
// 匹配格式: /Users/xxx/.gemini/tmp/[64位十六进制哈希] // 匹配格式: /Users/xxx/.gemini/tmp/[64位十六进制哈希]
var geminiCLITmpDirRegex = regexp.MustCompile(`/\.gemini/tmp/([A-Fa-f0-9]{64})`) var geminiCLITmpDirRegex = regexp.MustCompile(`/\.gemini/tmp/([A-Fa-f0-9]{64})`)
func isGeminiCLIRequest(c *gin.Context, body []byte) bool {
if strings.TrimSpace(c.GetHeader("x-gemini-api-privileged-user-id")) != "" {
return true
}
return geminiCLITmpDirRegex.Match(body)
}
// GeminiV1BetaListModels proxies: // GeminiV1BetaListModels proxies:
// GET /v1beta/models // GET /v1beta/models
func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) { func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
@@ -239,7 +233,14 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
sessionHash := extractGeminiCLISessionHash(c, body) sessionHash := extractGeminiCLISessionHash(c, body)
if sessionHash == "" { if sessionHash == "" {
// Fallback: 使用通用的会话哈希生成逻辑(适用于其他客户端) // Fallback: 使用通用的会话哈希生成逻辑(适用于其他客户端)
parsedReq, _ := service.ParseGatewayRequest(body) parsedReq, _ := service.ParseGatewayRequest(body, domain.PlatformGemini)
if parsedReq != nil {
parsedReq.SessionContext = &service.SessionContext{
ClientIP: ip.GetClientIP(c),
UserAgent: c.GetHeader("User-Agent"),
APIKeyID: apiKey.ID,
}
}
sessionHash = h.gatewayService.GenerateSessionHash(parsedReq) sessionHash = h.gatewayService.GenerateSessionHash(parsedReq)
} }
sessionKey := sessionHash sessionKey := sessionHash
@@ -258,6 +259,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
var geminiDigestChain string var geminiDigestChain string
var geminiPrefixHash string var geminiPrefixHash string
var geminiSessionUUID string var geminiSessionUUID string
var matchedDigestChain string
useDigestFallback := sessionBoundAccountID == 0 useDigestFallback := sessionBoundAccountID == 0
if useDigestFallback { if useDigestFallback {
@@ -284,13 +286,14 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
) )
// 查找会话 // 查找会话
foundUUID, foundAccountID, found := h.gatewayService.FindGeminiSession( foundUUID, foundAccountID, foundMatchedChain, found := h.gatewayService.FindGeminiSession(
c.Request.Context(), c.Request.Context(),
derefGroupID(apiKey.GroupID), derefGroupID(apiKey.GroupID),
geminiPrefixHash, geminiPrefixHash,
geminiDigestChain, geminiDigestChain,
) )
if found { if found {
matchedDigestChain = foundMatchedChain
sessionBoundAccountID = foundAccountID sessionBoundAccountID = foundAccountID
geminiSessionUUID = foundUUID geminiSessionUUID = foundUUID
log.Printf("[Gemini] Digest fallback matched: uuid=%s, accountID=%d, chain=%s", log.Printf("[Gemini] Digest fallback matched: uuid=%s, accountID=%d, chain=%s",
@@ -316,7 +319,6 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
// 判断是否真的绑定了粘性会话:有 sessionKey 且已经绑定到某个账号 // 判断是否真的绑定了粘性会话:有 sessionKey 且已经绑定到某个账号
hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0 hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0
isCLI := isGeminiCLIRequest(c, body)
cleanedForUnknownBinding := false cleanedForUnknownBinding := false
maxAccountSwitches := h.maxAccountSwitchesGemini maxAccountSwitches := h.maxAccountSwitchesGemini
@@ -344,10 +346,10 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
log.Printf("[Gemini] Sticky session account switched: %d -> %d, cleaning thoughtSignature", sessionBoundAccountID, account.ID) log.Printf("[Gemini] Sticky session account switched: %d -> %d, cleaning thoughtSignature", sessionBoundAccountID, account.ID)
body = service.CleanGeminiNativeThoughtSignatures(body) body = service.CleanGeminiNativeThoughtSignatures(body)
sessionBoundAccountID = account.ID sessionBoundAccountID = account.ID
} else if sessionKey != "" && sessionBoundAccountID == 0 && isCLI && !cleanedForUnknownBinding && bytes.Contains(body, []byte(`"thoughtSignature"`)) { } else if sessionKey != "" && sessionBoundAccountID == 0 && !cleanedForUnknownBinding && bytes.Contains(body, []byte(`"thoughtSignature"`)) {
// 无缓存绑定但请求里已有 thoughtSignature常见于缓存丢失/TTL 过期后,CLI 继续携带旧签名。 // 无缓存绑定但请求里已有 thoughtSignature常见于缓存丢失/TTL 过期后,客户端继续携带旧签名。
// 为避免第一次转发就 400这里做一次确定性清理让新账号重新生成签名链路。 // 为避免第一次转发就 400这里做一次确定性清理让新账号重新生成签名链路。
log.Printf("[Gemini] Sticky session binding missing for CLI request, cleaning thoughtSignature proactively") log.Printf("[Gemini] Sticky session binding missing, cleaning thoughtSignature proactively")
body = service.CleanGeminiNativeThoughtSignatures(body) body = service.CleanGeminiNativeThoughtSignatures(body)
cleanedForUnknownBinding = true cleanedForUnknownBinding = true
sessionBoundAccountID = account.ID sessionBoundAccountID = account.ID
@@ -410,7 +412,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
if switchCount > 0 { if switchCount > 0 {
requestCtx = context.WithValue(requestCtx, ctxkey.AccountSwitchCount, switchCount) requestCtx = context.WithValue(requestCtx, ctxkey.AccountSwitchCount, switchCount)
} }
if account.Platform == service.PlatformAntigravity { if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey {
result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, modelName, action, stream, body, hasBoundSession) result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, modelName, action, stream, body, hasBoundSession)
} else { } else {
result, err = h.geminiCompatService.ForwardNative(requestCtx, c, account, modelName, action, stream, body) result, err = h.geminiCompatService.ForwardNative(requestCtx, c, account, modelName, action, stream, body)
@@ -422,7 +424,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
var failoverErr *service.UpstreamFailoverError var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) { if errors.As(err, &failoverErr) {
failedAccountIDs[account.ID] = struct{}{} failedAccountIDs[account.ID] = struct{}{}
if failoverErr.ForceCacheBilling { if needForceCacheBilling(hasBoundSession, failoverErr) {
forceCacheBilling = true forceCacheBilling = true
} }
if switchCount >= maxAccountSwitches { if switchCount >= maxAccountSwitches {
@@ -433,6 +435,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
lastFailoverErr = failoverErr lastFailoverErr = failoverErr
switchCount++ switchCount++
log.Printf("Gemini account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches) log.Printf("Gemini account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
if account.Platform == service.PlatformAntigravity {
if !sleepFailoverDelay(c.Request.Context(), switchCount) {
return
}
}
continue continue
} }
// ForwardNative already wrote the response // ForwardNative already wrote the response
@@ -453,6 +460,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
geminiDigestChain, geminiDigestChain,
geminiSessionUUID, geminiSessionUUID,
account.ID, account.ID,
matchedDigestChain,
); err != nil { ); err != nil {
log.Printf("[Gemini] Failed to save digest session: %v", err) log.Printf("[Gemini] Failed to save digest session: %v", err)
} }

View File

@@ -282,6 +282,34 @@ func (r *accountRepository) GetByCRSAccountID(ctx context.Context, crsAccountID
return &accounts[0], nil return &accounts[0], nil
} }
func (r *accountRepository) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
rows, err := r.sql.QueryContext(ctx, `
SELECT id, extra->>'crs_account_id'
FROM accounts
WHERE deleted_at IS NULL
AND extra->>'crs_account_id' IS NOT NULL
AND extra->>'crs_account_id' != ''
`)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
result := make(map[string]int64)
for rows.Next() {
var id int64
var crsID string
if err := rows.Scan(&id, &crsID); err != nil {
return nil, err
}
result[crsID] = id
}
if err := rows.Err(); err != nil {
return nil, err
}
return result, nil
}
func (r *accountRepository) Update(ctx context.Context, account *service.Account) error { func (r *accountRepository) Update(ctx context.Context, account *service.Account) error {
if account == nil { if account == nil {
return nil return nil
@@ -798,53 +826,6 @@ func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetA
return nil return nil
} }
func (r *accountRepository) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope service.AntigravityQuotaScope, resetAt time.Time) error {
now := time.Now().UTC()
payload := map[string]string{
"rate_limited_at": now.Format(time.RFC3339),
"rate_limit_reset_at": resetAt.UTC().Format(time.RFC3339),
}
raw, err := json.Marshal(payload)
if err != nil {
return err
}
scopeKey := string(scope)
client := clientFromContext(ctx, r.client)
result, err := client.ExecContext(
ctx,
`UPDATE accounts SET
extra = jsonb_set(
jsonb_set(COALESCE(extra, '{}'::jsonb), '{antigravity_quota_scopes}'::text[], COALESCE(extra->'antigravity_quota_scopes', '{}'::jsonb), true),
ARRAY['antigravity_quota_scopes', $1]::text[],
$2::jsonb,
true
),
updated_at = NOW(),
last_used_at = NOW()
WHERE id = $3 AND deleted_at IS NULL`,
scopeKey,
raw,
id,
)
if err != nil {
return err
}
affected, err := result.RowsAffected()
if err != nil {
return err
}
if affected == 0 {
return service.ErrAccountNotFound
}
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
log.Printf("[SchedulerOutbox] enqueue quota scope failed: account=%d err=%v", id, err)
}
return nil
}
func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error { func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error {
if scope == "" { if scope == "" {
return nil return nil

View File

@@ -468,6 +468,7 @@ func groupEntityToService(g *dbent.Group) *service.Group {
ModelRoutingEnabled: g.ModelRoutingEnabled, ModelRoutingEnabled: g.ModelRoutingEnabled,
MCPXMLInject: g.McpXMLInject, MCPXMLInject: g.McpXMLInject,
SupportedModelScopes: g.SupportedModelScopes, SupportedModelScopes: g.SupportedModelScopes,
SortOrder: g.SortOrder,
CreatedAt: g.CreatedAt, CreatedAt: g.CreatedAt,
UpdatedAt: g.UpdatedAt, UpdatedAt: g.UpdatedAt,
} }

View File

@@ -11,63 +11,6 @@ import (
const stickySessionPrefix = "sticky_session:" const stickySessionPrefix = "sticky_session:"
// Gemini Trie Lua 脚本
const (
// geminiTrieFindScript 查找最长前缀匹配的 Lua 脚本
// KEYS[1] = trie key
// ARGV[1] = digestChain (如 "u:a-m:b-u:c-m:d")
// ARGV[2] = TTL seconds (用于刷新)
// 返回: 最长匹配的 value (uuid:accountID) 或 nil
// 查找成功时自动刷新 TTL防止活跃会话意外过期
geminiTrieFindScript = `
local chain = ARGV[1]
local ttl = tonumber(ARGV[2])
local lastMatch = nil
local path = ""
for part in string.gmatch(chain, "[^-]+") do
path = path == "" and part or path .. "-" .. part
local val = redis.call('HGET', KEYS[1], path)
if val and val ~= "" then
lastMatch = val
end
end
if lastMatch then
redis.call('EXPIRE', KEYS[1], ttl)
end
return lastMatch
`
// geminiTrieSaveScript 保存会话到 Trie 的 Lua 脚本
// KEYS[1] = trie key
// ARGV[1] = digestChain
// ARGV[2] = value (uuid:accountID)
// ARGV[3] = TTL seconds
geminiTrieSaveScript = `
local chain = ARGV[1]
local value = ARGV[2]
local ttl = tonumber(ARGV[3])
local path = ""
for part in string.gmatch(chain, "[^-]+") do
path = path == "" and part or path .. "-" .. part
end
redis.call('HSET', KEYS[1], path, value)
redis.call('EXPIRE', KEYS[1], ttl)
return "OK"
`
)
// 模型负载统计相关常量
const (
modelLoadKeyPrefix = "ag:model_load:" // 模型调用次数 key 前缀
modelLastUsedKeyPrefix = "ag:model_last_used:" // 模型最后调度时间 key 前缀
modelLoadTTL = 24 * time.Hour // 调用次数 TTL24 小时无调用后清零)
modelLastUsedTTL = 24 * time.Hour // 最后调度时间 TTL
)
type gatewayCache struct { type gatewayCache struct {
rdb *redis.Client rdb *redis.Client
} }
@@ -108,133 +51,3 @@ func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64
key := buildSessionKey(groupID, sessionHash) key := buildSessionKey(groupID, sessionHash)
return c.rdb.Del(ctx, key).Err() return c.rdb.Del(ctx, key).Err()
} }
// ============ Antigravity 模型负载统计方法 ============
// modelLoadKey 构建模型调用次数 key
// 格式: ag:model_load:{accountID}:{model}
func modelLoadKey(accountID int64, model string) string {
return fmt.Sprintf("%s%d:%s", modelLoadKeyPrefix, accountID, model)
}
// modelLastUsedKey 构建模型最后调度时间 key
// 格式: ag:model_last_used:{accountID}:{model}
func modelLastUsedKey(accountID int64, model string) string {
return fmt.Sprintf("%s%d:%s", modelLastUsedKeyPrefix, accountID, model)
}
// IncrModelCallCount 增加模型调用次数并更新最后调度时间
// 返回更新后的调用次数
func (c *gatewayCache) IncrModelCallCount(ctx context.Context, accountID int64, model string) (int64, error) {
loadKey := modelLoadKey(accountID, model)
lastUsedKey := modelLastUsedKey(accountID, model)
pipe := c.rdb.Pipeline()
incrCmd := pipe.Incr(ctx, loadKey)
pipe.Expire(ctx, loadKey, modelLoadTTL) // 每次调用刷新 TTL
pipe.Set(ctx, lastUsedKey, time.Now().Unix(), modelLastUsedTTL)
if _, err := pipe.Exec(ctx); err != nil {
return 0, err
}
return incrCmd.Val(), nil
}
// GetModelLoadBatch 批量获取账号的模型负载信息
func (c *gatewayCache) GetModelLoadBatch(ctx context.Context, accountIDs []int64, model string) (map[int64]*service.ModelLoadInfo, error) {
if len(accountIDs) == 0 {
return make(map[int64]*service.ModelLoadInfo), nil
}
loadCmds, lastUsedCmds := c.pipelineModelLoadGet(ctx, accountIDs, model)
return c.parseModelLoadResults(accountIDs, loadCmds, lastUsedCmds), nil
}
// pipelineModelLoadGet 批量获取模型负载的 Pipeline 操作
func (c *gatewayCache) pipelineModelLoadGet(
ctx context.Context,
accountIDs []int64,
model string,
) (map[int64]*redis.StringCmd, map[int64]*redis.StringCmd) {
pipe := c.rdb.Pipeline()
loadCmds := make(map[int64]*redis.StringCmd, len(accountIDs))
lastUsedCmds := make(map[int64]*redis.StringCmd, len(accountIDs))
for _, id := range accountIDs {
loadCmds[id] = pipe.Get(ctx, modelLoadKey(id, model))
lastUsedCmds[id] = pipe.Get(ctx, modelLastUsedKey(id, model))
}
_, _ = pipe.Exec(ctx) // 忽略错误key 不存在是正常的
return loadCmds, lastUsedCmds
}
// parseModelLoadResults 解析 Pipeline 结果
func (c *gatewayCache) parseModelLoadResults(
accountIDs []int64,
loadCmds map[int64]*redis.StringCmd,
lastUsedCmds map[int64]*redis.StringCmd,
) map[int64]*service.ModelLoadInfo {
result := make(map[int64]*service.ModelLoadInfo, len(accountIDs))
for _, id := range accountIDs {
result[id] = &service.ModelLoadInfo{
CallCount: getInt64OrZero(loadCmds[id]),
LastUsedAt: getTimeOrZero(lastUsedCmds[id]),
}
}
return result
}
// getInt64OrZero 从 StringCmd 获取 int64 值,失败返回 0
func getInt64OrZero(cmd *redis.StringCmd) int64 {
val, _ := cmd.Int64()
return val
}
// getTimeOrZero 从 StringCmd 获取 time.Time失败返回零值
func getTimeOrZero(cmd *redis.StringCmd) time.Time {
val, err := cmd.Int64()
if err != nil {
return time.Time{}
}
return time.Unix(val, 0)
}
// ============ Gemini 会话 Fallback 方法 (Trie 实现) ============
// FindGeminiSession 查找 Gemini 会话(使用 Trie + Lua 脚本实现 O(L) 查询)
// 返回最长匹配的会话信息,匹配成功时自动刷新 TTL
func (c *gatewayCache) FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
if digestChain == "" {
return "", 0, false
}
trieKey := service.BuildGeminiTrieKey(groupID, prefixHash)
ttlSeconds := int(service.GeminiSessionTTL().Seconds())
// 使用 Lua 脚本在 Redis 端执行 Trie 查找O(L) 次 HGET1 次网络往返
// 查找成功时自动刷新 TTL防止活跃会话意外过期
result, err := c.rdb.Eval(ctx, geminiTrieFindScript, []string{trieKey}, digestChain, ttlSeconds).Result()
if err != nil || result == nil {
return "", 0, false
}
value, ok := result.(string)
if !ok || value == "" {
return "", 0, false
}
uuid, accountID, ok = service.ParseGeminiSessionValue(value)
return uuid, accountID, ok
}
// SaveGeminiSession 保存 Gemini 会话(使用 Trie + Lua 脚本)
func (c *gatewayCache) SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
if digestChain == "" {
return nil
}
trieKey := service.BuildGeminiTrieKey(groupID, prefixHash)
value := service.FormatGeminiSessionValue(uuid, accountID)
ttlSeconds := int(service.GeminiSessionTTL().Seconds())
return c.rdb.Eval(ctx, geminiTrieSaveScript, []string{trieKey}, digestChain, value, ttlSeconds).Err()
}

View File

@@ -104,157 +104,6 @@ func (s *GatewayCacheSuite) TestGetSessionAccountID_CorruptedValue() {
require.False(s.T(), errors.Is(err, redis.Nil), "expected parsing error, not redis.Nil") require.False(s.T(), errors.Is(err, redis.Nil), "expected parsing error, not redis.Nil")
} }
// ============ Gemini Trie 会话测试 ============
func (s *GatewayCacheSuite) TestGeminiSessionTrie_SaveAndFind() {
groupID := int64(1)
prefixHash := "testprefix"
digestChain := "u:hash1-m:hash2-u:hash3"
uuid := "test-uuid-123"
accountID := int64(42)
// 保存会话
err := s.cache.SaveGeminiSession(s.ctx, groupID, prefixHash, digestChain, uuid, accountID)
require.NoError(s.T(), err, "SaveGeminiSession")
// 精确匹配查找
foundUUID, foundAccountID, found := s.cache.FindGeminiSession(s.ctx, groupID, prefixHash, digestChain)
require.True(s.T(), found, "should find exact match")
require.Equal(s.T(), uuid, foundUUID)
require.Equal(s.T(), accountID, foundAccountID)
}
func (s *GatewayCacheSuite) TestGeminiSessionTrie_PrefixMatch() {
groupID := int64(1)
prefixHash := "prefixmatch"
shortChain := "u:a-m:b"
longChain := "u:a-m:b-u:c-m:d"
uuid := "uuid-prefix"
accountID := int64(100)
// 保存短链
err := s.cache.SaveGeminiSession(s.ctx, groupID, prefixHash, shortChain, uuid, accountID)
require.NoError(s.T(), err)
// 用长链查找,应该匹配到短链(前缀匹配)
foundUUID, foundAccountID, found := s.cache.FindGeminiSession(s.ctx, groupID, prefixHash, longChain)
require.True(s.T(), found, "should find prefix match")
require.Equal(s.T(), uuid, foundUUID)
require.Equal(s.T(), accountID, foundAccountID)
}
func (s *GatewayCacheSuite) TestGeminiSessionTrie_LongestPrefixMatch() {
groupID := int64(1)
prefixHash := "longestmatch"
// 保存多个不同长度的链
err := s.cache.SaveGeminiSession(s.ctx, groupID, prefixHash, "u:a", "uuid-short", 1)
require.NoError(s.T(), err)
err = s.cache.SaveGeminiSession(s.ctx, groupID, prefixHash, "u:a-m:b", "uuid-medium", 2)
require.NoError(s.T(), err)
err = s.cache.SaveGeminiSession(s.ctx, groupID, prefixHash, "u:a-m:b-u:c", "uuid-long", 3)
require.NoError(s.T(), err)
// 查找更长的链,应该匹配到最长的前缀
foundUUID, foundAccountID, found := s.cache.FindGeminiSession(s.ctx, groupID, prefixHash, "u:a-m:b-u:c-m:d-u:e")
require.True(s.T(), found, "should find longest prefix match")
require.Equal(s.T(), "uuid-long", foundUUID)
require.Equal(s.T(), int64(3), foundAccountID)
// 查找中等长度的链
foundUUID, foundAccountID, found = s.cache.FindGeminiSession(s.ctx, groupID, prefixHash, "u:a-m:b-u:x")
require.True(s.T(), found)
require.Equal(s.T(), "uuid-medium", foundUUID)
require.Equal(s.T(), int64(2), foundAccountID)
}
func (s *GatewayCacheSuite) TestGeminiSessionTrie_NoMatch() {
groupID := int64(1)
prefixHash := "nomatch"
digestChain := "u:a-m:b"
// 保存一个会话
err := s.cache.SaveGeminiSession(s.ctx, groupID, prefixHash, digestChain, "uuid", 1)
require.NoError(s.T(), err)
// 用不同的链查找,应该找不到
_, _, found := s.cache.FindGeminiSession(s.ctx, groupID, prefixHash, "u:x-m:y")
require.False(s.T(), found, "should not find non-matching chain")
}
func (s *GatewayCacheSuite) TestGeminiSessionTrie_DifferentPrefixHash() {
groupID := int64(1)
digestChain := "u:a-m:b"
// 保存到 prefixHash1
err := s.cache.SaveGeminiSession(s.ctx, groupID, "prefix1", digestChain, "uuid1", 1)
require.NoError(s.T(), err)
// 用 prefixHash2 查找,应该找不到(不同用户/客户端隔离)
_, _, found := s.cache.FindGeminiSession(s.ctx, groupID, "prefix2", digestChain)
require.False(s.T(), found, "different prefixHash should be isolated")
}
func (s *GatewayCacheSuite) TestGeminiSessionTrie_DifferentGroupID() {
prefixHash := "sameprefix"
digestChain := "u:a-m:b"
// 保存到 groupID 1
err := s.cache.SaveGeminiSession(s.ctx, 1, prefixHash, digestChain, "uuid1", 1)
require.NoError(s.T(), err)
// 用 groupID 2 查找,应该找不到(分组隔离)
_, _, found := s.cache.FindGeminiSession(s.ctx, 2, prefixHash, digestChain)
require.False(s.T(), found, "different groupID should be isolated")
}
func (s *GatewayCacheSuite) TestGeminiSessionTrie_EmptyDigestChain() {
groupID := int64(1)
prefixHash := "emptytest"
// 空链不应该保存
err := s.cache.SaveGeminiSession(s.ctx, groupID, prefixHash, "", "uuid", 1)
require.NoError(s.T(), err, "empty chain should not error")
// 空链查找应该返回 false
_, _, found := s.cache.FindGeminiSession(s.ctx, groupID, prefixHash, "")
require.False(s.T(), found, "empty chain should not match")
}
func (s *GatewayCacheSuite) TestGeminiSessionTrie_MultipleSessions() {
groupID := int64(1)
prefixHash := "multisession"
// 保存多个不同会话(模拟 1000 个并发会话的场景)
sessions := []struct {
chain string
uuid string
accountID int64
}{
{"u:session1", "uuid-1", 1},
{"u:session2-m:reply2", "uuid-2", 2},
{"u:session3-m:reply3-u:msg3", "uuid-3", 3},
}
for _, sess := range sessions {
err := s.cache.SaveGeminiSession(s.ctx, groupID, prefixHash, sess.chain, sess.uuid, sess.accountID)
require.NoError(s.T(), err)
}
// 验证每个会话都能正确查找
for _, sess := range sessions {
foundUUID, foundAccountID, found := s.cache.FindGeminiSession(s.ctx, groupID, prefixHash, sess.chain)
require.True(s.T(), found, "should find session: %s", sess.chain)
require.Equal(s.T(), sess.uuid, foundUUID)
require.Equal(s.T(), sess.accountID, foundAccountID)
}
// 验证继续对话的场景
foundUUID, foundAccountID, found := s.cache.FindGeminiSession(s.ctx, groupID, prefixHash, "u:session2-m:reply2-u:newmsg")
require.True(s.T(), found)
require.Equal(s.T(), "uuid-2", foundUUID)
require.Equal(s.T(), int64(2), foundAccountID)
}
func TestGatewayCacheSuite(t *testing.T) { func TestGatewayCacheSuite(t *testing.T) {
suite.Run(t, new(GatewayCacheSuite)) suite.Run(t, new(GatewayCacheSuite))

View File

@@ -1,234 +0,0 @@
//go:build integration
package repository
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
)
// ============ Gateway Cache 模型负载统计集成测试 ============
type GatewayCacheModelLoadSuite struct {
suite.Suite
}
func TestGatewayCacheModelLoadSuite(t *testing.T) {
suite.Run(t, new(GatewayCacheModelLoadSuite))
}
func (s *GatewayCacheModelLoadSuite) TestIncrModelCallCount_Basic() {
t := s.T()
rdb := testRedis(t)
cache := &gatewayCache{rdb: rdb}
ctx := context.Background()
accountID := int64(123)
model := "claude-sonnet-4-20250514"
// 首次调用应返回 1
count1, err := cache.IncrModelCallCount(ctx, accountID, model)
require.NoError(t, err)
require.Equal(t, int64(1), count1)
// 第二次调用应返回 2
count2, err := cache.IncrModelCallCount(ctx, accountID, model)
require.NoError(t, err)
require.Equal(t, int64(2), count2)
// 第三次调用应返回 3
count3, err := cache.IncrModelCallCount(ctx, accountID, model)
require.NoError(t, err)
require.Equal(t, int64(3), count3)
}
func (s *GatewayCacheModelLoadSuite) TestIncrModelCallCount_DifferentModels() {
t := s.T()
rdb := testRedis(t)
cache := &gatewayCache{rdb: rdb}
ctx := context.Background()
accountID := int64(456)
model1 := "claude-sonnet-4-20250514"
model2 := "claude-opus-4-5-20251101"
// 不同模型应该独立计数
count1, err := cache.IncrModelCallCount(ctx, accountID, model1)
require.NoError(t, err)
require.Equal(t, int64(1), count1)
count2, err := cache.IncrModelCallCount(ctx, accountID, model2)
require.NoError(t, err)
require.Equal(t, int64(1), count2)
count1Again, err := cache.IncrModelCallCount(ctx, accountID, model1)
require.NoError(t, err)
require.Equal(t, int64(2), count1Again)
}
func (s *GatewayCacheModelLoadSuite) TestIncrModelCallCount_DifferentAccounts() {
t := s.T()
rdb := testRedis(t)
cache := &gatewayCache{rdb: rdb}
ctx := context.Background()
account1 := int64(111)
account2 := int64(222)
model := "gemini-2.5-pro"
// 不同账号应该独立计数
count1, err := cache.IncrModelCallCount(ctx, account1, model)
require.NoError(t, err)
require.Equal(t, int64(1), count1)
count2, err := cache.IncrModelCallCount(ctx, account2, model)
require.NoError(t, err)
require.Equal(t, int64(1), count2)
}
func (s *GatewayCacheModelLoadSuite) TestGetModelLoadBatch_Empty() {
t := s.T()
rdb := testRedis(t)
cache := &gatewayCache{rdb: rdb}
ctx := context.Background()
result, err := cache.GetModelLoadBatch(ctx, []int64{}, "any-model")
require.NoError(t, err)
require.NotNil(t, result)
require.Empty(t, result)
}
func (s *GatewayCacheModelLoadSuite) TestGetModelLoadBatch_NonExistent() {
t := s.T()
rdb := testRedis(t)
cache := &gatewayCache{rdb: rdb}
ctx := context.Background()
// 查询不存在的账号应返回零值
result, err := cache.GetModelLoadBatch(ctx, []int64{9999, 9998}, "claude-sonnet-4-20250514")
require.NoError(t, err)
require.Len(t, result, 2)
require.Equal(t, int64(0), result[9999].CallCount)
require.True(t, result[9999].LastUsedAt.IsZero())
require.Equal(t, int64(0), result[9998].CallCount)
require.True(t, result[9998].LastUsedAt.IsZero())
}
func (s *GatewayCacheModelLoadSuite) TestGetModelLoadBatch_AfterIncrement() {
t := s.T()
rdb := testRedis(t)
cache := &gatewayCache{rdb: rdb}
ctx := context.Background()
accountID := int64(789)
model := "claude-sonnet-4-20250514"
// 先增加调用次数
beforeIncr := time.Now()
_, err := cache.IncrModelCallCount(ctx, accountID, model)
require.NoError(t, err)
_, err = cache.IncrModelCallCount(ctx, accountID, model)
require.NoError(t, err)
_, err = cache.IncrModelCallCount(ctx, accountID, model)
require.NoError(t, err)
afterIncr := time.Now()
// 获取负载信息
result, err := cache.GetModelLoadBatch(ctx, []int64{accountID}, model)
require.NoError(t, err)
require.Len(t, result, 1)
loadInfo := result[accountID]
require.NotNil(t, loadInfo)
require.Equal(t, int64(3), loadInfo.CallCount)
require.False(t, loadInfo.LastUsedAt.IsZero())
// LastUsedAt 应该在 beforeIncr 和 afterIncr 之间
require.True(t, loadInfo.LastUsedAt.After(beforeIncr.Add(-time.Second)) || loadInfo.LastUsedAt.Equal(beforeIncr))
require.True(t, loadInfo.LastUsedAt.Before(afterIncr.Add(time.Second)) || loadInfo.LastUsedAt.Equal(afterIncr))
}
func (s *GatewayCacheModelLoadSuite) TestGetModelLoadBatch_MultipleAccounts() {
t := s.T()
rdb := testRedis(t)
cache := &gatewayCache{rdb: rdb}
ctx := context.Background()
model := "claude-opus-4-5-20251101"
account1 := int64(1001)
account2 := int64(1002)
account3 := int64(1003) // 不调用
// account1 调用 2 次
_, err := cache.IncrModelCallCount(ctx, account1, model)
require.NoError(t, err)
_, err = cache.IncrModelCallCount(ctx, account1, model)
require.NoError(t, err)
// account2 调用 5 次
for i := 0; i < 5; i++ {
_, err = cache.IncrModelCallCount(ctx, account2, model)
require.NoError(t, err)
}
// 批量获取
result, err := cache.GetModelLoadBatch(ctx, []int64{account1, account2, account3}, model)
require.NoError(t, err)
require.Len(t, result, 3)
require.Equal(t, int64(2), result[account1].CallCount)
require.False(t, result[account1].LastUsedAt.IsZero())
require.Equal(t, int64(5), result[account2].CallCount)
require.False(t, result[account2].LastUsedAt.IsZero())
require.Equal(t, int64(0), result[account3].CallCount)
require.True(t, result[account3].LastUsedAt.IsZero())
}
func (s *GatewayCacheModelLoadSuite) TestGetModelLoadBatch_ModelIsolation() {
t := s.T()
rdb := testRedis(t)
cache := &gatewayCache{rdb: rdb}
ctx := context.Background()
accountID := int64(2001)
model1 := "claude-sonnet-4-20250514"
model2 := "gemini-2.5-pro"
// 对 model1 调用 3 次
for i := 0; i < 3; i++ {
_, err := cache.IncrModelCallCount(ctx, accountID, model1)
require.NoError(t, err)
}
// 获取 model1 的负载
result1, err := cache.GetModelLoadBatch(ctx, []int64{accountID}, model1)
require.NoError(t, err)
require.Equal(t, int64(3), result1[accountID].CallCount)
// 获取 model2 的负载(应该为 0
result2, err := cache.GetModelLoadBatch(ctx, []int64{accountID}, model2)
require.NoError(t, err)
require.Equal(t, int64(0), result2[accountID].CallCount)
}
// ============ 辅助函数测试 ============
func (s *GatewayCacheModelLoadSuite) TestModelLoadKey_Format() {
t := s.T()
key := modelLoadKey(123, "claude-sonnet-4")
require.Equal(t, "ag:model_load:123:claude-sonnet-4", key)
}
func (s *GatewayCacheModelLoadSuite) TestModelLastUsedKey_Format() {
t := s.T()
key := modelLastUsedKey(456, "gemini-2.5-pro")
require.Equal(t, "ag:model_last_used:456:gemini-2.5-pro", key)
}

View File

@@ -191,7 +191,7 @@ func (r *groupRepository) ListWithFilters(ctx context.Context, params pagination
groups, err := q. groups, err := q.
Offset(params.Offset()). Offset(params.Offset()).
Limit(params.Limit()). Limit(params.Limit()).
Order(dbent.Asc(group.FieldID)). Order(dbent.Asc(group.FieldSortOrder), dbent.Asc(group.FieldID)).
All(ctx) All(ctx)
if err != nil { if err != nil {
return nil, nil, err return nil, nil, err
@@ -218,7 +218,7 @@ func (r *groupRepository) ListWithFilters(ctx context.Context, params pagination
func (r *groupRepository) ListActive(ctx context.Context) ([]service.Group, error) { func (r *groupRepository) ListActive(ctx context.Context) ([]service.Group, error) {
groups, err := r.client.Group.Query(). groups, err := r.client.Group.Query().
Where(group.StatusEQ(service.StatusActive)). Where(group.StatusEQ(service.StatusActive)).
Order(dbent.Asc(group.FieldID)). Order(dbent.Asc(group.FieldSortOrder), dbent.Asc(group.FieldID)).
All(ctx) All(ctx)
if err != nil { if err != nil {
return nil, err return nil, err
@@ -245,7 +245,7 @@ func (r *groupRepository) ListActive(ctx context.Context) ([]service.Group, erro
func (r *groupRepository) ListActiveByPlatform(ctx context.Context, platform string) ([]service.Group, error) { func (r *groupRepository) ListActiveByPlatform(ctx context.Context, platform string) ([]service.Group, error) {
groups, err := r.client.Group.Query(). groups, err := r.client.Group.Query().
Where(group.StatusEQ(service.StatusActive), group.PlatformEQ(platform)). Where(group.StatusEQ(service.StatusActive), group.PlatformEQ(platform)).
Order(dbent.Asc(group.FieldID)). Order(dbent.Asc(group.FieldSortOrder), dbent.Asc(group.FieldID)).
All(ctx) All(ctx)
if err != nil { if err != nil {
return nil, err return nil, err
@@ -497,3 +497,29 @@ func (r *groupRepository) BindAccountsToGroup(ctx context.Context, groupID int64
return nil return nil
} }
// UpdateSortOrders 批量更新分组排序
func (r *groupRepository) UpdateSortOrders(ctx context.Context, updates []service.GroupSortOrderUpdate) error {
if len(updates) == 0 {
return nil
}
// 使用事务批量更新
tx, err := r.client.Tx(ctx)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
for _, u := range updates {
if _, err := tx.Group.UpdateOneID(u.ID).SetSortOrder(u.SortOrder).Save(ctx); err != nil {
return translatePersistenceError(err, service.ErrGroupNotFound, nil)
}
}
if err := tx.Commit(); err != nil {
return err
}
return nil
}

View File

@@ -896,6 +896,10 @@ func (stubGroupRepo) GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []int
return nil, errors.New("not implemented") return nil, errors.New("not implemented")
} }
func (stubGroupRepo) UpdateSortOrders(ctx context.Context, updates []service.GroupSortOrderUpdate) error {
return nil
}
type stubAccountRepo struct { type stubAccountRepo struct {
bulkUpdateIDs []int64 bulkUpdateIDs []int64
} }
@@ -1004,10 +1008,6 @@ func (s *stubAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt
return errors.New("not implemented") return errors.New("not implemented")
} }
func (s *stubAccountRepo) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope service.AntigravityQuotaScope, resetAt time.Time) error {
return errors.New("not implemented")
}
func (s *stubAccountRepo) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error { func (s *stubAccountRepo) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error {
return errors.New("not implemented") return errors.New("not implemented")
} }
@@ -1049,6 +1049,10 @@ func (s *stubAccountRepo) BulkUpdate(ctx context.Context, ids []int64, updates s
return int64(len(ids)), nil return int64(len(ids)), nil
} }
func (s *stubAccountRepo) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
return nil, errors.New("not implemented")
}
type stubProxyRepo struct{} type stubProxyRepo struct{}
func (stubProxyRepo) Create(ctx context.Context, proxy *service.Proxy) error { func (stubProxyRepo) Create(ctx context.Context, proxy *service.Proxy) error {

View File

@@ -192,6 +192,7 @@ func registerGroupRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
{ {
groups.GET("", h.Admin.Group.List) groups.GET("", h.Admin.Group.List)
groups.GET("/all", h.Admin.Group.GetAll) groups.GET("/all", h.Admin.Group.GetAll)
groups.PUT("/sort-order", h.Admin.Group.UpdateSortOrder)
groups.GET("/:id", h.Admin.Group.GetByID) groups.GET("/:id", h.Admin.Group.GetByID)
groups.POST("", h.Admin.Group.Create) groups.POST("", h.Admin.Group.Create)
groups.PUT("/:id", h.Admin.Group.Update) groups.PUT("/:id", h.Admin.Group.Update)
@@ -208,6 +209,7 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
accounts.GET("/:id", h.Admin.Account.GetByID) accounts.GET("/:id", h.Admin.Account.GetByID)
accounts.POST("", h.Admin.Account.Create) accounts.POST("", h.Admin.Account.Create)
accounts.POST("/sync/crs", h.Admin.Account.SyncFromCRS) accounts.POST("/sync/crs", h.Admin.Account.SyncFromCRS)
accounts.POST("/sync/crs/preview", h.Admin.Account.PreviewFromCRS)
accounts.PUT("/:id", h.Admin.Account.Update) accounts.PUT("/:id", h.Admin.Account.Update)
accounts.DELETE("/:id", h.Admin.Account.Delete) accounts.DELETE("/:id", h.Admin.Account.Delete)
accounts.POST("/:id/test", h.Admin.Account.Test) accounts.POST("/:id/test", h.Admin.Account.Test)

View File

@@ -425,6 +425,22 @@ func (a *Account) GetBaseURL() string {
if baseURL == "" { if baseURL == "" {
return "https://api.anthropic.com" return "https://api.anthropic.com"
} }
if a.Platform == PlatformAntigravity {
return strings.TrimRight(baseURL, "/") + "/antigravity"
}
return baseURL
}
// GetGeminiBaseURL 返回 Gemini 兼容端点的 base URL。
// Antigravity 平台的 APIKey 账号自动拼接 /antigravity。
func (a *Account) GetGeminiBaseURL(defaultBaseURL string) string {
baseURL := strings.TrimSpace(a.GetCredential("base_url"))
if baseURL == "" {
return defaultBaseURL
}
if a.Platform == PlatformAntigravity && a.Type == AccountTypeAPIKey {
return strings.TrimRight(baseURL, "/") + "/antigravity"
}
return baseURL return baseURL
} }

View File

@@ -0,0 +1,160 @@
//go:build unit
package service
import (
"testing"
)
func TestGetBaseURL(t *testing.T) {
tests := []struct {
name string
account Account
expected string
}{
{
name: "non-apikey type returns empty",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformAnthropic,
},
expected: "",
},
{
name: "apikey without base_url returns default anthropic",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAnthropic,
Credentials: map[string]any{},
},
expected: "https://api.anthropic.com",
},
{
name: "apikey with custom base_url",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAnthropic,
Credentials: map[string]any{"base_url": "https://custom.example.com"},
},
expected: "https://custom.example.com",
},
{
name: "antigravity apikey auto-appends /antigravity",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
},
expected: "https://upstream.example.com/antigravity",
},
{
name: "antigravity apikey trims trailing slash before appending",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com/"},
},
expected: "https://upstream.example.com/antigravity",
},
{
name: "antigravity non-apikey returns empty",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
},
expected: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.account.GetBaseURL()
if result != tt.expected {
t.Errorf("GetBaseURL() = %q, want %q", result, tt.expected)
}
})
}
}
func TestGetGeminiBaseURL(t *testing.T) {
const defaultGeminiURL = "https://generativelanguage.googleapis.com"
tests := []struct {
name string
account Account
expected string
}{
{
name: "apikey without base_url returns default",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{},
},
expected: defaultGeminiURL,
},
{
name: "apikey with custom base_url",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{"base_url": "https://custom-gemini.example.com"},
},
expected: "https://custom-gemini.example.com",
},
{
name: "antigravity apikey auto-appends /antigravity",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
},
expected: "https://upstream.example.com/antigravity",
},
{
name: "antigravity apikey trims trailing slash",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com/"},
},
expected: "https://upstream.example.com/antigravity",
},
{
name: "antigravity oauth does NOT append /antigravity",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
},
expected: "https://upstream.example.com",
},
{
name: "oauth without base_url returns default",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{},
},
expected: defaultGeminiURL,
},
{
name: "nil credentials returns default",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
},
expected: defaultGeminiURL,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.account.GetGeminiBaseURL(defaultGeminiURL)
if result != tt.expected {
t.Errorf("GetGeminiBaseURL() = %q, want %q", result, tt.expected)
}
})
}
}

View File

@@ -25,6 +25,9 @@ type AccountRepository interface {
// GetByCRSAccountID finds an account previously synced from CRS. // GetByCRSAccountID finds an account previously synced from CRS.
// Returns (nil, nil) if not found. // Returns (nil, nil) if not found.
GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error)
// ListCRSAccountIDs returns a map of crs_account_id -> local account ID
// for all accounts that have been synced from CRS.
ListCRSAccountIDs(ctx context.Context) (map[string]int64, error)
Update(ctx context.Context, account *Account) error Update(ctx context.Context, account *Account) error
Delete(ctx context.Context, id int64) error Delete(ctx context.Context, id int64) error
@@ -50,7 +53,6 @@ type AccountRepository interface {
ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error)
SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error
SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error
SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error
SetOverloaded(ctx context.Context, id int64, until time.Time) error SetOverloaded(ctx context.Context, id int64, until time.Time) error
SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error

View File

@@ -54,6 +54,10 @@ func (s *accountRepoStub) GetByCRSAccountID(ctx context.Context, crsAccountID st
panic("unexpected GetByCRSAccountID call") panic("unexpected GetByCRSAccountID call")
} }
func (s *accountRepoStub) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
panic("unexpected ListCRSAccountIDs call")
}
func (s *accountRepoStub) Update(ctx context.Context, account *Account) error { func (s *accountRepoStub) Update(ctx context.Context, account *Account) error {
panic("unexpected Update call") panic("unexpected Update call")
} }
@@ -143,10 +147,6 @@ func (s *accountRepoStub) SetRateLimited(ctx context.Context, id int64, resetAt
panic("unexpected SetRateLimited call") panic("unexpected SetRateLimited call")
} }
func (s *accountRepoStub) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error {
panic("unexpected SetAntigravityQuotaScopeLimit call")
}
func (s *accountRepoStub) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error { func (s *accountRepoStub) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error {
panic("unexpected SetModelRateLimit call") panic("unexpected SetModelRateLimit call")
} }

View File

@@ -245,7 +245,6 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
// Set common headers // Set common headers
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
req.Header.Set("anthropic-version", "2023-06-01") req.Header.Set("anthropic-version", "2023-06-01")
req.Header.Set("anthropic-beta", claude.DefaultBetaHeader)
// Apply Claude Code client headers // Apply Claude Code client headers
for key, value := range claude.DefaultHeaders { for key, value := range claude.DefaultHeaders {
@@ -254,8 +253,10 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
// Set authentication header // Set authentication header
if useBearer { if useBearer {
req.Header.Set("anthropic-beta", claude.DefaultBetaHeader)
req.Header.Set("Authorization", "Bearer "+authToken) req.Header.Set("Authorization", "Bearer "+authToken)
} else { } else {
req.Header.Set("anthropic-beta", claude.APIKeyBetaHeader)
req.Header.Set("x-api-key", authToken) req.Header.Set("x-api-key", authToken)
} }

View File

@@ -36,6 +36,7 @@ type AdminService interface {
UpdateGroup(ctx context.Context, id int64, input *UpdateGroupInput) (*Group, error) UpdateGroup(ctx context.Context, id int64, input *UpdateGroupInput) (*Group, error)
DeleteGroup(ctx context.Context, id int64) error DeleteGroup(ctx context.Context, id int64) error
GetGroupAPIKeys(ctx context.Context, groupID int64, page, pageSize int) ([]APIKey, int64, error) GetGroupAPIKeys(ctx context.Context, groupID int64, page, pageSize int) ([]APIKey, int64, error)
UpdateGroupSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error
// Account management // Account management
ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string) ([]Account, int64, error) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string) ([]Account, int64, error)
@@ -1015,6 +1016,10 @@ func (s *adminServiceImpl) GetGroupAPIKeys(ctx context.Context, groupID int64, p
return keys, result.Total, nil return keys, result.Total, nil
} }
func (s *adminServiceImpl) UpdateGroupSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
return s.groupRepo.UpdateSortOrders(ctx, updates)
}
// Account management implementations // Account management implementations
func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string) ([]Account, int64, error) { func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string) ([]Account, int64, error) {
params := pagination.PaginationParams{Page: page, PageSize: pageSize} params := pagination.PaginationParams{Page: page, PageSize: pageSize}

View File

@@ -172,6 +172,10 @@ func (s *groupRepoStub) GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []
panic("unexpected GetAccountIDsByGroupIDs call") panic("unexpected GetAccountIDsByGroupIDs call")
} }
func (s *groupRepoStub) UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
return nil
}
type proxyRepoStub struct { type proxyRepoStub struct {
deleteErr error deleteErr error
countErr error countErr error

View File

@@ -116,6 +116,10 @@ func (s *groupRepoStubForAdmin) GetAccountIDsByGroupIDs(_ context.Context, _ []i
panic("unexpected GetAccountIDsByGroupIDs call") panic("unexpected GetAccountIDsByGroupIDs call")
} }
func (s *groupRepoStubForAdmin) UpdateSortOrders(_ context.Context, _ []GroupSortOrderUpdate) error {
return nil
}
// TestAdminService_CreateGroup_WithImagePricing 测试创建分组时 ImagePrice 字段正确传递 // TestAdminService_CreateGroup_WithImagePricing 测试创建分组时 ImagePrice 字段正确传递
func TestAdminService_CreateGroup_WithImagePricing(t *testing.T) { func TestAdminService_CreateGroup_WithImagePricing(t *testing.T) {
repo := &groupRepoStubForAdmin{} repo := &groupRepoStubForAdmin{}
@@ -395,6 +399,10 @@ func (s *groupRepoStubForFallbackCycle) GetAccountIDsByGroupIDs(_ context.Contex
panic("unexpected GetAccountIDsByGroupIDs call") panic("unexpected GetAccountIDsByGroupIDs call")
} }
func (s *groupRepoStubForFallbackCycle) UpdateSortOrders(_ context.Context, _ []GroupSortOrderUpdate) error {
return nil
}
type groupRepoStubForInvalidRequestFallback struct { type groupRepoStubForInvalidRequestFallback struct {
groups map[int64]*Group groups map[int64]*Group
created *Group created *Group
@@ -466,6 +474,10 @@ func (s *groupRepoStubForInvalidRequestFallback) BindAccountsToGroup(_ context.C
panic("unexpected BindAccountsToGroup call") panic("unexpected BindAccountsToGroup call")
} }
func (s *groupRepoStubForInvalidRequestFallback) UpdateSortOrders(_ context.Context, _ []GroupSortOrderUpdate) error {
return nil
}
func TestAdminService_CreateGroup_InvalidRequestFallbackRejectsUnsupportedPlatform(t *testing.T) { func TestAdminService_CreateGroup_InvalidRequestFallbackRejectsUnsupportedPlatform(t *testing.T) {
fallbackID := int64(10) fallbackID := int64(10)
repo := &groupRepoStubForInvalidRequestFallback{ repo := &groupRepoStubForInvalidRequestFallback{

View File

@@ -0,0 +1,79 @@
package service
import (
"encoding/json"
"strings"
"time"
)
// Anthropic 会话 Fallback 相关常量
const (
// anthropicSessionTTLSeconds Anthropic 会话缓存 TTL5 分钟)
anthropicSessionTTLSeconds = 300
// anthropicDigestSessionKeyPrefix Anthropic 摘要 fallback 会话 key 前缀
anthropicDigestSessionKeyPrefix = "anthropic:digest:"
)
// AnthropicSessionTTL 返回 Anthropic 会话缓存 TTL
func AnthropicSessionTTL() time.Duration {
return anthropicSessionTTLSeconds * time.Second
}
// BuildAnthropicDigestChain 根据 Anthropic 请求生成摘要链
// 格式: s:<hash>-u:<hash>-a:<hash>-u:<hash>-...
// s = system, u = user, a = assistant
func BuildAnthropicDigestChain(parsed *ParsedRequest) string {
if parsed == nil {
return ""
}
var parts []string
// 1. system prompt
if parsed.System != nil {
systemData, _ := json.Marshal(parsed.System)
if len(systemData) > 0 && string(systemData) != "null" {
parts = append(parts, "s:"+shortHash(systemData))
}
}
// 2. messages
for _, msg := range parsed.Messages {
msgMap, ok := msg.(map[string]any)
if !ok {
continue
}
role, _ := msgMap["role"].(string)
prefix := rolePrefix(role)
content := msgMap["content"]
contentData, _ := json.Marshal(content)
parts = append(parts, prefix+":"+shortHash(contentData))
}
return strings.Join(parts, "-")
}
// rolePrefix 将 Anthropic 的 role 映射为单字符前缀
func rolePrefix(role string) string {
switch role {
case "assistant":
return "a"
default:
return "u"
}
}
// GenerateAnthropicDigestSessionKey 生成 Anthropic 摘要 fallback 的 sessionKey
// 组合 prefixHash 前 8 位 + uuid 前 8 位,确保不同会话产生不同的 sessionKey
func GenerateAnthropicDigestSessionKey(prefixHash, uuid string) string {
prefix := prefixHash
if len(prefixHash) >= 8 {
prefix = prefixHash[:8]
}
uuidPart := uuid
if len(uuid) >= 8 {
uuidPart = uuid[:8]
}
return anthropicDigestSessionKeyPrefix + prefix + ":" + uuidPart
}

View File

@@ -0,0 +1,320 @@
package service
import (
"strings"
"testing"
)
func TestBuildAnthropicDigestChain_NilRequest(t *testing.T) {
result := BuildAnthropicDigestChain(nil)
if result != "" {
t.Errorf("expected empty string for nil request, got: %s", result)
}
}
func TestBuildAnthropicDigestChain_EmptyMessages(t *testing.T) {
parsed := &ParsedRequest{
Messages: []any{},
}
result := BuildAnthropicDigestChain(parsed)
if result != "" {
t.Errorf("expected empty string for empty messages, got: %s", result)
}
}
func TestBuildAnthropicDigestChain_SingleUserMessage(t *testing.T) {
parsed := &ParsedRequest{
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 1 {
t.Fatalf("expected 1 part, got %d: %s", len(parts), result)
}
if !strings.HasPrefix(parts[0], "u:") {
t.Errorf("expected prefix 'u:', got: %s", parts[0])
}
}
func TestBuildAnthropicDigestChain_UserAndAssistant(t *testing.T) {
parsed := &ParsedRequest{
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "hi there"},
},
}
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 2 {
t.Fatalf("expected 2 parts, got %d: %s", len(parts), result)
}
if !strings.HasPrefix(parts[0], "u:") {
t.Errorf("part[0] expected prefix 'u:', got: %s", parts[0])
}
if !strings.HasPrefix(parts[1], "a:") {
t.Errorf("part[1] expected prefix 'a:', got: %s", parts[1])
}
}
func TestBuildAnthropicDigestChain_WithSystemString(t *testing.T) {
parsed := &ParsedRequest{
System: "You are a helpful assistant",
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 2 {
t.Fatalf("expected 2 parts (s + u), got %d: %s", len(parts), result)
}
if !strings.HasPrefix(parts[0], "s:") {
t.Errorf("part[0] expected prefix 's:', got: %s", parts[0])
}
if !strings.HasPrefix(parts[1], "u:") {
t.Errorf("part[1] expected prefix 'u:', got: %s", parts[1])
}
}
func TestBuildAnthropicDigestChain_WithSystemContentBlocks(t *testing.T) {
parsed := &ParsedRequest{
System: []any{
map[string]any{"type": "text", "text": "You are a helpful assistant"},
},
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 2 {
t.Fatalf("expected 2 parts (s + u), got %d: %s", len(parts), result)
}
if !strings.HasPrefix(parts[0], "s:") {
t.Errorf("part[0] expected prefix 's:', got: %s", parts[0])
}
}
func TestBuildAnthropicDigestChain_ConversationPrefixRelationship(t *testing.T) {
// 核心测试:验证对话增长时链的前缀关系
// 上一轮的完整链一定是下一轮链的前缀
system := "You are a helpful assistant"
// 第 1 轮: system + user
round1 := &ParsedRequest{
System: system,
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
chain1 := BuildAnthropicDigestChain(round1)
// 第 2 轮: system + user + assistant + user
round2 := &ParsedRequest{
System: system,
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "hi there"},
map[string]any{"role": "user", "content": "how are you?"},
},
}
chain2 := BuildAnthropicDigestChain(round2)
// 第 3 轮: system + user + assistant + user + assistant + user
round3 := &ParsedRequest{
System: system,
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "hi there"},
map[string]any{"role": "user", "content": "how are you?"},
map[string]any{"role": "assistant", "content": "I'm doing well"},
map[string]any{"role": "user", "content": "great"},
},
}
chain3 := BuildAnthropicDigestChain(round3)
t.Logf("Chain1: %s", chain1)
t.Logf("Chain2: %s", chain2)
t.Logf("Chain3: %s", chain3)
// chain1 是 chain2 的前缀
if !strings.HasPrefix(chain2, chain1) {
t.Errorf("chain1 should be prefix of chain2:\n chain1: %s\n chain2: %s", chain1, chain2)
}
// chain2 是 chain3 的前缀
if !strings.HasPrefix(chain3, chain2) {
t.Errorf("chain2 should be prefix of chain3:\n chain2: %s\n chain3: %s", chain2, chain3)
}
// chain1 也是 chain3 的前缀(传递性)
if !strings.HasPrefix(chain3, chain1) {
t.Errorf("chain1 should be prefix of chain3:\n chain1: %s\n chain3: %s", chain1, chain3)
}
}
func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *testing.T) {
parsed1 := &ParsedRequest{
System: "System A",
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
parsed2 := &ParsedRequest{
System: "System B",
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
chain1 := BuildAnthropicDigestChain(parsed1)
chain2 := BuildAnthropicDigestChain(parsed2)
if chain1 == chain2 {
t.Error("Different system prompts should produce different chains")
}
// 但 user 部分的 hash 应该相同
parts1 := splitChain(chain1)
parts2 := splitChain(chain2)
if parts1[1] != parts2[1] {
t.Error("Same user message should produce same hash regardless of system")
}
}
func TestBuildAnthropicDigestChain_DifferentContentProducesDifferentChain(t *testing.T) {
parsed1 := &ParsedRequest{
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "ORIGINAL reply"},
map[string]any{"role": "user", "content": "next"},
},
}
parsed2 := &ParsedRequest{
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "TAMPERED reply"},
map[string]any{"role": "user", "content": "next"},
},
}
chain1 := BuildAnthropicDigestChain(parsed1)
chain2 := BuildAnthropicDigestChain(parsed2)
if chain1 == chain2 {
t.Error("Different content should produce different chains")
}
parts1 := splitChain(chain1)
parts2 := splitChain(chain2)
// 第一个 user message hash 应该相同
if parts1[0] != parts2[0] {
t.Error("First user message hash should be the same")
}
// assistant reply hash 应该不同
if parts1[1] == parts2[1] {
t.Error("Assistant reply hash should differ")
}
}
func TestBuildAnthropicDigestChain_Deterministic(t *testing.T) {
parsed := &ParsedRequest{
System: "test system",
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "hi"},
},
}
chain1 := BuildAnthropicDigestChain(parsed)
chain2 := BuildAnthropicDigestChain(parsed)
if chain1 != chain2 {
t.Errorf("BuildAnthropicDigestChain not deterministic: %s vs %s", chain1, chain2)
}
}
func TestGenerateAnthropicDigestSessionKey(t *testing.T) {
tests := []struct {
name string
prefixHash string
uuid string
want string
}{
{
name: "normal 16 char hash with uuid",
prefixHash: "abcdefgh12345678",
uuid: "550e8400-e29b-41d4-a716-446655440000",
want: "anthropic:digest:abcdefgh:550e8400",
},
{
name: "exactly 8 chars",
prefixHash: "12345678",
uuid: "abcdefgh",
want: "anthropic:digest:12345678:abcdefgh",
},
{
name: "short values",
prefixHash: "abc",
uuid: "xyz",
want: "anthropic:digest:abc:xyz",
},
{
name: "empty values",
prefixHash: "",
uuid: "",
want: "anthropic:digest::",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := GenerateAnthropicDigestSessionKey(tt.prefixHash, tt.uuid)
if got != tt.want {
t.Errorf("GenerateAnthropicDigestSessionKey(%q, %q) = %q, want %q", tt.prefixHash, tt.uuid, got, tt.want)
}
})
}
// 验证不同 uuid 产生不同 sessionKey
t.Run("different uuid different key", func(t *testing.T) {
hash := "sameprefix123456"
result1 := GenerateAnthropicDigestSessionKey(hash, "uuid0001-session-a")
result2 := GenerateAnthropicDigestSessionKey(hash, "uuid0002-session-b")
if result1 == result2 {
t.Errorf("Different UUIDs should produce different session keys: %s vs %s", result1, result2)
}
})
}
func TestAnthropicSessionTTL(t *testing.T) {
ttl := AnthropicSessionTTL()
if ttl.Seconds() != 300 {
t.Errorf("expected 300 seconds, got: %v", ttl.Seconds())
}
}
func TestBuildAnthropicDigestChain_ContentBlocks(t *testing.T) {
// 测试 content 为 content blocks 数组的情况
parsed := &ParsedRequest{
Messages: []any{
map[string]any{
"role": "user",
"content": []any{
map[string]any{"type": "text", "text": "describe this image"},
map[string]any{"type": "image", "source": map[string]any{"type": "base64"}},
},
},
},
}
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 1 {
t.Fatalf("expected 1 part, got %d: %s", len(parts), result)
}
if !strings.HasPrefix(parts[0], "u:") {
t.Errorf("expected prefix 'u:', got: %s", parts[0])
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -4,18 +4,42 @@ import (
"bytes" "bytes"
"context" "context"
"encoding/json" "encoding/json"
"errors"
"fmt"
"io" "io"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"strings"
"testing" "testing"
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
// antigravityFailingWriter 模拟客户端断开连接的 gin.ResponseWriter
type antigravityFailingWriter struct {
gin.ResponseWriter
failAfter int // 允许成功写入的次数,之后所有写入返回错误
writes int
}
func (w *antigravityFailingWriter) Write(p []byte) (int, error) {
if w.writes >= w.failAfter {
return 0, errors.New("write failed: client disconnected")
}
w.writes++
return w.ResponseWriter.Write(p)
}
// newAntigravityTestService 创建用于流式测试的 AntigravityGatewayService
func newAntigravityTestService(cfg *config.Config) *AntigravityGatewayService {
return &AntigravityGatewayService{
settingService: &SettingService{cfg: cfg},
}
}
func TestStripSignatureSensitiveBlocksFromClaudeRequest(t *testing.T) { func TestStripSignatureSensitiveBlocksFromClaudeRequest(t *testing.T) {
req := &antigravity.ClaudeRequest{ req := &antigravity.ClaudeRequest{
Model: "claude-sonnet-4-5", Model: "claude-sonnet-4-5",
@@ -338,8 +362,8 @@ func TestAntigravityGatewayService_Forward_StickySessionForceCacheBilling(t *tes
require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch") require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch")
} }
// TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling // TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling verifies
// 验证:ForwardGemini 粘性会话切换时UpstreamFailoverError.ForceCacheBilling 应为 true // that ForwardGemini sets ForceCacheBilling=true for sticky session switch.
func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling(t *testing.T) { func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling(t *testing.T) {
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
writer := httptest.NewRecorder() writer := httptest.NewRecorder()
@@ -393,10 +417,16 @@ func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling(
require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch") require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch")
} }
func TestAntigravityStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) { // TestStreamUpstreamResponse_UsageAndFirstToken
// 验证usage 字段可被累积/覆盖更新,并且能记录首 token 时间
func TestStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) {
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
writer := httptest.NewRecorder() svc := newAntigravityTestService(&config.Config{
c, _ := gin.CreateTestContext(writer) Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
pr, pw := io.Pipe() pr, pw := io.Pipe()
@@ -404,25 +434,458 @@ func TestAntigravityStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) {
go func() { go func() {
defer func() { _ = pw.Close() }() defer func() { _ = pw.Close() }()
_, _ = pw.Write([]byte("data: {\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"cache_read_input_tokens\":3,\"cache_creation_input_tokens\":4}}\n")) fmt.Fprintln(pw, `data: {"usage":{"input_tokens":1,"output_tokens":2,"cache_read_input_tokens":3,"cache_creation_input_tokens":4}}`)
_, _ = pw.Write([]byte("data: {\"usage\":{\"output_tokens\":5}}\n")) fmt.Fprintln(pw, `data: {"usage":{"output_tokens":5}}`)
}() }()
svc := &AntigravityGatewayService{}
start := time.Now().Add(-10 * time.Millisecond) start := time.Now().Add(-10 * time.Millisecond)
usage, firstTokenMs := svc.streamUpstreamResponse(c, resp, start) result := svc.streamUpstreamResponse(c, resp, start)
_ = pr.Close() _ = pr.Close()
require.NotNil(t, usage) require.NotNil(t, result)
require.Equal(t, 1, usage.InputTokens) require.NotNil(t, result.usage)
require.Equal(t, 1, result.usage.InputTokens)
// 第二次事件覆盖 output_tokens // 第二次事件覆盖 output_tokens
require.Equal(t, 5, usage.OutputTokens) require.Equal(t, 5, result.usage.OutputTokens)
require.Equal(t, 3, usage.CacheReadInputTokens) require.Equal(t, 3, result.usage.CacheReadInputTokens)
require.Equal(t, 4, usage.CacheCreationInputTokens) require.Equal(t, 4, result.usage.CacheCreationInputTokens)
require.NotNil(t, result.firstTokenMs)
if firstTokenMs == nil {
t.Fatalf("expected firstTokenMs to be set")
}
// 确保有透传输出 // 确保有透传输出
require.True(t, strings.Contains(writer.Body.String(), "data:")) require.Contains(t, rec.Body.String(), "data:")
}
// --- 流式 happy path 测试 ---
// TestStreamUpstreamResponse_NormalComplete
// 验证正常流式转发完成时数据正确透传、usage 正确收集、clientDisconnect=false
func TestStreamUpstreamResponse_NormalComplete(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
go func() {
defer func() { _ = pw.Close() }()
fmt.Fprintln(pw, `event: message_start`)
fmt.Fprintln(pw, `data: {"type":"message_start","message":{"usage":{"input_tokens":10}}}`)
fmt.Fprintln(pw, "")
fmt.Fprintln(pw, `event: content_block_delta`)
fmt.Fprintln(pw, `data: {"type":"content_block_delta","delta":{"text":"hello"}}`)
fmt.Fprintln(pw, "")
fmt.Fprintln(pw, `event: message_delta`)
fmt.Fprintln(pw, `data: {"type":"message_delta","usage":{"output_tokens":5}}`)
fmt.Fprintln(pw, "")
}()
result := svc.streamUpstreamResponse(c, resp, time.Now())
_ = pr.Close()
require.NotNil(t, result)
require.False(t, result.clientDisconnect, "normal completion should not set clientDisconnect")
require.NotNil(t, result.usage)
require.Equal(t, 5, result.usage.OutputTokens, "should collect output_tokens from message_delta")
require.NotNil(t, result.firstTokenMs, "should record first token time")
// 验证数据被透传到客户端
body := rec.Body.String()
require.Contains(t, body, "event: message_start")
require.Contains(t, body, "content_block_delta")
require.Contains(t, body, "message_delta")
}
// TestHandleGeminiStreamingResponse_NormalComplete
// 验证:正常 Gemini 流式转发数据正确透传、usage 正确收集
func TestHandleGeminiStreamingResponse_NormalComplete(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
go func() {
defer func() { _ = pw.Close() }()
// 第一个 chunk部分内容
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":"Hello"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":3}}`)
fmt.Fprintln(pw, "")
// 第二个 chunk最终内容+完整 usage
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":" world"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":8,"cachedContentTokenCount":2}}`)
fmt.Fprintln(pw, "")
}()
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
_ = pr.Close()
require.NoError(t, err)
require.NotNil(t, result)
require.False(t, result.clientDisconnect, "normal completion should not set clientDisconnect")
require.NotNil(t, result.usage)
// Gemini usage: promptTokenCount=10, candidatesTokenCount=8, cachedContentTokenCount=2
// → InputTokens=10-2=8, OutputTokens=8, CacheReadInputTokens=2
require.Equal(t, 8, result.usage.InputTokens)
require.Equal(t, 8, result.usage.OutputTokens)
require.Equal(t, 2, result.usage.CacheReadInputTokens)
require.NotNil(t, result.firstTokenMs, "should record first token time")
// 验证数据被透传到客户端
body := rec.Body.String()
require.Contains(t, body, "Hello")
require.Contains(t, body, "world")
// 不应包含错误事件
require.NotContains(t, body, "event: error")
}
// TestHandleClaudeStreamingResponse_NormalComplete
// 验证:正常 Claude 流式转发Gemini→Claude 转换),数据正确转换并输出
func TestHandleClaudeStreamingResponse_NormalComplete(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
go func() {
defer func() { _ = pw.Close() }()
// v1internal 包装格式Gemini 数据嵌套在 "response" 字段下
// ProcessLine 先尝试反序列化为 V1InternalResponse裸格式会导致 Response.UsageMetadata 为空
fmt.Fprintln(pw, `data: {"response":{"candidates":[{"content":{"parts":[{"text":"Hi there"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3}}}`)
fmt.Fprintln(pw, "")
}()
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "claude-sonnet-4-5")
_ = pr.Close()
require.NoError(t, err)
require.NotNil(t, result)
require.False(t, result.clientDisconnect, "normal completion should not set clientDisconnect")
require.NotNil(t, result.usage)
// Gemini→Claude 转换的 usagepromptTokenCount=5→InputTokens=5, candidatesTokenCount=3→OutputTokens=3
require.Equal(t, 5, result.usage.InputTokens)
require.Equal(t, 3, result.usage.OutputTokens)
require.NotNil(t, result.firstTokenMs, "should record first token time")
// 验证输出是 Claude SSE 格式processor 会转换)
body := rec.Body.String()
require.Contains(t, body, "event: message_start", "should contain Claude message_start event")
require.Contains(t, body, "event: message_stop", "should contain Claude message_stop event")
// 不应包含错误事件
require.NotContains(t, body, "event: error")
}
// --- 流式客户端断开检测测试 ---
// TestStreamUpstreamResponse_ClientDisconnectDrainsUsage
// 验证客户端写入失败后streamUpstreamResponse 继续读取上游以收集 usage
func TestStreamUpstreamResponse_ClientDisconnectDrainsUsage(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
go func() {
defer func() { _ = pw.Close() }()
fmt.Fprintln(pw, `event: message_start`)
fmt.Fprintln(pw, `data: {"type":"message_start","message":{"usage":{"input_tokens":10}}}`)
fmt.Fprintln(pw, "")
fmt.Fprintln(pw, `event: message_delta`)
fmt.Fprintln(pw, `data: {"type":"message_delta","usage":{"output_tokens":20}}`)
fmt.Fprintln(pw, "")
}()
result := svc.streamUpstreamResponse(c, resp, time.Now())
_ = pr.Close()
require.NotNil(t, result)
require.True(t, result.clientDisconnect)
require.NotNil(t, result.usage)
require.Equal(t, 20, result.usage.OutputTokens)
}
// TestStreamUpstreamResponse_ContextCanceled
// 验证context 取消时返回 usage 且标记 clientDisconnect
func TestStreamUpstreamResponse_ContextCanceled(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
ctx, cancel := context.WithCancel(context.Background())
cancel()
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
resp := &http.Response{StatusCode: http.StatusOK, Body: cancelReadCloser{}, Header: http.Header{}}
result := svc.streamUpstreamResponse(c, resp, time.Now())
require.NotNil(t, result)
require.True(t, result.clientDisconnect)
require.NotContains(t, rec.Body.String(), "event: error")
}
// TestStreamUpstreamResponse_Timeout
// 验证:上游超时时返回已收集的 usage
func TestStreamUpstreamResponse_Timeout(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1, MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
result := svc.streamUpstreamResponse(c, resp, time.Now())
_ = pw.Close()
_ = pr.Close()
require.NotNil(t, result)
require.False(t, result.clientDisconnect)
}
// TestStreamUpstreamResponse_TimeoutAfterClientDisconnect
// 验证:客户端断开后上游超时,返回 usage 并标记 clientDisconnect
func TestStreamUpstreamResponse_TimeoutAfterClientDisconnect(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1, MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
go func() {
fmt.Fprintln(pw, `data: {"type":"message_start","message":{"usage":{"input_tokens":5}}}`)
fmt.Fprintln(pw, "")
// 不关闭 pw → 等待超时
}()
result := svc.streamUpstreamResponse(c, resp, time.Now())
_ = pw.Close()
_ = pr.Close()
require.NotNil(t, result)
require.True(t, result.clientDisconnect)
}
// TestHandleGeminiStreamingResponse_ClientDisconnect
// 验证Gemini 流式转发中客户端断开后继续 drain 上游
func TestHandleGeminiStreamingResponse_ClientDisconnect(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
go func() {
defer func() { _ = pw.Close() }()
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":10}}`)
fmt.Fprintln(pw, "")
}()
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
_ = pr.Close()
require.NoError(t, err)
require.NotNil(t, result)
require.True(t, result.clientDisconnect)
require.NotContains(t, rec.Body.String(), "write_failed")
}
// TestHandleGeminiStreamingResponse_ContextCanceled
// 验证context 取消时不注入错误事件
func TestHandleGeminiStreamingResponse_ContextCanceled(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
ctx, cancel := context.WithCancel(context.Background())
cancel()
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
resp := &http.Response{StatusCode: http.StatusOK, Body: cancelReadCloser{}, Header: http.Header{}}
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
require.NoError(t, err)
require.NotNil(t, result)
require.True(t, result.clientDisconnect)
require.NotContains(t, rec.Body.String(), "event: error")
}
// TestHandleClaudeStreamingResponse_ClientDisconnect
// 验证Claude 流式转发中客户端断开后继续 drain 上游
func TestHandleClaudeStreamingResponse_ClientDisconnect(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
go func() {
defer func() { _ = pw.Close() }()
// v1internal 包装格式
fmt.Fprintln(pw, `data: {"response":{"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":15}}}`)
fmt.Fprintln(pw, "")
}()
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "claude-sonnet-4-5")
_ = pr.Close()
require.NoError(t, err)
require.NotNil(t, result)
require.True(t, result.clientDisconnect)
}
// TestHandleClaudeStreamingResponse_ContextCanceled
// 验证context 取消时不注入错误事件
func TestHandleClaudeStreamingResponse_ContextCanceled(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
ctx, cancel := context.WithCancel(context.Background())
cancel()
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
resp := &http.Response{StatusCode: http.StatusOK, Body: cancelReadCloser{}, Header: http.Header{}}
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "claude-sonnet-4-5")
require.NoError(t, err)
require.NotNil(t, result)
require.True(t, result.clientDisconnect)
require.NotContains(t, rec.Body.String(), "event: error")
}
// TestExtractSSEUsage 验证 extractSSEUsage 从 SSE data 行正确提取 usage
func TestExtractSSEUsage(t *testing.T) {
svc := &AntigravityGatewayService{}
tests := []struct {
name string
line string
expected ClaudeUsage
}{
{
name: "message_delta with output_tokens",
line: `data: {"type":"message_delta","usage":{"output_tokens":42}}`,
expected: ClaudeUsage{OutputTokens: 42},
},
{
name: "non-data line ignored",
line: `event: message_start`,
expected: ClaudeUsage{},
},
{
name: "top-level usage with all fields",
line: `data: {"usage":{"input_tokens":10,"output_tokens":20,"cache_read_input_tokens":5,"cache_creation_input_tokens":3}}`,
expected: ClaudeUsage{InputTokens: 10, OutputTokens: 20, CacheReadInputTokens: 5, CacheCreationInputTokens: 3},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
usage := &ClaudeUsage{}
svc.extractSSEUsage(tt.line, usage)
require.Equal(t, tt.expected, *usage)
})
}
}
// TestAntigravityClientWriter 验证 antigravityClientWriter 的断开检测
func TestAntigravityClientWriter(t *testing.T) {
t.Run("normal write succeeds", func(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
flusher, _ := c.Writer.(http.Flusher)
cw := newAntigravityClientWriter(c.Writer, flusher, "test")
ok := cw.Write([]byte("hello"))
require.True(t, ok)
require.False(t, cw.Disconnected())
require.Contains(t, rec.Body.String(), "hello")
})
t.Run("write failure marks disconnected", func(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
fw := &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
flusher, _ := c.Writer.(http.Flusher)
cw := newAntigravityClientWriter(fw, flusher, "test")
ok := cw.Write([]byte("hello"))
require.False(t, ok)
require.True(t, cw.Disconnected())
})
t.Run("subsequent writes are no-op", func(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
fw := &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
flusher, _ := c.Writer.(http.Flusher)
cw := newAntigravityClientWriter(fw, flusher, "test")
cw.Write([]byte("first"))
ok := cw.Fprintf("second %d", 2)
require.False(t, ok)
require.True(t, cw.Disconnected())
})
} }

View File

@@ -2,63 +2,23 @@ package service
import ( import (
"context" "context"
"slices"
"strings" "strings"
"time" "time"
) )
const antigravityQuotaScopesKey = "antigravity_quota_scopes"
// AntigravityQuotaScope 表示 Antigravity 的配额域
type AntigravityQuotaScope string
const (
AntigravityQuotaScopeClaude AntigravityQuotaScope = "claude"
AntigravityQuotaScopeGeminiText AntigravityQuotaScope = "gemini_text"
AntigravityQuotaScopeGeminiImage AntigravityQuotaScope = "gemini_image"
)
// IsScopeSupported 检查给定的 scope 是否在分组支持的 scope 列表中
func IsScopeSupported(supportedScopes []string, scope AntigravityQuotaScope) bool {
if len(supportedScopes) == 0 {
// 未配置时默认全部支持
return true
}
supported := slices.Contains(supportedScopes, string(scope))
return supported
}
// ResolveAntigravityQuotaScope 根据模型名称解析配额域(导出版本)
func ResolveAntigravityQuotaScope(requestedModel string) (AntigravityQuotaScope, bool) {
return resolveAntigravityQuotaScope(requestedModel)
}
// resolveAntigravityQuotaScope 根据模型名称解析配额域
func resolveAntigravityQuotaScope(requestedModel string) (AntigravityQuotaScope, bool) {
model := normalizeAntigravityModelName(requestedModel)
if model == "" {
return "", false
}
switch {
case strings.HasPrefix(model, "claude-"):
return AntigravityQuotaScopeClaude, true
case strings.HasPrefix(model, "gemini-"):
if isImageGenerationModel(model) {
return AntigravityQuotaScopeGeminiImage, true
}
return AntigravityQuotaScopeGeminiText, true
default:
return "", false
}
}
func normalizeAntigravityModelName(model string) string { func normalizeAntigravityModelName(model string) string {
normalized := strings.ToLower(strings.TrimSpace(model)) normalized := strings.ToLower(strings.TrimSpace(model))
normalized = strings.TrimPrefix(normalized, "models/") normalized = strings.TrimPrefix(normalized, "models/")
return normalized return normalized
} }
// IsSchedulableForModel 结合 Antigravity 配额域限流判断是否可调度。 // resolveAntigravityModelKey 根据请求的模型名解析限流 key
// 返回空字符串表示无法解析
func resolveAntigravityModelKey(requestedModel string) string {
return normalizeAntigravityModelName(requestedModel)
}
// IsSchedulableForModel 结合模型级限流判断是否可调度。
// 保持旧签名以兼容既有调用方;默认使用 context.Background()。 // 保持旧签名以兼容既有调用方;默认使用 context.Background()。
func (a *Account) IsSchedulableForModel(requestedModel string) bool { func (a *Account) IsSchedulableForModel(requestedModel string) bool {
return a.IsSchedulableForModelWithContext(context.Background(), requestedModel) return a.IsSchedulableForModelWithContext(context.Background(), requestedModel)
@@ -74,107 +34,20 @@ func (a *Account) IsSchedulableForModelWithContext(ctx context.Context, requeste
if a.isModelRateLimitedWithContext(ctx, requestedModel) { if a.isModelRateLimitedWithContext(ctx, requestedModel) {
return false return false
} }
if a.Platform != PlatformAntigravity { return true
return true
}
scope, ok := resolveAntigravityQuotaScope(requestedModel)
if !ok {
return true
}
resetAt := a.antigravityQuotaScopeResetAt(scope)
if resetAt == nil {
return true
}
now := time.Now()
return !now.Before(*resetAt)
} }
func (a *Account) antigravityQuotaScopeResetAt(scope AntigravityQuotaScope) *time.Time { // GetRateLimitRemainingTime 获取限流剩余时间(模型级限流)
if a == nil || a.Extra == nil || scope == "" {
return nil
}
rawScopes, ok := a.Extra[antigravityQuotaScopesKey].(map[string]any)
if !ok {
return nil
}
rawScope, ok := rawScopes[string(scope)].(map[string]any)
if !ok {
return nil
}
resetAtRaw, ok := rawScope["rate_limit_reset_at"].(string)
if !ok || strings.TrimSpace(resetAtRaw) == "" {
return nil
}
resetAt, err := time.Parse(time.RFC3339, resetAtRaw)
if err != nil {
return nil
}
return &resetAt
}
var antigravityAllScopes = []AntigravityQuotaScope{
AntigravityQuotaScopeClaude,
AntigravityQuotaScopeGeminiText,
AntigravityQuotaScopeGeminiImage,
}
func (a *Account) GetAntigravityScopeRateLimits() map[string]int64 {
if a == nil || a.Platform != PlatformAntigravity {
return nil
}
now := time.Now()
result := make(map[string]int64)
for _, scope := range antigravityAllScopes {
resetAt := a.antigravityQuotaScopeResetAt(scope)
if resetAt != nil && now.Before(*resetAt) {
remainingSec := int64(time.Until(*resetAt).Seconds())
if remainingSec > 0 {
result[string(scope)] = remainingSec
}
}
}
if len(result) == 0 {
return nil
}
return result
}
// GetQuotaScopeRateLimitRemainingTime 获取模型域限流剩余时间
// 返回 0 表示未限流或已过期
func (a *Account) GetQuotaScopeRateLimitRemainingTime(requestedModel string) time.Duration {
if a == nil || a.Platform != PlatformAntigravity {
return 0
}
scope, ok := resolveAntigravityQuotaScope(requestedModel)
if !ok {
return 0
}
resetAt := a.antigravityQuotaScopeResetAt(scope)
if resetAt == nil {
return 0
}
if remaining := time.Until(*resetAt); remaining > 0 {
return remaining
}
return 0
}
// GetRateLimitRemainingTime 获取限流剩余时间(模型限流和模型域限流取最大值)
// 返回 0 表示未限流或已过期 // 返回 0 表示未限流或已过期
func (a *Account) GetRateLimitRemainingTime(requestedModel string) time.Duration { func (a *Account) GetRateLimitRemainingTime(requestedModel string) time.Duration {
return a.GetRateLimitRemainingTimeWithContext(context.Background(), requestedModel) return a.GetRateLimitRemainingTimeWithContext(context.Background(), requestedModel)
} }
// GetRateLimitRemainingTimeWithContext 获取限流剩余时间(模型限流和模型域限流取最大值 // GetRateLimitRemainingTimeWithContext 获取限流剩余时间(模型限流)
// 返回 0 表示未限流或已过期 // 返回 0 表示未限流或已过期
func (a *Account) GetRateLimitRemainingTimeWithContext(ctx context.Context, requestedModel string) time.Duration { func (a *Account) GetRateLimitRemainingTimeWithContext(ctx context.Context, requestedModel string) time.Duration {
if a == nil { if a == nil {
return 0 return 0
} }
modelRemaining := a.GetModelRateLimitRemainingTimeWithContext(ctx, requestedModel) return a.GetModelRateLimitRemainingTimeWithContext(ctx, requestedModel)
scopeRemaining := a.GetQuotaScopeRateLimitRemainingTime(requestedModel)
if modelRemaining > scopeRemaining {
return modelRemaining
}
return scopeRemaining
} }

View File

@@ -59,12 +59,6 @@ func (s *stubAntigravityUpstream) DoWithTLS(req *http.Request, proxyURL string,
return s.Do(req, proxyURL, accountID, accountConcurrency) return s.Do(req, proxyURL, accountID, accountConcurrency)
} }
type scopeLimitCall struct {
accountID int64
scope AntigravityQuotaScope
resetAt time.Time
}
type rateLimitCall struct { type rateLimitCall struct {
accountID int64 accountID int64
resetAt time.Time resetAt time.Time
@@ -78,16 +72,10 @@ type modelRateLimitCall struct {
type stubAntigravityAccountRepo struct { type stubAntigravityAccountRepo struct {
AccountRepository AccountRepository
scopeCalls []scopeLimitCall
rateCalls []rateLimitCall rateCalls []rateLimitCall
modelRateLimitCalls []modelRateLimitCall modelRateLimitCalls []modelRateLimitCall
} }
func (s *stubAntigravityAccountRepo) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error {
s.scopeCalls = append(s.scopeCalls, scopeLimitCall{accountID: id, scope: scope, resetAt: resetAt})
return nil
}
func (s *stubAntigravityAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { func (s *stubAntigravityAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
s.rateCalls = append(s.rateCalls, rateLimitCall{accountID: id, resetAt: resetAt}) s.rateCalls = append(s.rateCalls, rateLimitCall{accountID: id, resetAt: resetAt})
return nil return nil
@@ -131,10 +119,9 @@ func TestAntigravityRetryLoop_URLFallback_UsesLatestSuccess(t *testing.T) {
accessToken: "token", accessToken: "token",
action: "generateContent", action: "generateContent",
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
quotaScope: AntigravityQuotaScopeClaude,
httpUpstream: upstream, httpUpstream: upstream,
requestedModel: "claude-sonnet-4-5", requestedModel: "claude-sonnet-4-5",
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
handleErrorCalled = true handleErrorCalled = true
return nil return nil
}, },
@@ -155,23 +142,6 @@ func TestAntigravityRetryLoop_URLFallback_UsesLatestSuccess(t *testing.T) {
require.Equal(t, base2, available[0]) require.Equal(t, base2, available[0])
} }
func TestAntigravityHandleUpstreamError_UsesScopeLimit(t *testing.T) {
// 分区限流始终开启,不再支持通过环境变量关闭
repo := &stubAntigravityAccountRepo{}
svc := &AntigravityGatewayService{accountRepo: repo}
account := &Account{ID: 9, Name: "acc-9", Platform: PlatformAntigravity}
body := buildGeminiRateLimitBody("3s")
svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, AntigravityQuotaScopeClaude, 0, "", false)
require.Len(t, repo.scopeCalls, 1)
require.Empty(t, repo.rateCalls)
call := repo.scopeCalls[0]
require.Equal(t, account.ID, call.accountID)
require.Equal(t, AntigravityQuotaScopeClaude, call.scope)
require.WithinDuration(t, time.Now().Add(3*time.Second), call.resetAt, 2*time.Second)
}
// TestHandleUpstreamError_429_ModelRateLimit 测试 429 模型限流场景 // TestHandleUpstreamError_429_ModelRateLimit 测试 429 模型限流场景
func TestHandleUpstreamError_429_ModelRateLimit(t *testing.T) { func TestHandleUpstreamError_429_ModelRateLimit(t *testing.T) {
repo := &stubAntigravityAccountRepo{} repo := &stubAntigravityAccountRepo{}
@@ -189,7 +159,7 @@ func TestHandleUpstreamError_429_ModelRateLimit(t *testing.T) {
} }
}`) }`)
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, AntigravityQuotaScopeClaude, 0, "", false) result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, "claude-sonnet-4-5", 0, "", false)
// 应该触发模型限流 // 应该触发模型限流
require.NotNil(t, result) require.NotNil(t, result)
@@ -200,22 +170,22 @@ func TestHandleUpstreamError_429_ModelRateLimit(t *testing.T) {
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey) require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
} }
// TestHandleUpstreamError_429_NonModelRateLimit 测试 429 非模型限流场景(走 scope 限流 // TestHandleUpstreamError_429_NonModelRateLimit 测试 429 非模型限流场景(走模型级限流兜底
func TestHandleUpstreamError_429_NonModelRateLimit(t *testing.T) { func TestHandleUpstreamError_429_NonModelRateLimit(t *testing.T) {
repo := &stubAntigravityAccountRepo{} repo := &stubAntigravityAccountRepo{}
svc := &AntigravityGatewayService{accountRepo: repo} svc := &AntigravityGatewayService{accountRepo: repo}
account := &Account{ID: 2, Name: "acc-2", Platform: PlatformAntigravity} account := &Account{ID: 2, Name: "acc-2", Platform: PlatformAntigravity}
// 429 + 普通限流响应(无 RATE_LIMIT_EXCEEDED reasonscope 限流 // 429 + 普通限流响应(无 RATE_LIMIT_EXCEEDED reason走模型级限流兜底
body := buildGeminiRateLimitBody("5s") body := buildGeminiRateLimitBody("5s")
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, AntigravityQuotaScopeClaude, 0, "", false) result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, "claude-sonnet-4-5", 0, "", false)
// 不应该触发模型限流,应该走 scope 限流 // handleModelRateLimit 不会处理(因为没有 RATE_LIMIT_EXCEEDED
// 但 429 兜底逻辑会使用 requestedModel 设置模型级限流
require.Nil(t, result) require.Nil(t, result)
require.Empty(t, repo.modelRateLimitCalls) require.Len(t, repo.modelRateLimitCalls, 1)
require.Len(t, repo.scopeCalls, 1) require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
require.Equal(t, AntigravityQuotaScopeClaude, repo.scopeCalls[0].scope)
} }
// TestHandleUpstreamError_503_ModelRateLimit 测试 503 模型限流场景 // TestHandleUpstreamError_503_ModelRateLimit 测试 503 模型限流场景
@@ -235,7 +205,7 @@ func TestHandleUpstreamError_503_ModelRateLimit(t *testing.T) {
} }
}`) }`)
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, AntigravityQuotaScopeGeminiText, 0, "", false) result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, "gemini-3-pro-high", 0, "", false)
// 应该触发模型限流 // 应该触发模型限流
require.NotNil(t, result) require.NotNil(t, result)
@@ -263,12 +233,11 @@ func TestHandleUpstreamError_503_NonModelRateLimit(t *testing.T) {
} }
}`) }`)
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, AntigravityQuotaScopeGeminiText, 0, "", false) result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, "gemini-3-pro-high", 0, "", false)
// 503 非模型限流不应该做任何处理 // 503 非模型限流不应该做任何处理
require.Nil(t, result) require.Nil(t, result)
require.Empty(t, repo.modelRateLimitCalls, "503 non-model rate limit should not trigger model rate limit") require.Empty(t, repo.modelRateLimitCalls, "503 non-model rate limit should not trigger model rate limit")
require.Empty(t, repo.scopeCalls, "503 non-model rate limit should not trigger scope rate limit")
require.Empty(t, repo.rateCalls, "503 non-model rate limit should not trigger account rate limit") require.Empty(t, repo.rateCalls, "503 non-model rate limit should not trigger account rate limit")
} }
@@ -281,12 +250,11 @@ func TestHandleUpstreamError_503_EmptyBody(t *testing.T) {
// 503 + 空响应体 → 不做任何处理 // 503 + 空响应体 → 不做任何处理
body := []byte(`{}`) body := []byte(`{}`)
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, AntigravityQuotaScopeGeminiText, 0, "", false) result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, "gemini-3-pro-high", 0, "", false)
// 503 空响应不应该做任何处理 // 503 空响应不应该做任何处理
require.Nil(t, result) require.Nil(t, result)
require.Empty(t, repo.modelRateLimitCalls) require.Empty(t, repo.modelRateLimitCalls)
require.Empty(t, repo.scopeCalls)
require.Empty(t, repo.rateCalls) require.Empty(t, repo.rateCalls)
} }
@@ -307,15 +275,7 @@ func TestAccountIsSchedulableForModel_AntigravityRateLimits(t *testing.T) {
require.False(t, account.IsSchedulableForModel("gemini-3-flash")) require.False(t, account.IsSchedulableForModel("gemini-3-flash"))
account.RateLimitResetAt = nil account.RateLimitResetAt = nil
account.Extra = map[string]any{ require.True(t, account.IsSchedulableForModel("claude-sonnet-4-5"))
antigravityQuotaScopesKey: map[string]any{
"claude": map[string]any{
"rate_limit_reset_at": future.Format(time.RFC3339),
},
},
}
require.False(t, account.IsSchedulableForModel("claude-sonnet-4-5"))
require.True(t, account.IsSchedulableForModel("gemini-3-flash")) require.True(t, account.IsSchedulableForModel("gemini-3-flash"))
} }
@@ -635,6 +595,7 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
}`, }`,
expectedShouldRetry: false, expectedShouldRetry: false,
expectedShouldRateLimit: true, expectedShouldRateLimit: true,
minWait: 7 * time.Second,
modelName: "gemini-pro", modelName: "gemini-pro",
}, },
{ {
@@ -652,6 +613,7 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
}`, }`,
expectedShouldRetry: false, expectedShouldRetry: false,
expectedShouldRateLimit: true, expectedShouldRateLimit: true,
minWait: 39 * time.Second,
modelName: "gemini-3-pro-high", modelName: "gemini-3-pro-high",
}, },
{ {
@@ -669,6 +631,7 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
}`, }`,
expectedShouldRetry: false, expectedShouldRetry: false,
expectedShouldRateLimit: true, expectedShouldRateLimit: true,
minWait: 30 * time.Second,
modelName: "gemini-2.5-flash", modelName: "gemini-2.5-flash",
}, },
{ {
@@ -686,6 +649,7 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
}`, }`,
expectedShouldRetry: false, expectedShouldRetry: false,
expectedShouldRateLimit: true, expectedShouldRateLimit: true,
minWait: 30 * time.Second,
modelName: "claude-sonnet-4-5", modelName: "claude-sonnet-4-5",
}, },
} }
@@ -704,6 +668,11 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
t.Errorf("wait = %v, want >= %v", wait, tt.minWait) t.Errorf("wait = %v, want >= %v", wait, tt.minWait)
} }
} }
if shouldRateLimit && tt.minWait > 0 {
if wait < tt.minWait {
t.Errorf("rate limit wait = %v, want >= %v", wait, tt.minWait)
}
}
if (shouldRetry || shouldRateLimit) && model != tt.modelName { if (shouldRetry || shouldRateLimit) && model != tt.modelName {
t.Errorf("modelName = %q, want %q", model, tt.modelName) t.Errorf("modelName = %q, want %q", model, tt.modelName)
} }
@@ -803,7 +772,7 @@ func TestSetModelRateLimitByModelName_NotConvertToScope(t *testing.T) {
require.NotEqual(t, "claude_sonnet", call.modelKey, "should NOT be scope") require.NotEqual(t, "claude_sonnet", call.modelKey, "should NOT be scope")
} }
func TestAntigravityRetryLoop_PreCheck_WaitsWhenRemainingBelowThreshold(t *testing.T) { func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRateLimited(t *testing.T) {
upstream := &recordingOKUpstream{} upstream := &recordingOKUpstream{}
account := &Account{ account := &Account{
ID: 1, ID: 1,
@@ -815,19 +784,15 @@ func TestAntigravityRetryLoop_PreCheck_WaitsWhenRemainingBelowThreshold(t *testi
Extra: map[string]any{ Extra: map[string]any{
modelRateLimitsKey: map[string]any{ modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{ "claude-sonnet-4-5": map[string]any{
// RFC3339 here is second-precision; keep it safely in the future.
"rate_limit_reset_at": time.Now().Add(2 * time.Second).Format(time.RFC3339), "rate_limit_reset_at": time.Now().Add(2 * time.Second).Format(time.RFC3339),
}, },
}, },
}, },
} }
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
defer cancel()
svc := &AntigravityGatewayService{} svc := &AntigravityGatewayService{}
result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{ result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
ctx: ctx, ctx: context.Background(),
prefix: "[test]", prefix: "[test]",
account: account, account: account,
accessToken: "token", accessToken: "token",
@@ -836,17 +801,21 @@ func TestAntigravityRetryLoop_PreCheck_WaitsWhenRemainingBelowThreshold(t *testi
requestedModel: "claude-sonnet-4-5", requestedModel: "claude-sonnet-4-5",
httpUpstream: upstream, httpUpstream: upstream,
isStickySession: true, isStickySession: true,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
}) })
require.ErrorIs(t, err, context.DeadlineExceeded)
require.Nil(t, result) require.Nil(t, result)
require.Equal(t, 0, upstream.calls, "should not call upstream while waiting on pre-check") var switchErr *AntigravityAccountSwitchError
require.ErrorAs(t, err, &switchErr)
require.Equal(t, account.ID, switchErr.OriginalAccountID)
require.Equal(t, "claude-sonnet-4-5", switchErr.RateLimitedModel)
require.True(t, switchErr.IsStickySession)
require.Equal(t, 0, upstream.calls, "should not call upstream when switching on pre-check")
} }
func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRemainingAtOrAboveThreshold(t *testing.T) { func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRemainingLong(t *testing.T) {
upstream := &recordingOKUpstream{} upstream := &recordingOKUpstream{}
account := &Account{ account := &Account{
ID: 2, ID: 2,
@@ -875,7 +844,7 @@ func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRemainingAtOrAboveThreshold(t
requestedModel: "claude-sonnet-4-5", requestedModel: "claude-sonnet-4-5",
httpUpstream: upstream, httpUpstream: upstream,
isStickySession: true, isStickySession: true,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
}) })

View File

@@ -13,6 +13,23 @@ import (
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
// stubSmartRetryCache 用于 handleSmartRetry 测试的 GatewayCache mock
// 仅关注 DeleteSessionAccountID 的调用记录
type stubSmartRetryCache struct {
GatewayCache // 嵌入接口,未实现的方法 panic确保只调用预期方法
deleteCalls []deleteSessionCall
}
type deleteSessionCall struct {
groupID int64
sessionHash string
}
func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID int64, sessionHash string) error {
c.deleteCalls = append(c.deleteCalls, deleteSessionCall{groupID: groupID, sessionHash: sessionHash})
return nil
}
// mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream // mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream
type mockSmartRetryUpstream struct { type mockSmartRetryUpstream struct {
responses []*http.Response responses []*http.Response
@@ -58,7 +75,7 @@ func TestHandleSmartRetry_URLLevelRateLimit(t *testing.T) {
accessToken: "token", accessToken: "token",
action: "generateContent", action: "generateContent",
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -110,7 +127,7 @@ func TestHandleSmartRetry_LongDelay_ReturnsSwitchError(t *testing.T) {
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
accountRepo: repo, accountRepo: repo,
isStickySession: true, isStickySession: true,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -177,7 +194,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetrySuccess(t *testing.T) {
action: "generateContent", action: "generateContent",
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
httpUpstream: upstream, httpUpstream: upstream,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -198,7 +215,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetrySuccess(t *testing.T) {
// TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError 测试智能重试失败后返回 switchError // TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError 测试智能重试失败后返回 switchError
func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *testing.T) { func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *testing.T) {
// 智能重试后仍然返回 429需要提供 3 个响应,因为智能重试最多 3 次) // 智能重试后仍然返回 429需要提供 1 个响应,因为智能重试最多 1 次)
failRespBody := `{ failRespBody := `{
"error": { "error": {
"status": "RESOURCE_EXHAUSTED", "status": "RESOURCE_EXHAUSTED",
@@ -213,19 +230,9 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test
Header: http.Header{}, Header: http.Header{},
Body: io.NopCloser(strings.NewReader(failRespBody)), Body: io.NopCloser(strings.NewReader(failRespBody)),
} }
failResp2 := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(failRespBody)),
}
failResp3 := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(failRespBody)),
}
upstream := &mockSmartRetryUpstream{ upstream := &mockSmartRetryUpstream{
responses: []*http.Response{failResp1, failResp2, failResp3}, responses: []*http.Response{failResp1},
errors: []error{nil, nil, nil}, errors: []error{nil},
} }
repo := &stubAntigravityAccountRepo{} repo := &stubAntigravityAccountRepo{}
@@ -236,7 +243,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test
Platform: PlatformAntigravity, Platform: PlatformAntigravity,
} }
// 3s < 7s 阈值,应该触发智能重试(最多 3 次) // 3s < 7s 阈值,应该触发智能重试(最多 1 次)
respBody := []byte(`{ respBody := []byte(`{
"error": { "error": {
"status": "RESOURCE_EXHAUSTED", "status": "RESOURCE_EXHAUSTED",
@@ -262,7 +269,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test
httpUpstream: upstream, httpUpstream: upstream,
accountRepo: repo, accountRepo: repo,
isStickySession: false, isStickySession: false,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -284,7 +291,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test
// 验证模型限流已设置 // 验证模型限流已设置
require.Len(t, repo.modelRateLimitCalls, 1) require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey)
require.Len(t, upstream.calls, 3, "should have made three retry calls (max attempts)") require.Len(t, upstream.calls, 1, "should have made one retry call (max attempts)")
} }
// TestHandleSmartRetry_503_ModelCapacityExhausted_ReturnsSwitchError 测试 503 MODEL_CAPACITY_EXHAUSTED 返回 switchError // TestHandleSmartRetry_503_ModelCapacityExhausted_ReturnsSwitchError 测试 503 MODEL_CAPACITY_EXHAUSTED 返回 switchError
@@ -324,7 +331,7 @@ func TestHandleSmartRetry_503_ModelCapacityExhausted_ReturnsSwitchError(t *testi
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
accountRepo: repo, accountRepo: repo,
isStickySession: true, isStickySession: true,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -380,7 +387,7 @@ func TestHandleSmartRetry_NonAntigravityAccount_ContinuesDefaultLogic(t *testing
accessToken: "token", accessToken: "token",
action: "generateContent", action: "generateContent",
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -429,7 +436,7 @@ func TestHandleSmartRetry_NonModelRateLimit_ContinuesDefaultLogic(t *testing.T)
accessToken: "token", accessToken: "token",
action: "generateContent", action: "generateContent",
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -480,7 +487,7 @@ func TestHandleSmartRetry_ExactlyAtThreshold_ReturnsSwitchError(t *testing.T) {
action: "generateContent", action: "generateContent",
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
accountRepo: repo, accountRepo: repo,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -541,7 +548,7 @@ func TestAntigravityRetryLoop_HandleSmartRetry_SwitchError_Propagates(t *testing
httpUpstream: upstream, httpUpstream: upstream,
accountRepo: repo, accountRepo: repo,
isStickySession: true, isStickySession: true,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
}) })
@@ -556,19 +563,15 @@ func TestAntigravityRetryLoop_HandleSmartRetry_SwitchError_Propagates(t *testing
require.True(t, switchErr.IsStickySession) require.True(t, switchErr.IsStickySession)
} }
// TestHandleSmartRetry_NetworkError_ContinuesRetry 测试网络错误时继续重试 // TestHandleSmartRetry_NetworkError_ExhaustsRetry 测试网络错误时maxAttempts=1直接耗尽重试并切换账号
func TestHandleSmartRetry_NetworkError_ContinuesRetry(t *testing.T) { func TestHandleSmartRetry_NetworkError_ExhaustsRetry(t *testing.T) {
// 第一次网络错误,第二次成功 // 唯一一次重试遇到网络错误nil response
successResp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
}
upstream := &mockSmartRetryUpstream{ upstream := &mockSmartRetryUpstream{
responses: []*http.Response{nil, successResp}, // 第一次返回 nil模拟网络错误 responses: []*http.Response{nil}, // 返回 nil模拟网络错误
errors: []error{nil, nil}, // mock 不返回 error靠 nil response 触发 errors: []error{nil}, // mock 不返回 error靠 nil response 触发
} }
repo := &stubAntigravityAccountRepo{}
account := &Account{ account := &Account{
ID: 8, ID: 8,
Name: "acc-8", Name: "acc-8",
@@ -600,7 +603,8 @@ func TestHandleSmartRetry_NetworkError_ContinuesRetry(t *testing.T) {
action: "generateContent", action: "generateContent",
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
httpUpstream: upstream, httpUpstream: upstream,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { accountRepo: repo,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -612,10 +616,15 @@ func TestHandleSmartRetry_NetworkError_ContinuesRetry(t *testing.T) {
require.NotNil(t, result) require.NotNil(t, result)
require.Equal(t, smartRetryActionBreakWithResp, result.action) require.Equal(t, smartRetryActionBreakWithResp, result.action)
require.NotNil(t, result.resp, "should return successful response after network error recovery") require.Nil(t, result.resp, "should not return resp when switchError is set")
require.Equal(t, http.StatusOK, result.resp.StatusCode) require.NotNil(t, result.switchError, "should return switchError after network error exhausted retry")
require.Nil(t, result.switchError, "should not return switchError on success") require.Equal(t, account.ID, result.switchError.OriginalAccountID)
require.Len(t, upstream.calls, 2, "should have made two retry calls") require.Equal(t, "claude-sonnet-4-5", result.switchError.RateLimitedModel)
require.Len(t, upstream.calls, 1, "should have made one retry call")
// 验证模型限流已设置
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
} }
// TestHandleSmartRetry_NoRetryDelay_UsesDefaultRateLimit 测试无 retryDelay 时使用默认 1 分钟限流 // TestHandleSmartRetry_NoRetryDelay_UsesDefaultRateLimit 测试无 retryDelay 时使用默认 1 分钟限流
@@ -653,7 +662,7 @@ func TestHandleSmartRetry_NoRetryDelay_UsesDefaultRateLimit(t *testing.T) {
body: []byte(`{"input":"test"}`), body: []byte(`{"input":"test"}`),
accountRepo: repo, accountRepo: repo,
isStickySession: true, isStickySession: true,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil return nil
}, },
} }
@@ -674,3 +683,617 @@ func TestHandleSmartRetry_NoRetryDelay_UsesDefaultRateLimit(t *testing.T) {
require.Len(t, repo.modelRateLimitCalls, 1) require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey) require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
} }
// ---------------------------------------------------------------------------
// 以下测试覆盖本次改动:
// 1. antigravitySmartRetryMaxAttempts = 1仅重试 1 次)
// 2. 智能重试失败后清除粘性会话绑定DeleteSessionAccountID
// ---------------------------------------------------------------------------
// TestSmartRetryMaxAttempts_VerifyConstant 验证常量值为 1
func TestSmartRetryMaxAttempts_VerifyConstant(t *testing.T) {
require.Equal(t, 1, antigravitySmartRetryMaxAttempts,
"antigravitySmartRetryMaxAttempts should be 1 to prevent repeated rate limiting")
}
// TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_ClearsSession
// 核心场景:粘性会话 + 短延迟重试失败 → 必须清除粘性绑定
func TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_ClearsSession(t *testing.T) {
failRespBody := `{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
failResp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(failRespBody)),
}
upstream := &mockSmartRetryUpstream{
responses: []*http.Response{failResp},
errors: []error{nil},
}
repo := &stubAntigravityAccountRepo{}
cache := &stubSmartRetryCache{}
account := &Account{
ID: 10,
Name: "acc-10",
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
}
respBody := []byte(`{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`)
resp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(bytes.NewReader(respBody)),
}
params := antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: account,
accessToken: "token",
action: "generateContent",
body: []byte(`{"input":"test"}`),
httpUpstream: upstream,
accountRepo: repo,
isStickySession: true,
groupID: 42,
sessionHash: "sticky-hash-abc",
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil
},
}
availableURLs := []string{"https://ag-1.test"}
svc := &AntigravityGatewayService{cache: cache}
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
// 验证返回 switchError
require.NotNil(t, result)
require.Equal(t, smartRetryActionBreakWithResp, result.action)
require.NotNil(t, result.switchError)
require.True(t, result.switchError.IsStickySession, "switchError should carry IsStickySession=true")
require.Equal(t, account.ID, result.switchError.OriginalAccountID)
// 核心断言DeleteSessionAccountID 被调用,且参数正确
require.Len(t, cache.deleteCalls, 1, "should call DeleteSessionAccountID exactly once")
require.Equal(t, int64(42), cache.deleteCalls[0].groupID)
require.Equal(t, "sticky-hash-abc", cache.deleteCalls[0].sessionHash)
// 验证仅重试 1 次
require.Len(t, upstream.calls, 1, "should make exactly 1 retry call (maxAttempts=1)")
// 验证模型限流已设置
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
}
// TestHandleSmartRetry_ShortDelay_NonStickySession_FailedRetry_NoDeleteSession
// 非粘性会话 + 短延迟重试失败 → 不应调用 DeleteSessionAccountIDsessionHash 为空)
func TestHandleSmartRetry_ShortDelay_NonStickySession_FailedRetry_NoDeleteSession(t *testing.T) {
failRespBody := `{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
failResp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(failRespBody)),
}
upstream := &mockSmartRetryUpstream{
responses: []*http.Response{failResp},
errors: []error{nil},
}
repo := &stubAntigravityAccountRepo{}
cache := &stubSmartRetryCache{}
account := &Account{
ID: 11,
Name: "acc-11",
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
}
respBody := []byte(`{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`)
resp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(bytes.NewReader(respBody)),
}
params := antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: account,
accessToken: "token",
action: "generateContent",
body: []byte(`{"input":"test"}`),
httpUpstream: upstream,
accountRepo: repo,
isStickySession: false,
groupID: 42,
sessionHash: "", // 非粘性会话sessionHash 为空
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil
},
}
availableURLs := []string{"https://ag-1.test"}
svc := &AntigravityGatewayService{cache: cache}
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
require.NotNil(t, result)
require.Equal(t, smartRetryActionBreakWithResp, result.action)
require.NotNil(t, result.switchError)
require.False(t, result.switchError.IsStickySession)
// 核心断言sessionHash 为空时不应调用 DeleteSessionAccountID
require.Len(t, cache.deleteCalls, 0, "should NOT call DeleteSessionAccountID when sessionHash is empty")
}
// TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_NilCache_NoPanic
// 边界cache 为 nil 时不应 panic
func TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_NilCache_NoPanic(t *testing.T) {
failRespBody := `{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
failResp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(failRespBody)),
}
upstream := &mockSmartRetryUpstream{
responses: []*http.Response{failResp},
errors: []error{nil},
}
repo := &stubAntigravityAccountRepo{}
account := &Account{
ID: 12,
Name: "acc-12",
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
}
respBody := []byte(`{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`)
resp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(bytes.NewReader(respBody)),
}
params := antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: account,
accessToken: "token",
action: "generateContent",
body: []byte(`{"input":"test"}`),
httpUpstream: upstream,
accountRepo: repo,
isStickySession: true,
groupID: 42,
sessionHash: "sticky-hash-nil-cache",
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil
},
}
availableURLs := []string{"https://ag-1.test"}
// cache 为 nil不应 panic
svc := &AntigravityGatewayService{cache: nil}
require.NotPanics(t, func() {
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
require.NotNil(t, result)
require.Equal(t, smartRetryActionBreakWithResp, result.action)
require.NotNil(t, result.switchError)
require.True(t, result.switchError.IsStickySession)
})
}
// TestHandleSmartRetry_ShortDelay_StickySession_SuccessRetry_NoDeleteSession
// 重试成功时不应清除粘性会话(只有失败才清除)
func TestHandleSmartRetry_ShortDelay_StickySession_SuccessRetry_NoDeleteSession(t *testing.T) {
successResp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
}
upstream := &mockSmartRetryUpstream{
responses: []*http.Response{successResp},
errors: []error{nil},
}
cache := &stubSmartRetryCache{}
account := &Account{
ID: 13,
Name: "acc-13",
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
}
respBody := []byte(`{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
]
}
}`)
resp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(bytes.NewReader(respBody)),
}
params := antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: account,
accessToken: "token",
action: "generateContent",
body: []byte(`{"input":"test"}`),
httpUpstream: upstream,
isStickySession: true,
groupID: 42,
sessionHash: "sticky-hash-success",
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil
},
}
availableURLs := []string{"https://ag-1.test"}
svc := &AntigravityGatewayService{cache: cache}
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
require.NotNil(t, result)
require.Equal(t, smartRetryActionBreakWithResp, result.action)
require.NotNil(t, result.resp, "should return successful response")
require.Equal(t, http.StatusOK, result.resp.StatusCode)
require.Nil(t, result.switchError, "should not return switchError on success")
// 核心断言:重试成功时不应清除粘性会话
require.Len(t, cache.deleteCalls, 0, "should NOT call DeleteSessionAccountID on successful retry")
}
// TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry
// 长延迟路径情况1在 handleSmartRetry 中不直接调用 DeleteSessionAccountID
// (清除由 handler 层的 shouldClearStickySession 在下次请求时处理)
func TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry(t *testing.T) {
repo := &stubAntigravityAccountRepo{}
cache := &stubSmartRetryCache{}
account := &Account{
ID: 14,
Name: "acc-14",
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
}
// 15s >= 7s 阈值 → 走长延迟路径
respBody := []byte(`{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"}
]
}
}`)
resp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(bytes.NewReader(respBody)),
}
params := antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: account,
accessToken: "token",
action: "generateContent",
body: []byte(`{"input":"test"}`),
accountRepo: repo,
isStickySession: true,
groupID: 42,
sessionHash: "sticky-hash-long-delay",
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil
},
}
availableURLs := []string{"https://ag-1.test"}
svc := &AntigravityGatewayService{cache: cache}
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
require.NotNil(t, result)
require.Equal(t, smartRetryActionBreakWithResp, result.action)
require.NotNil(t, result.switchError)
require.True(t, result.switchError.IsStickySession)
// 长延迟路径不在 handleSmartRetry 中调用 DeleteSessionAccountID
// (由上游 handler 的 shouldClearStickySession 处理)
require.Len(t, cache.deleteCalls, 0,
"long delay path should NOT call DeleteSessionAccountID in handleSmartRetry (handled by handler layer)")
}
// TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession
// 网络错误耗尽重试 + 粘性会话 → 也应清除粘性绑定
func TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession(t *testing.T) {
upstream := &mockSmartRetryUpstream{
responses: []*http.Response{nil}, // 网络错误
errors: []error{nil},
}
repo := &stubAntigravityAccountRepo{}
cache := &stubSmartRetryCache{}
account := &Account{
ID: 15,
Name: "acc-15",
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
}
respBody := []byte(`{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`)
resp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(bytes.NewReader(respBody)),
}
params := antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: account,
accessToken: "token",
action: "generateContent",
body: []byte(`{"input":"test"}`),
httpUpstream: upstream,
accountRepo: repo,
isStickySession: true,
groupID: 99,
sessionHash: "sticky-net-error",
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil
},
}
availableURLs := []string{"https://ag-1.test"}
svc := &AntigravityGatewayService{cache: cache}
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
require.NotNil(t, result)
require.NotNil(t, result.switchError)
require.True(t, result.switchError.IsStickySession)
// 核心断言:网络错误耗尽重试后也应清除粘性绑定
require.Len(t, cache.deleteCalls, 1, "should call DeleteSessionAccountID after network error exhausts retry")
require.Equal(t, int64(99), cache.deleteCalls[0].groupID)
require.Equal(t, "sticky-net-error", cache.deleteCalls[0].sessionHash)
}
// TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession
// 503 + 短延迟 + 粘性会话 + 重试失败 → 清除粘性绑定
func TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession(t *testing.T) {
failRespBody := `{
"error": {
"code": 503,
"status": "UNAVAILABLE",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
]
}
}`
failResp := &http.Response{
StatusCode: http.StatusServiceUnavailable,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(failRespBody)),
}
upstream := &mockSmartRetryUpstream{
responses: []*http.Response{failResp},
errors: []error{nil},
}
repo := &stubAntigravityAccountRepo{}
cache := &stubSmartRetryCache{}
account := &Account{
ID: 16,
Name: "acc-16",
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
}
respBody := []byte(`{
"error": {
"code": 503,
"status": "UNAVAILABLE",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
]
}
}`)
resp := &http.Response{
StatusCode: http.StatusServiceUnavailable,
Header: http.Header{},
Body: io.NopCloser(bytes.NewReader(respBody)),
}
params := antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: account,
accessToken: "token",
action: "generateContent",
body: []byte(`{"input":"test"}`),
httpUpstream: upstream,
accountRepo: repo,
isStickySession: true,
groupID: 77,
sessionHash: "sticky-503-short",
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil
},
}
availableURLs := []string{"https://ag-1.test"}
svc := &AntigravityGatewayService{cache: cache}
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
require.NotNil(t, result)
require.NotNil(t, result.switchError)
require.True(t, result.switchError.IsStickySession)
// 验证粘性绑定被清除
require.Len(t, cache.deleteCalls, 1)
require.Equal(t, int64(77), cache.deleteCalls[0].groupID)
require.Equal(t, "sticky-503-short", cache.deleteCalls[0].sessionHash)
// 验证模型限流已设置
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "gemini-3-pro", repo.modelRateLimitCalls[0].modelKey)
}
// TestAntigravityRetryLoop_SmartRetryFailed_StickySession_SwitchErrorPropagates
// 集成测试antigravityRetryLoop → handleSmartRetry → switchError 传播
// 验证 IsStickySession 正确传递到上层,且粘性绑定被清除
func TestAntigravityRetryLoop_SmartRetryFailed_StickySession_SwitchErrorPropagates(t *testing.T) {
// 初始 429 响应
initialRespBody := []byte(`{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4-6"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`)
initialResp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(bytes.NewReader(initialRespBody)),
}
// 智能重试也返回 429
retryRespBody := `{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4-6"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
retryResp := &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(retryRespBody)),
}
upstream := &mockSmartRetryUpstream{
responses: []*http.Response{initialResp, retryResp},
errors: []error{nil, nil},
}
repo := &stubAntigravityAccountRepo{}
cache := &stubSmartRetryCache{}
account := &Account{
ID: 17,
Name: "acc-17",
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Schedulable: true,
Status: StatusActive,
Concurrency: 1,
}
svc := &AntigravityGatewayService{cache: cache}
result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: account,
accessToken: "token",
action: "generateContent",
body: []byte(`{"input":"test"}`),
httpUpstream: upstream,
accountRepo: repo,
isStickySession: true,
groupID: 55,
sessionHash: "sticky-loop-test",
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
return nil
},
})
require.Nil(t, result, "should not return result when switchError")
require.NotNil(t, err, "should return error")
var switchErr *AntigravityAccountSwitchError
require.ErrorAs(t, err, &switchErr, "error should be AntigravityAccountSwitchError")
require.Equal(t, account.ID, switchErr.OriginalAccountID)
require.Equal(t, "claude-opus-4-6", switchErr.RateLimitedModel)
require.True(t, switchErr.IsStickySession, "IsStickySession must propagate through retryLoop")
// 验证粘性绑定被清除
require.Len(t, cache.deleteCalls, 1, "should clear sticky session in handleSmartRetry")
require.Equal(t, int64(55), cache.deleteCalls[0].groupID)
require.Equal(t, "sticky-loop-test", cache.deleteCalls[0].sessionHash)
}

View File

@@ -0,0 +1,112 @@
package service
import (
"testing"
)
func TestBuildSelectedSet(t *testing.T) {
tests := []struct {
name string
ids []string
wantNil bool
wantSize int
}{
{
name: "nil input returns nil (backward compatible: create all)",
ids: nil,
wantNil: true,
},
{
name: "empty slice returns empty map (create none)",
ids: []string{},
wantNil: false,
wantSize: 0,
},
{
name: "single ID",
ids: []string{"abc-123"},
wantNil: false,
wantSize: 1,
},
{
name: "multiple IDs",
ids: []string{"a", "b", "c"},
wantNil: false,
wantSize: 3,
},
{
name: "duplicate IDs are deduplicated",
ids: []string{"a", "a", "b"},
wantNil: false,
wantSize: 2,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := buildSelectedSet(tt.ids)
if tt.wantNil {
if got != nil {
t.Errorf("buildSelectedSet(%v) = %v, want nil", tt.ids, got)
}
return
}
if got == nil {
t.Fatalf("buildSelectedSet(%v) = nil, want non-nil map", tt.ids)
}
if len(got) != tt.wantSize {
t.Errorf("buildSelectedSet(%v) has %d entries, want %d", tt.ids, len(got), tt.wantSize)
}
// Verify all unique IDs are present
for _, id := range tt.ids {
if _, ok := got[id]; !ok {
t.Errorf("buildSelectedSet(%v) missing key %q", tt.ids, id)
}
}
})
}
}
func TestShouldCreateAccount(t *testing.T) {
tests := []struct {
name string
crsID string
selectedSet map[string]struct{}
want bool
}{
{
name: "nil set allows all (backward compatible)",
crsID: "any-id",
selectedSet: nil,
want: true,
},
{
name: "empty set blocks all",
crsID: "any-id",
selectedSet: map[string]struct{}{},
want: false,
},
{
name: "ID in set is allowed",
crsID: "abc-123",
selectedSet: map[string]struct{}{"abc-123": {}, "def-456": {}},
want: true,
},
{
name: "ID not in set is blocked",
crsID: "xyz-789",
selectedSet: map[string]struct{}{"abc-123": {}, "def-456": {}},
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := shouldCreateAccount(tt.crsID, tt.selectedSet)
if got != tt.want {
t.Errorf("shouldCreateAccount(%q, %v) = %v, want %v",
tt.crsID, tt.selectedSet, got, tt.want)
}
})
}
}

View File

@@ -45,10 +45,11 @@ func NewCRSSyncService(
} }
type SyncFromCRSInput struct { type SyncFromCRSInput struct {
BaseURL string BaseURL string
Username string Username string
Password string Password string
SyncProxies bool SyncProxies bool
SelectedAccountIDs []string // if non-empty, only create new accounts with these CRS IDs
} }
type SyncFromCRSItemResult struct { type SyncFromCRSItemResult struct {
@@ -190,25 +191,27 @@ type crsGeminiAPIKeyAccount struct {
Extra map[string]any `json:"extra"` Extra map[string]any `json:"extra"`
} }
func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput) (*SyncFromCRSResult, error) { // fetchCRSExport validates the connection parameters, authenticates with CRS,
// and returns the exported accounts. Shared by SyncFromCRS and PreviewFromCRS.
func (s *CRSSyncService) fetchCRSExport(ctx context.Context, baseURL, username, password string) (*crsExportResponse, error) {
if s.cfg == nil { if s.cfg == nil {
return nil, errors.New("config is not available") return nil, errors.New("config is not available")
} }
baseURL := strings.TrimSpace(input.BaseURL) normalizedURL := strings.TrimSpace(baseURL)
if s.cfg.Security.URLAllowlist.Enabled { if s.cfg.Security.URLAllowlist.Enabled {
normalized, err := normalizeBaseURL(baseURL, s.cfg.Security.URLAllowlist.CRSHosts, s.cfg.Security.URLAllowlist.AllowPrivateHosts) normalized, err := normalizeBaseURL(normalizedURL, s.cfg.Security.URLAllowlist.CRSHosts, s.cfg.Security.URLAllowlist.AllowPrivateHosts)
if err != nil { if err != nil {
return nil, err return nil, err
} }
baseURL = normalized normalizedURL = normalized
} else { } else {
normalized, err := urlvalidator.ValidateURLFormat(baseURL, s.cfg.Security.URLAllowlist.AllowInsecureHTTP) normalized, err := urlvalidator.ValidateURLFormat(normalizedURL, s.cfg.Security.URLAllowlist.AllowInsecureHTTP)
if err != nil { if err != nil {
return nil, fmt.Errorf("invalid base_url: %w", err) return nil, fmt.Errorf("invalid base_url: %w", err)
} }
baseURL = normalized normalizedURL = normalized
} }
if strings.TrimSpace(input.Username) == "" || strings.TrimSpace(input.Password) == "" { if strings.TrimSpace(username) == "" || strings.TrimSpace(password) == "" {
return nil, errors.New("username and password are required") return nil, errors.New("username and password are required")
} }
@@ -221,12 +224,16 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
client = &http.Client{Timeout: 20 * time.Second} client = &http.Client{Timeout: 20 * time.Second}
} }
adminToken, err := crsLogin(ctx, client, baseURL, input.Username, input.Password) adminToken, err := crsLogin(ctx, client, normalizedURL, username, password)
if err != nil { if err != nil {
return nil, err return nil, err
} }
exported, err := crsExportAccounts(ctx, client, baseURL, adminToken) return crsExportAccounts(ctx, client, normalizedURL, adminToken)
}
func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput) (*SyncFromCRSResult, error) {
exported, err := s.fetchCRSExport(ctx, input.BaseURL, input.Username, input.Password)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -241,6 +248,8 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
), ),
} }
selectedSet := buildSelectedSet(input.SelectedAccountIDs)
var proxies []Proxy var proxies []Proxy
if input.SyncProxies { if input.SyncProxies {
proxies, _ = s.proxyRepo.ListActive(ctx) proxies, _ = s.proxyRepo.ListActive(ctx)
@@ -329,6 +338,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
} }
if existing == nil { if existing == nil {
if !shouldCreateAccount(src.ID, selectedSet) {
item.Action = "skipped"
item.Error = "not selected"
result.Skipped++
result.Items = append(result.Items, item)
continue
}
account := &Account{ account := &Account{
Name: defaultName(src.Name, src.ID), Name: defaultName(src.Name, src.ID),
Platform: PlatformAnthropic, Platform: PlatformAnthropic,
@@ -446,6 +462,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
} }
if existing == nil { if existing == nil {
if !shouldCreateAccount(src.ID, selectedSet) {
item.Action = "skipped"
item.Error = "not selected"
result.Skipped++
result.Items = append(result.Items, item)
continue
}
account := &Account{ account := &Account{
Name: defaultName(src.Name, src.ID), Name: defaultName(src.Name, src.ID),
Platform: PlatformAnthropic, Platform: PlatformAnthropic,
@@ -569,6 +592,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
} }
if existing == nil { if existing == nil {
if !shouldCreateAccount(src.ID, selectedSet) {
item.Action = "skipped"
item.Error = "not selected"
result.Skipped++
result.Items = append(result.Items, item)
continue
}
account := &Account{ account := &Account{
Name: defaultName(src.Name, src.ID), Name: defaultName(src.Name, src.ID),
Platform: PlatformOpenAI, Platform: PlatformOpenAI,
@@ -690,6 +720,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
} }
if existing == nil { if existing == nil {
if !shouldCreateAccount(src.ID, selectedSet) {
item.Action = "skipped"
item.Error = "not selected"
result.Skipped++
result.Items = append(result.Items, item)
continue
}
account := &Account{ account := &Account{
Name: defaultName(src.Name, src.ID), Name: defaultName(src.Name, src.ID),
Platform: PlatformOpenAI, Platform: PlatformOpenAI,
@@ -798,6 +835,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
} }
if existing == nil { if existing == nil {
if !shouldCreateAccount(src.ID, selectedSet) {
item.Action = "skipped"
item.Error = "not selected"
result.Skipped++
result.Items = append(result.Items, item)
continue
}
account := &Account{ account := &Account{
Name: defaultName(src.Name, src.ID), Name: defaultName(src.Name, src.ID),
Platform: PlatformGemini, Platform: PlatformGemini,
@@ -909,6 +953,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
} }
if existing == nil { if existing == nil {
if !shouldCreateAccount(src.ID, selectedSet) {
item.Action = "skipped"
item.Error = "not selected"
result.Skipped++
result.Items = append(result.Items, item)
continue
}
account := &Account{ account := &Account{
Name: defaultName(src.Name, src.ID), Name: defaultName(src.Name, src.ID),
Platform: PlatformGemini, Platform: PlatformGemini,
@@ -1253,3 +1304,102 @@ func (s *CRSSyncService) refreshOAuthToken(ctx context.Context, account *Account
return newCredentials return newCredentials
} }
// buildSelectedSet converts a slice of selected CRS account IDs to a set for O(1) lookup.
// Returns nil if ids is nil (field not sent → backward compatible: create all).
// Returns an empty map if ids is non-nil but empty (user selected none → create none).
func buildSelectedSet(ids []string) map[string]struct{} {
if ids == nil {
return nil
}
set := make(map[string]struct{}, len(ids))
for _, id := range ids {
set[id] = struct{}{}
}
return set
}
// shouldCreateAccount checks if a new CRS account should be created based on user selection.
// Returns true if selectedSet is nil (backward compatible: create all) or if crsID is in the set.
func shouldCreateAccount(crsID string, selectedSet map[string]struct{}) bool {
if selectedSet == nil {
return true
}
_, ok := selectedSet[crsID]
return ok
}
// PreviewFromCRSResult contains the preview of accounts from CRS before sync.
type PreviewFromCRSResult struct {
NewAccounts []CRSPreviewAccount `json:"new_accounts"`
ExistingAccounts []CRSPreviewAccount `json:"existing_accounts"`
}
// CRSPreviewAccount represents a single account in the preview result.
type CRSPreviewAccount struct {
CRSAccountID string `json:"crs_account_id"`
Kind string `json:"kind"`
Name string `json:"name"`
Platform string `json:"platform"`
Type string `json:"type"`
}
// PreviewFromCRS connects to CRS, fetches all accounts, and classifies them
// as new or existing by batch-querying local crs_account_id mappings.
func (s *CRSSyncService) PreviewFromCRS(ctx context.Context, input SyncFromCRSInput) (*PreviewFromCRSResult, error) {
exported, err := s.fetchCRSExport(ctx, input.BaseURL, input.Username, input.Password)
if err != nil {
return nil, err
}
// Batch query all existing CRS account IDs
existingCRSIDs, err := s.accountRepo.ListCRSAccountIDs(ctx)
if err != nil {
return nil, fmt.Errorf("failed to list existing CRS accounts: %w", err)
}
result := &PreviewFromCRSResult{
NewAccounts: make([]CRSPreviewAccount, 0),
ExistingAccounts: make([]CRSPreviewAccount, 0),
}
classify := func(crsID, kind, name, platform, accountType string) {
preview := CRSPreviewAccount{
CRSAccountID: crsID,
Kind: kind,
Name: defaultName(name, crsID),
Platform: platform,
Type: accountType,
}
if _, exists := existingCRSIDs[crsID]; exists {
result.ExistingAccounts = append(result.ExistingAccounts, preview)
} else {
result.NewAccounts = append(result.NewAccounts, preview)
}
}
for _, src := range exported.Data.ClaudeAccounts {
authType := strings.TrimSpace(src.AuthType)
if authType == "" {
authType = AccountTypeOAuth
}
classify(src.ID, src.Kind, src.Name, PlatformAnthropic, authType)
}
for _, src := range exported.Data.ClaudeConsoleAccounts {
classify(src.ID, src.Kind, src.Name, PlatformAnthropic, AccountTypeAPIKey)
}
for _, src := range exported.Data.OpenAIOAuthAccounts {
classify(src.ID, src.Kind, src.Name, PlatformOpenAI, AccountTypeOAuth)
}
for _, src := range exported.Data.OpenAIResponsesAccounts {
classify(src.ID, src.Kind, src.Name, PlatformOpenAI, AccountTypeAPIKey)
}
for _, src := range exported.Data.GeminiOAuthAccounts {
classify(src.ID, src.Kind, src.Name, PlatformGemini, AccountTypeOAuth)
}
for _, src := range exported.Data.GeminiAPIKeyAccounts {
classify(src.ID, src.Kind, src.Name, PlatformGemini, AccountTypeAPIKey)
}
return result, nil
}

View File

@@ -0,0 +1,69 @@
package service
import (
"strconv"
"strings"
"time"
gocache "github.com/patrickmn/go-cache"
)
// digestSessionTTL 摘要会话默认 TTL
const digestSessionTTL = 5 * time.Minute
// sessionEntry flat cache 条目
type sessionEntry struct {
uuid string
accountID int64
}
// DigestSessionStore 内存摘要会话存储flat cache 实现)
// key: "{groupID}:{prefixHash}|{digestChain}" → *sessionEntry
type DigestSessionStore struct {
cache *gocache.Cache
}
// NewDigestSessionStore 创建内存摘要会话存储
func NewDigestSessionStore() *DigestSessionStore {
return &DigestSessionStore{
cache: gocache.New(digestSessionTTL, time.Minute),
}
}
// Save 保存摘要会话。oldDigestChain 为 Find 返回的 matchedChain用于删旧 key。
func (s *DigestSessionStore) Save(groupID int64, prefixHash, digestChain, uuid string, accountID int64, oldDigestChain string) {
if digestChain == "" {
return
}
ns := buildNS(groupID, prefixHash)
s.cache.Set(ns+digestChain, &sessionEntry{uuid: uuid, accountID: accountID}, gocache.DefaultExpiration)
if oldDigestChain != "" && oldDigestChain != digestChain {
s.cache.Delete(ns + oldDigestChain)
}
}
// Find 查找摘要会话,从完整 chain 逐段截断,返回最长匹配及对应 matchedChain。
func (s *DigestSessionStore) Find(groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, matchedChain string, found bool) {
if digestChain == "" {
return "", 0, "", false
}
ns := buildNS(groupID, prefixHash)
chain := digestChain
for {
if val, ok := s.cache.Get(ns + chain); ok {
if e, ok := val.(*sessionEntry); ok {
return e.uuid, e.accountID, chain, true
}
}
i := strings.LastIndex(chain, "-")
if i < 0 {
return "", 0, "", false
}
chain = chain[:i]
}
}
// buildNS 构建 namespace 前缀
func buildNS(groupID int64, prefixHash string) string {
return strconv.FormatInt(groupID, 10) + ":" + prefixHash + "|"
}

View File

@@ -0,0 +1,312 @@
//go:build unit
package service
import (
"fmt"
"sync"
"testing"
"time"
gocache "github.com/patrickmn/go-cache"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestDigestSessionStore_SaveAndFind(t *testing.T) {
store := NewDigestSessionStore()
store.Save(1, "prefix", "s:a1-u:b2-m:c3", "uuid-1", 100, "")
uuid, accountID, _, found := store.Find(1, "prefix", "s:a1-u:b2-m:c3")
require.True(t, found)
assert.Equal(t, "uuid-1", uuid)
assert.Equal(t, int64(100), accountID)
}
func TestDigestSessionStore_PrefixMatch(t *testing.T) {
store := NewDigestSessionStore()
// 保存短链
store.Save(1, "prefix", "u:a-m:b", "uuid-short", 10, "")
// 用长链查找,应前缀匹配到短链
uuid, accountID, matchedChain, found := store.Find(1, "prefix", "u:a-m:b-u:c-m:d")
require.True(t, found)
assert.Equal(t, "uuid-short", uuid)
assert.Equal(t, int64(10), accountID)
assert.Equal(t, "u:a-m:b", matchedChain)
}
func TestDigestSessionStore_LongestPrefixMatch(t *testing.T) {
store := NewDigestSessionStore()
store.Save(1, "prefix", "u:a", "uuid-1", 1, "")
store.Save(1, "prefix", "u:a-m:b", "uuid-2", 2, "")
store.Save(1, "prefix", "u:a-m:b-u:c", "uuid-3", 3, "")
// 应匹配最深的 "u:a-m:b-u:c"(从完整 chain 逐段截断,先命中最长的)
uuid, accountID, _, found := store.Find(1, "prefix", "u:a-m:b-u:c-m:d-u:e")
require.True(t, found)
assert.Equal(t, "uuid-3", uuid)
assert.Equal(t, int64(3), accountID)
// 查找中等长度,应匹配到 "u:a-m:b"
uuid, accountID, _, found = store.Find(1, "prefix", "u:a-m:b-u:x")
require.True(t, found)
assert.Equal(t, "uuid-2", uuid)
assert.Equal(t, int64(2), accountID)
}
func TestDigestSessionStore_SaveDeletesOldChain(t *testing.T) {
store := NewDigestSessionStore()
// 第一轮:保存 "u:a-m:b"
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
// 第二轮:同一 uuid 保存更长的链,传入旧 chain
store.Save(1, "prefix", "u:a-m:b-u:c-m:d", "uuid-1", 100, "u:a-m:b")
// 旧链 "u:a-m:b" 应已被删除
_, _, _, found := store.Find(1, "prefix", "u:a-m:b")
assert.False(t, found, "old chain should be deleted")
// 新链应能找到
uuid, accountID, _, found := store.Find(1, "prefix", "u:a-m:b-u:c-m:d")
require.True(t, found)
assert.Equal(t, "uuid-1", uuid)
assert.Equal(t, int64(100), accountID)
}
func TestDigestSessionStore_DifferentSessionsNoInterference(t *testing.T) {
store := NewDigestSessionStore()
// 相同系统提示词,不同用户提示词
store.Save(1, "prefix", "s:sys-u:user1", "uuid-1", 100, "")
store.Save(1, "prefix", "s:sys-u:user2", "uuid-2", 200, "")
uuid, accountID, _, found := store.Find(1, "prefix", "s:sys-u:user1-m:reply1")
require.True(t, found)
assert.Equal(t, "uuid-1", uuid)
assert.Equal(t, int64(100), accountID)
uuid, accountID, _, found = store.Find(1, "prefix", "s:sys-u:user2-m:reply2")
require.True(t, found)
assert.Equal(t, "uuid-2", uuid)
assert.Equal(t, int64(200), accountID)
}
func TestDigestSessionStore_NoMatch(t *testing.T) {
store := NewDigestSessionStore()
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
// 完全不同的 chain
_, _, _, found := store.Find(1, "prefix", "u:x-m:y")
assert.False(t, found)
}
func TestDigestSessionStore_DifferentPrefixHash(t *testing.T) {
store := NewDigestSessionStore()
store.Save(1, "prefix1", "u:a-m:b", "uuid-1", 100, "")
// 不同 prefixHash 应隔离
_, _, _, found := store.Find(1, "prefix2", "u:a-m:b")
assert.False(t, found)
}
func TestDigestSessionStore_DifferentGroupID(t *testing.T) {
store := NewDigestSessionStore()
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
// 不同 groupID 应隔离
_, _, _, found := store.Find(2, "prefix", "u:a-m:b")
assert.False(t, found)
}
func TestDigestSessionStore_EmptyDigestChain(t *testing.T) {
store := NewDigestSessionStore()
// 空链不应保存
store.Save(1, "prefix", "", "uuid-1", 100, "")
_, _, _, found := store.Find(1, "prefix", "")
assert.False(t, found)
}
func TestDigestSessionStore_TTLExpiration(t *testing.T) {
store := &DigestSessionStore{
cache: gocache.New(100*time.Millisecond, 50*time.Millisecond),
}
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
// 立即应该能找到
_, _, _, found := store.Find(1, "prefix", "u:a-m:b")
require.True(t, found)
// 等待过期 + 清理周期
time.Sleep(300 * time.Millisecond)
// 过期后应找不到
_, _, _, found = store.Find(1, "prefix", "u:a-m:b")
assert.False(t, found)
}
func TestDigestSessionStore_ConcurrentSafety(t *testing.T) {
store := NewDigestSessionStore()
var wg sync.WaitGroup
const goroutines = 50
const operations = 100
wg.Add(goroutines)
for g := 0; g < goroutines; g++ {
go func(id int) {
defer wg.Done()
prefix := fmt.Sprintf("prefix-%d", id%5)
for i := 0; i < operations; i++ {
chain := fmt.Sprintf("u:%d-m:%d", id, i)
uuid := fmt.Sprintf("uuid-%d-%d", id, i)
store.Save(1, prefix, chain, uuid, int64(id), "")
store.Find(1, prefix, chain)
}
}(g)
}
wg.Wait()
}
func TestDigestSessionStore_MultipleSessions(t *testing.T) {
store := NewDigestSessionStore()
sessions := []struct {
chain string
uuid string
accountID int64
}{
{"u:session1", "uuid-1", 1},
{"u:session2-m:reply2", "uuid-2", 2},
{"u:session3-m:reply3-u:msg3", "uuid-3", 3},
}
for _, sess := range sessions {
store.Save(1, "prefix", sess.chain, sess.uuid, sess.accountID, "")
}
// 验证每个会话都能正确查找
for _, sess := range sessions {
uuid, accountID, _, found := store.Find(1, "prefix", sess.chain)
require.True(t, found, "should find session: %s", sess.chain)
assert.Equal(t, sess.uuid, uuid)
assert.Equal(t, sess.accountID, accountID)
}
// 验证继续对话的场景
uuid, accountID, _, found := store.Find(1, "prefix", "u:session2-m:reply2-u:newmsg")
require.True(t, found)
assert.Equal(t, "uuid-2", uuid)
assert.Equal(t, int64(2), accountID)
}
func TestDigestSessionStore_Performance1000Sessions(t *testing.T) {
store := NewDigestSessionStore()
// 插入 1000 个会话
for i := 0; i < 1000; i++ {
chain := fmt.Sprintf("s:sys-u:user%d-m:reply%d", i, i)
store.Save(1, "prefix", chain, fmt.Sprintf("uuid-%d", i), int64(i), "")
}
// 查找性能测试
start := time.Now()
const lookups = 10000
for i := 0; i < lookups; i++ {
idx := i % 1000
chain := fmt.Sprintf("s:sys-u:user%d-m:reply%d-u:newmsg", idx, idx)
_, _, _, found := store.Find(1, "prefix", chain)
assert.True(t, found)
}
elapsed := time.Since(start)
t.Logf("%d lookups in %v (%.0f ns/op)", lookups, elapsed, float64(elapsed.Nanoseconds())/lookups)
}
func TestDigestSessionStore_FindReturnsMatchedChain(t *testing.T) {
store := NewDigestSessionStore()
store.Save(1, "prefix", "u:a-m:b-u:c", "uuid-1", 100, "")
// 精确匹配
_, _, matchedChain, found := store.Find(1, "prefix", "u:a-m:b-u:c")
require.True(t, found)
assert.Equal(t, "u:a-m:b-u:c", matchedChain)
// 前缀匹配(截断后命中)
_, _, matchedChain, found = store.Find(1, "prefix", "u:a-m:b-u:c-m:d-u:e")
require.True(t, found)
assert.Equal(t, "u:a-m:b-u:c", matchedChain)
}
func TestDigestSessionStore_CacheItemCountStable(t *testing.T) {
store := NewDigestSessionStore()
// 模拟 100 个独立会话,每个进行 10 轮对话
// 正确传递 oldDigestChain 时,每个会话始终只保留 1 个 key
for conv := 0; conv < 100; conv++ {
var prevMatchedChain string
for round := 0; round < 10; round++ {
chain := fmt.Sprintf("s:sys-u:user%d", conv)
for r := 0; r < round; r++ {
chain += fmt.Sprintf("-m:a%d-u:q%d", r, r+1)
}
uuid := fmt.Sprintf("uuid-conv%d", conv)
_, _, matched, _ := store.Find(1, "prefix", chain)
store.Save(1, "prefix", chain, uuid, int64(conv), matched)
prevMatchedChain = matched
_ = prevMatchedChain
}
}
// 100 个会话 × 1 key/会话 = 应该 ≤ 100 个 key
// 允许少量并发残留,但绝不能接近 100×10=1000
itemCount := store.cache.ItemCount()
assert.LessOrEqual(t, itemCount, 100, "cache should have at most 100 items (1 per conversation), got %d", itemCount)
t.Logf("Cache item count after 100 conversations × 10 rounds: %d", itemCount)
}
func TestDigestSessionStore_TTLPreventsUnboundedGrowth(t *testing.T) {
// 使用极短 TTL 验证大量写入后 cache 能被清理
store := &DigestSessionStore{
cache: gocache.New(100*time.Millisecond, 50*time.Millisecond),
}
// 插入 500 个不同的 key无 oldDigestChain模拟最坏场景全是新会话首轮
for i := 0; i < 500; i++ {
chain := fmt.Sprintf("u:user%d", i)
store.Save(1, "prefix", chain, fmt.Sprintf("uuid-%d", i), int64(i), "")
}
assert.Equal(t, 500, store.cache.ItemCount())
// 等待 TTL + 清理周期
time.Sleep(300 * time.Millisecond)
assert.Equal(t, 0, store.cache.ItemCount(), "all items should be expired and cleaned up")
}
func TestDigestSessionStore_SaveSameChainNoDelete(t *testing.T) {
store := NewDigestSessionStore()
// 保存 chain
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
// 用户重发相同消息oldDigestChain == digestChain不应删掉刚设置的 key
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "u:a-m:b")
// 仍然能找到
uuid, accountID, _, found := store.Find(1, "prefix", "u:a-m:b")
require.True(t, found)
assert.Equal(t, "uuid-1", uuid)
assert.Equal(t, int64(100), accountID)
}

View File

@@ -0,0 +1,366 @@
//go:build unit
package service
import (
"context"
"io"
"net/http"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// Mocks (scoped to this file by naming convention)
// ---------------------------------------------------------------------------
// epFixedUpstream returns a fixed response for every request.
type epFixedUpstream struct {
statusCode int
body string
calls int
}
func (u *epFixedUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
u.calls++
return &http.Response{
StatusCode: u.statusCode,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(u.body)),
}, nil
}
func (u *epFixedUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, enableTLSFingerprint bool) (*http.Response, error) {
return u.Do(req, proxyURL, accountID, accountConcurrency)
}
// epAccountRepo records SetTempUnschedulable / SetError calls.
type epAccountRepo struct {
mockAccountRepoForGemini
tempCalls int
setErrCalls int
}
func (r *epAccountRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, _ string) error {
r.tempCalls++
return nil
}
func (r *epAccountRepo) SetError(_ context.Context, _ int64, _ string) error {
r.setErrCalls++
return nil
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
func saveAndSetBaseURLs(t *testing.T) {
t.Helper()
oldBaseURLs := append([]string(nil), antigravity.BaseURLs...)
oldAvail := antigravity.DefaultURLAvailability
antigravity.BaseURLs = []string{"https://ep-test.example"}
antigravity.DefaultURLAvailability = antigravity.NewURLAvailability(time.Minute)
t.Cleanup(func() {
antigravity.BaseURLs = oldBaseURLs
antigravity.DefaultURLAvailability = oldAvail
})
}
func newRetryParams(account *Account, upstream HTTPUpstream, handleError func(context.Context, string, *Account, int, http.Header, []byte, string, int64, string, bool) *handleModelRateLimitResult) antigravityRetryLoopParams {
return antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[ep-test]",
account: account,
accessToken: "token",
action: "generateContent",
body: []byte(`{"input":"test"}`),
httpUpstream: upstream,
requestedModel: "claude-sonnet-4-5",
handleError: handleError,
}
}
// ---------------------------------------------------------------------------
// TestRetryLoop_ErrorPolicy_CustomErrorCodes
// ---------------------------------------------------------------------------
func TestRetryLoop_ErrorPolicy_CustomErrorCodes(t *testing.T) {
tests := []struct {
name string
upstreamStatus int
upstreamBody string
customCodes []any
expectHandleError int
expectUpstream int
expectStatusCode int
}{
{
name: "429_in_custom_codes_matched",
upstreamStatus: 429,
upstreamBody: `{"error":"rate limited"}`,
customCodes: []any{float64(429)},
expectHandleError: 1,
expectUpstream: 1,
expectStatusCode: 429,
},
{
name: "429_not_in_custom_codes_skipped",
upstreamStatus: 429,
upstreamBody: `{"error":"rate limited"}`,
customCodes: []any{float64(500)},
expectHandleError: 0,
expectUpstream: 1,
expectStatusCode: 429,
},
{
name: "500_in_custom_codes_matched",
upstreamStatus: 500,
upstreamBody: `{"error":"internal"}`,
customCodes: []any{float64(500)},
expectHandleError: 1,
expectUpstream: 1,
expectStatusCode: 500,
},
{
name: "500_not_in_custom_codes_skipped",
upstreamStatus: 500,
upstreamBody: `{"error":"internal"}`,
customCodes: []any{float64(429)},
expectHandleError: 0,
expectUpstream: 1,
expectStatusCode: 500,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
saveAndSetBaseURLs(t)
upstream := &epFixedUpstream{statusCode: tt.upstreamStatus, body: tt.upstreamBody}
repo := &epAccountRepo{}
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
account := &Account{
ID: 100,
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Schedulable: true,
Status: StatusActive,
Concurrency: 1,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": tt.customCodes,
},
}
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
var handleErrorCount int
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
handleErrorCount++
return nil
})
result, err := svc.antigravityRetryLoop(p)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.resp)
defer func() { _ = result.resp.Body.Close() }()
require.Equal(t, tt.expectStatusCode, result.resp.StatusCode)
require.Equal(t, tt.expectHandleError, handleErrorCount, "handleError call count")
require.Equal(t, tt.expectUpstream, upstream.calls, "upstream call count")
})
}
}
// ---------------------------------------------------------------------------
// TestRetryLoop_ErrorPolicy_TempUnschedulable
// ---------------------------------------------------------------------------
func TestRetryLoop_ErrorPolicy_TempUnschedulable(t *testing.T) {
tempRulesAccount := func(rules []any) *Account {
return &Account{
ID: 200,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Schedulable: true,
Status: StatusActive,
Concurrency: 1,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": rules,
},
}
}
overloadedRule := map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded"},
"duration_minutes": float64(10),
}
rateLimitRule := map[string]any{
"error_code": float64(429),
"keywords": []any{"rate limited keyword"},
"duration_minutes": float64(5),
}
t.Run("503_overloaded_matches_rule", func(t *testing.T) {
saveAndSetBaseURLs(t)
upstream := &epFixedUpstream{statusCode: 503, body: `overloaded`}
repo := &epAccountRepo{}
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
account := tempRulesAccount([]any{overloadedRule})
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
t.Error("handleError should not be called for temp unschedulable")
return nil
})
result, err := svc.antigravityRetryLoop(p)
require.Nil(t, result)
var switchErr *AntigravityAccountSwitchError
require.ErrorAs(t, err, &switchErr)
require.Equal(t, account.ID, switchErr.OriginalAccountID)
require.Equal(t, 1, upstream.calls, "should not retry")
})
t.Run("429_rate_limited_keyword_matches_rule", func(t *testing.T) {
saveAndSetBaseURLs(t)
upstream := &epFixedUpstream{statusCode: 429, body: `rate limited keyword`}
repo := &epAccountRepo{}
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
account := tempRulesAccount([]any{rateLimitRule})
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
t.Error("handleError should not be called for temp unschedulable")
return nil
})
result, err := svc.antigravityRetryLoop(p)
require.Nil(t, result)
var switchErr *AntigravityAccountSwitchError
require.ErrorAs(t, err, &switchErr)
require.Equal(t, account.ID, switchErr.OriginalAccountID)
require.Equal(t, 1, upstream.calls, "should not retry")
})
t.Run("503_body_no_match_continues_default_retry", func(t *testing.T) {
saveAndSetBaseURLs(t)
upstream := &epFixedUpstream{statusCode: 503, body: `random`}
repo := &epAccountRepo{}
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
account := tempRulesAccount([]any{overloadedRule})
// Use a short-lived context: the backoff sleep (~1s) will be
// interrupted, proving the code entered the default retry path
// instead of breaking early via error policy.
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
return nil
})
p.ctx = ctx
result, err := svc.antigravityRetryLoop(p)
// Context cancellation during backoff proves default retry was entered
require.Nil(t, result)
require.ErrorIs(t, err, context.DeadlineExceeded)
require.GreaterOrEqual(t, upstream.calls, 1, "should have called upstream at least once")
})
}
// ---------------------------------------------------------------------------
// TestRetryLoop_ErrorPolicy_NilRateLimitService
// ---------------------------------------------------------------------------
func TestRetryLoop_ErrorPolicy_NilRateLimitService(t *testing.T) {
saveAndSetBaseURLs(t)
upstream := &epFixedUpstream{statusCode: 429, body: `{"error":"rate limited"}`}
// rateLimitService is nil — must not panic
svc := &AntigravityGatewayService{rateLimitService: nil}
account := &Account{
ID: 300,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Schedulable: true,
Status: StatusActive,
Concurrency: 1,
}
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
return nil
})
p.ctx = ctx
// Should not panic; enters the default retry path (eventually times out)
result, err := svc.antigravityRetryLoop(p)
require.Nil(t, result)
require.ErrorIs(t, err, context.DeadlineExceeded)
require.GreaterOrEqual(t, upstream.calls, 1)
}
// ---------------------------------------------------------------------------
// TestRetryLoop_ErrorPolicy_NoPolicy_OriginalBehavior
// ---------------------------------------------------------------------------
func TestRetryLoop_ErrorPolicy_NoPolicy_OriginalBehavior(t *testing.T) {
saveAndSetBaseURLs(t)
upstream := &epFixedUpstream{statusCode: 429, body: `{"error":"rate limited"}`}
repo := &epAccountRepo{}
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
// Plain OAuth account with no error policy configured
account := &Account{
ID: 400,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Schedulable: true,
Status: StatusActive,
Concurrency: 1,
}
var handleErrorCount int
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
handleErrorCount++
return nil
})
result, err := svc.antigravityRetryLoop(p)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.resp)
defer func() { _ = result.resp.Body.Close() }()
require.Equal(t, http.StatusTooManyRequests, result.resp.StatusCode)
require.Equal(t, antigravityMaxRetries, upstream.calls, "should exhaust all retries")
require.Equal(t, 1, handleErrorCount, "handleError should be called once after retries exhausted")
}

View File

@@ -0,0 +1,289 @@
//go:build unit
package service
import (
"context"
"net/http"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// TestCheckErrorPolicy — 6 table-driven cases for the pure logic function
// ---------------------------------------------------------------------------
func TestCheckErrorPolicy(t *testing.T) {
tests := []struct {
name string
account *Account
statusCode int
body []byte
expected ErrorPolicyResult
}{
{
name: "no_policy_oauth_returns_none",
account: &Account{
ID: 1,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
// no custom error codes, no temp rules
},
statusCode: 500,
body: []byte(`"error"`),
expected: ErrorPolicyNone,
},
{
name: "custom_error_codes_hit_returns_matched",
account: &Account{
ID: 2,
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(429), float64(500)},
},
},
statusCode: 500,
body: []byte(`"error"`),
expected: ErrorPolicyMatched,
},
{
name: "custom_error_codes_miss_returns_skipped",
account: &Account{
ID: 3,
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(429), float64(500)},
},
},
statusCode: 503,
body: []byte(`"error"`),
expected: ErrorPolicySkipped,
},
{
name: "temp_unschedulable_hit_returns_temp_unscheduled",
account: &Account{
ID: 4,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded"},
"duration_minutes": float64(10),
"description": "overloaded rule",
},
},
},
},
statusCode: 503,
body: []byte(`overloaded service`),
expected: ErrorPolicyTempUnscheduled,
},
{
name: "temp_unschedulable_body_miss_returns_none",
account: &Account{
ID: 5,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded"},
"duration_minutes": float64(10),
"description": "overloaded rule",
},
},
},
},
statusCode: 503,
body: []byte(`random msg`),
expected: ErrorPolicyNone,
},
{
name: "custom_error_codes_override_temp_unschedulable",
account: &Account{
ID: 6,
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(503)},
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded"},
"duration_minutes": float64(10),
"description": "overloaded rule",
},
},
},
},
statusCode: 503,
body: []byte(`overloaded`),
expected: ErrorPolicyMatched, // custom codes take precedence
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &errorPolicyRepoStub{}
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
result := svc.CheckErrorPolicy(context.Background(), tt.account, tt.statusCode, tt.body)
require.Equal(t, tt.expected, result, "unexpected ErrorPolicyResult")
})
}
}
// ---------------------------------------------------------------------------
// TestApplyErrorPolicy — 4 table-driven cases for the wrapper method
// ---------------------------------------------------------------------------
func TestApplyErrorPolicy(t *testing.T) {
tests := []struct {
name string
account *Account
statusCode int
body []byte
expectedHandled bool
expectedSwitchErr bool // expect *AntigravityAccountSwitchError
handleErrorCalls int
}{
{
name: "none_not_handled",
account: &Account{
ID: 10,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
},
statusCode: 500,
body: []byte(`"error"`),
expectedHandled: false,
handleErrorCalls: 0,
},
{
name: "skipped_handled_no_handleError",
account: &Account{
ID: 11,
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(429)},
},
},
statusCode: 500, // not in custom codes
body: []byte(`"error"`),
expectedHandled: true,
handleErrorCalls: 0,
},
{
name: "matched_handled_calls_handleError",
account: &Account{
ID: 12,
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(500)},
},
},
statusCode: 500,
body: []byte(`"error"`),
expectedHandled: true,
handleErrorCalls: 1,
},
{
name: "temp_unscheduled_returns_switch_error",
account: &Account{
ID: 13,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded"},
"duration_minutes": float64(10),
},
},
},
},
statusCode: 503,
body: []byte(`overloaded`),
expectedHandled: true,
expectedSwitchErr: true,
handleErrorCalls: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &errorPolicyRepoStub{}
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
svc := &AntigravityGatewayService{
rateLimitService: rlSvc,
}
var handleErrorCount int
p := antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: tt.account,
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
handleErrorCount++
return nil
},
isStickySession: true,
}
handled, retErr := svc.applyErrorPolicy(p, tt.statusCode, http.Header{}, tt.body)
require.Equal(t, tt.expectedHandled, handled, "handled mismatch")
require.Equal(t, tt.handleErrorCalls, handleErrorCount, "handleError call count mismatch")
if tt.expectedSwitchErr {
var switchErr *AntigravityAccountSwitchError
require.ErrorAs(t, retErr, &switchErr)
require.Equal(t, tt.account.ID, switchErr.OriginalAccountID)
} else {
require.NoError(t, retErr)
}
})
}
}
// ---------------------------------------------------------------------------
// errorPolicyRepoStub — minimal AccountRepository stub for error policy tests
// ---------------------------------------------------------------------------
type errorPolicyRepoStub struct {
mockAccountRepoForGemini
tempCalls int
setErrCalls int
lastErrorMsg string
}
func (r *errorPolicyRepoStub) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
r.tempCalls++
return nil
}
func (r *errorPolicyRepoStub) SetError(ctx context.Context, id int64, errorMsg string) error {
r.setErrCalls++
r.lastErrorMsg = errorMsg
return nil
}

View File

@@ -77,6 +77,9 @@ func (m *mockAccountRepoForPlatform) Create(ctx context.Context, account *Accoun
func (m *mockAccountRepoForPlatform) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) { func (m *mockAccountRepoForPlatform) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) {
return nil, nil return nil, nil
} }
func (m *mockAccountRepoForPlatform) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
return nil, nil
}
func (m *mockAccountRepoForPlatform) Update(ctx context.Context, account *Account) error { func (m *mockAccountRepoForPlatform) Update(ctx context.Context, account *Account) error {
return nil return nil
} }
@@ -142,9 +145,6 @@ func (m *mockAccountRepoForPlatform) ListSchedulableByGroupIDAndPlatforms(ctx co
func (m *mockAccountRepoForPlatform) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { func (m *mockAccountRepoForPlatform) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
return nil return nil
} }
func (m *mockAccountRepoForPlatform) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error {
return nil
}
func (m *mockAccountRepoForPlatform) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error { func (m *mockAccountRepoForPlatform) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error {
return nil return nil
} }
@@ -216,22 +216,6 @@ func (m *mockGatewayCacheForPlatform) DeleteSessionAccountID(ctx context.Context
return nil return nil
} }
func (m *mockGatewayCacheForPlatform) IncrModelCallCount(ctx context.Context, accountID int64, model string) (int64, error) {
return 0, nil
}
func (m *mockGatewayCacheForPlatform) GetModelLoadBatch(ctx context.Context, accountIDs []int64, model string) (map[int64]*ModelLoadInfo, error) {
return nil, nil
}
func (m *mockGatewayCacheForPlatform) FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
return "", 0, false
}
func (m *mockGatewayCacheForPlatform) SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
return nil
}
type mockGroupRepoForGateway struct { type mockGroupRepoForGateway struct {
groups map[int64]*Group groups map[int64]*Group
getByIDCalls int getByIDCalls int
@@ -290,6 +274,10 @@ func (m *mockGroupRepoForGateway) GetAccountIDsByGroupIDs(ctx context.Context, g
return nil, nil return nil, nil
} }
func (m *mockGroupRepoForGateway) UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
return nil
}
func ptr[T any](v T) *T { func ptr[T any](v T) *T {
return &v return &v
} }

View File

@@ -6,9 +6,19 @@ import (
"fmt" "fmt"
"math" "math"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
) )
// SessionContext 粘性会话上下文,用于区分不同来源的请求。
// 仅在 GenerateSessionHash 第 3 级 fallback消息内容 hash时混入
// 避免不同用户发送相同消息产生相同 hash 导致账号集中。
type SessionContext struct {
ClientIP string
UserAgent string
APIKeyID int64
}
// ParsedRequest 保存网关请求的预解析结果 // ParsedRequest 保存网关请求的预解析结果
// //
// 性能优化说明: // 性能优化说明:
@@ -22,20 +32,22 @@ import (
// 2. 将解析结果 ParsedRequest 传递给 Service 层 // 2. 将解析结果 ParsedRequest 传递给 Service 层
// 3. 避免重复 json.Unmarshal减少 CPU 和内存开销 // 3. 避免重复 json.Unmarshal减少 CPU 和内存开销
type ParsedRequest struct { type ParsedRequest struct {
Body []byte // 原始请求体(保留用于转发) Body []byte // 原始请求体(保留用于转发)
Model string // 请求的模型名称 Model string // 请求的模型名称
Stream bool // 是否为流式请求 Stream bool // 是否为流式请求
MetadataUserID string // metadata.user_id用于会话亲和 MetadataUserID string // metadata.user_id用于会话亲和
System any // system 字段内容 System any // system 字段内容
Messages []any // messages 数组 Messages []any // messages 数组
HasSystem bool // 是否包含 system 字段(包含 null 也视为显式传入) HasSystem bool // 是否包含 system 字段(包含 null 也视为显式传入)
ThinkingEnabled bool // 是否开启 thinking部分平台会影响最终模型名 ThinkingEnabled bool // 是否开启 thinking部分平台会影响最终模型名
MaxTokens int // max_tokens 值(用于探测请求拦截) MaxTokens int // max_tokens 值(用于探测请求拦截)
SessionContext *SessionContext // 可选请求上下文区分因子nil 时行为不变)
} }
// ParseGatewayRequest 解析网关请求体并返回结构化结果 // ParseGatewayRequest 解析网关请求体并返回结构化结果
// 性能优化:一次解析提取所有需要的字段,避免重复 Unmarshal // protocol 指定请求协议格式domain.PlatformAnthropic / domain.PlatformGemini
func ParseGatewayRequest(body []byte) (*ParsedRequest, error) { // 不同协议使用不同的 system/messages 字段名。
func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) {
var req map[string]any var req map[string]any
if err := json.Unmarshal(body, &req); err != nil { if err := json.Unmarshal(body, &req); err != nil {
return nil, err return nil, err
@@ -64,14 +76,29 @@ func ParseGatewayRequest(body []byte) (*ParsedRequest, error) {
parsed.MetadataUserID = userID parsed.MetadataUserID = userID
} }
} }
// system 字段只要存在就视为显式提供(即使为 null
// 以避免客户端传 null 时被默认 system 误注入。 switch protocol {
if system, ok := req["system"]; ok { case domain.PlatformGemini:
parsed.HasSystem = true // Gemini 原生格式: systemInstruction.parts / contents
parsed.System = system if sysInst, ok := req["systemInstruction"].(map[string]any); ok {
} if parts, ok := sysInst["parts"].([]any); ok {
if messages, ok := req["messages"].([]any); ok { parsed.System = parts
parsed.Messages = messages }
}
if contents, ok := req["contents"].([]any); ok {
parsed.Messages = contents
}
default:
// Anthropic / OpenAI 格式: system / messages
// system 字段只要存在就视为显式提供(即使为 null
// 以避免客户端传 null 时被默认 system 误注入。
if system, ok := req["system"]; ok {
parsed.HasSystem = true
parsed.System = system
}
if messages, ok := req["messages"].([]any); ok {
parsed.Messages = messages
}
} }
// thinking: {type: "enabled"} // thinking: {type: "enabled"}

View File

@@ -4,12 +4,13 @@ import (
"encoding/json" "encoding/json"
"testing" "testing"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
func TestParseGatewayRequest(t *testing.T) { func TestParseGatewayRequest(t *testing.T) {
body := []byte(`{"model":"claude-3-7-sonnet","stream":true,"metadata":{"user_id":"session_123e4567-e89b-12d3-a456-426614174000"},"system":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}],"messages":[{"content":"hi"}]}`) body := []byte(`{"model":"claude-3-7-sonnet","stream":true,"metadata":{"user_id":"session_123e4567-e89b-12d3-a456-426614174000"},"system":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}],"messages":[{"content":"hi"}]}`)
parsed, err := ParseGatewayRequest(body) parsed, err := ParseGatewayRequest(body, "")
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, "claude-3-7-sonnet", parsed.Model) require.Equal(t, "claude-3-7-sonnet", parsed.Model)
require.True(t, parsed.Stream) require.True(t, parsed.Stream)
@@ -22,7 +23,7 @@ func TestParseGatewayRequest(t *testing.T) {
func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) { func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) {
body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"enabled"},"messages":[{"content":"hi"}]}`) body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"enabled"},"messages":[{"content":"hi"}]}`)
parsed, err := ParseGatewayRequest(body) parsed, err := ParseGatewayRequest(body, "")
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, "claude-sonnet-4-5", parsed.Model) require.Equal(t, "claude-sonnet-4-5", parsed.Model)
require.True(t, parsed.ThinkingEnabled) require.True(t, parsed.ThinkingEnabled)
@@ -30,21 +31,21 @@ func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) {
func TestParseGatewayRequest_MaxTokens(t *testing.T) { func TestParseGatewayRequest_MaxTokens(t *testing.T) {
body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1}`) body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1}`)
parsed, err := ParseGatewayRequest(body) parsed, err := ParseGatewayRequest(body, "")
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, 1, parsed.MaxTokens) require.Equal(t, 1, parsed.MaxTokens)
} }
func TestParseGatewayRequest_MaxTokensNonIntegralIgnored(t *testing.T) { func TestParseGatewayRequest_MaxTokensNonIntegralIgnored(t *testing.T) {
body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1.5}`) body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1.5}`)
parsed, err := ParseGatewayRequest(body) parsed, err := ParseGatewayRequest(body, "")
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, 0, parsed.MaxTokens) require.Equal(t, 0, parsed.MaxTokens)
} }
func TestParseGatewayRequest_SystemNull(t *testing.T) { func TestParseGatewayRequest_SystemNull(t *testing.T) {
body := []byte(`{"model":"claude-3","system":null}`) body := []byte(`{"model":"claude-3","system":null}`)
parsed, err := ParseGatewayRequest(body) parsed, err := ParseGatewayRequest(body, "")
require.NoError(t, err) require.NoError(t, err)
// 显式传入 system:null 也应视为“字段已存在”,避免默认 system 被注入。 // 显式传入 system:null 也应视为“字段已存在”,避免默认 system 被注入。
require.True(t, parsed.HasSystem) require.True(t, parsed.HasSystem)
@@ -53,16 +54,112 @@ func TestParseGatewayRequest_SystemNull(t *testing.T) {
func TestParseGatewayRequest_InvalidModelType(t *testing.T) { func TestParseGatewayRequest_InvalidModelType(t *testing.T) {
body := []byte(`{"model":123}`) body := []byte(`{"model":123}`)
_, err := ParseGatewayRequest(body) _, err := ParseGatewayRequest(body, "")
require.Error(t, err) require.Error(t, err)
} }
func TestParseGatewayRequest_InvalidStreamType(t *testing.T) { func TestParseGatewayRequest_InvalidStreamType(t *testing.T) {
body := []byte(`{"stream":"true"}`) body := []byte(`{"stream":"true"}`)
_, err := ParseGatewayRequest(body) _, err := ParseGatewayRequest(body, "")
require.Error(t, err) require.Error(t, err)
} }
// ============ Gemini 原生格式解析测试 ============
func TestParseGatewayRequest_GeminiContents(t *testing.T) {
body := []byte(`{
"contents": [
{"role": "user", "parts": [{"text": "Hello"}]},
{"role": "model", "parts": [{"text": "Hi there"}]},
{"role": "user", "parts": [{"text": "How are you?"}]}
]
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
require.NoError(t, err)
require.Len(t, parsed.Messages, 3, "should parse contents as Messages")
require.False(t, parsed.HasSystem, "Gemini format should not set HasSystem")
require.Nil(t, parsed.System, "no systemInstruction means nil System")
}
func TestParseGatewayRequest_GeminiSystemInstruction(t *testing.T) {
body := []byte(`{
"systemInstruction": {
"parts": [{"text": "You are a helpful assistant."}]
},
"contents": [
{"role": "user", "parts": [{"text": "Hello"}]}
]
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
require.NoError(t, err)
require.NotNil(t, parsed.System, "should parse systemInstruction.parts as System")
parts, ok := parsed.System.([]any)
require.True(t, ok)
require.Len(t, parts, 1)
partMap, ok := parts[0].(map[string]any)
require.True(t, ok)
require.Equal(t, "You are a helpful assistant.", partMap["text"])
require.Len(t, parsed.Messages, 1)
}
func TestParseGatewayRequest_GeminiWithModel(t *testing.T) {
body := []byte(`{
"model": "gemini-2.5-pro",
"contents": [{"role": "user", "parts": [{"text": "test"}]}]
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
require.NoError(t, err)
require.Equal(t, "gemini-2.5-pro", parsed.Model)
require.Len(t, parsed.Messages, 1)
}
func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) {
// Gemini 格式下 system/messages 字段应被忽略
body := []byte(`{
"system": "should be ignored",
"messages": [{"role": "user", "content": "ignored"}],
"contents": [{"role": "user", "parts": [{"text": "real content"}]}]
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
require.NoError(t, err)
require.False(t, parsed.HasSystem, "Gemini protocol should not parse Anthropic system field")
require.Nil(t, parsed.System, "no systemInstruction = nil System")
require.Len(t, parsed.Messages, 1, "should use contents, not messages")
}
func TestParseGatewayRequest_GeminiEmptyContents(t *testing.T) {
body := []byte(`{"contents": []}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
require.NoError(t, err)
require.Empty(t, parsed.Messages)
}
func TestParseGatewayRequest_GeminiNoContents(t *testing.T) {
body := []byte(`{"model": "gemini-2.5-flash"}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
require.NoError(t, err)
require.Nil(t, parsed.Messages)
require.Equal(t, "gemini-2.5-flash", parsed.Model)
}
func TestParseGatewayRequest_AnthropicIgnoresGeminiFields(t *testing.T) {
// Anthropic 格式下 contents/systemInstruction 字段应被忽略
body := []byte(`{
"system": "real system",
"messages": [{"role": "user", "content": "real content"}],
"contents": [{"role": "user", "parts": [{"text": "ignored"}]}],
"systemInstruction": {"parts": [{"text": "ignored"}]}
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformAnthropic)
require.NoError(t, err)
require.True(t, parsed.HasSystem)
require.Equal(t, "real system", parsed.System)
require.Len(t, parsed.Messages, 1)
msg, ok := parsed.Messages[0].(map[string]any)
require.True(t, ok)
require.Equal(t, "real content", msg["content"])
}
func TestFilterThinkingBlocks(t *testing.T) { func TestFilterThinkingBlocks(t *testing.T) {
containsThinkingBlock := func(body []byte) bool { containsThinkingBlock := func(body []byte) bool {
var req map[string]any var req map[string]any

View File

@@ -5,7 +5,6 @@ import (
"bytes" "bytes"
"context" "context"
"crypto/sha256" "crypto/sha256"
"encoding/hex"
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
@@ -17,6 +16,7 @@ import (
"os" "os"
"regexp" "regexp"
"sort" "sort"
"strconv"
"strings" "strings"
"sync/atomic" "sync/atomic"
"time" "time"
@@ -26,6 +26,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
"github.com/cespare/xxhash/v2"
"github.com/google/uuid" "github.com/google/uuid"
"github.com/tidwall/gjson" "github.com/tidwall/gjson"
"github.com/tidwall/sjson" "github.com/tidwall/sjson"
@@ -245,9 +246,6 @@ var (
// ErrClaudeCodeOnly 表示分组仅允许 Claude Code 客户端访问 // ErrClaudeCodeOnly 表示分组仅允许 Claude Code 客户端访问
var ErrClaudeCodeOnly = errors.New("this group only allows Claude Code clients") var ErrClaudeCodeOnly = errors.New("this group only allows Claude Code clients")
// ErrModelScopeNotSupported 表示请求的模型系列不在分组支持的范围内
var ErrModelScopeNotSupported = errors.New("model scope not supported by this group")
// allowedHeaders 白名单headers参考CRS项目 // allowedHeaders 白名单headers参考CRS项目
var allowedHeaders = map[string]bool{ var allowedHeaders = map[string]bool{
"accept": true, "accept": true,
@@ -273,13 +271,6 @@ var allowedHeaders = map[string]bool{
// GatewayCache 定义网关服务的缓存操作接口。 // GatewayCache 定义网关服务的缓存操作接口。
// 提供粘性会话Sticky Session的存储、查询、刷新和删除功能。 // 提供粘性会话Sticky Session的存储、查询、刷新和删除功能。
// //
// ModelLoadInfo 模型负载信息(用于 Antigravity 调度)
// Model load info for Antigravity scheduling
type ModelLoadInfo struct {
CallCount int64 // 当前分钟调用次数 / Call count in current minute
LastUsedAt time.Time // 最后调度时间(零值表示未调度过)/ Last scheduling time (zero means never scheduled)
}
// GatewayCache defines cache operations for gateway service. // GatewayCache defines cache operations for gateway service.
// Provides sticky session storage, retrieval, refresh and deletion capabilities. // Provides sticky session storage, retrieval, refresh and deletion capabilities.
type GatewayCache interface { type GatewayCache interface {
@@ -295,24 +286,6 @@ type GatewayCache interface {
// DeleteSessionAccountID 删除粘性会话绑定,用于账号不可用时主动清理 // DeleteSessionAccountID 删除粘性会话绑定,用于账号不可用时主动清理
// Delete sticky session binding, used to proactively clean up when account becomes unavailable // Delete sticky session binding, used to proactively clean up when account becomes unavailable
DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error
// IncrModelCallCount 增加模型调用次数并更新最后调度时间Antigravity 专用)
// Increment model call count and update last scheduling time (Antigravity only)
// 返回更新后的调用次数
IncrModelCallCount(ctx context.Context, accountID int64, model string) (int64, error)
// GetModelLoadBatch 批量获取账号的模型负载信息Antigravity 专用)
// Batch get model load info for accounts (Antigravity only)
GetModelLoadBatch(ctx context.Context, accountIDs []int64, model string) (map[int64]*ModelLoadInfo, error)
// FindGeminiSession 查找 Gemini 会话MGET 倒序匹配)
// Find Gemini session using MGET reverse order matching
// 返回最长匹配的会话信息uuid, accountID
FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool)
// SaveGeminiSession 保存 Gemini 会话
// Save Gemini session binding
SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error
} }
// derefGroupID safely dereferences *int64 to int64, returning 0 if nil // derefGroupID safely dereferences *int64 to int64, returning 0 if nil
@@ -323,21 +296,15 @@ func derefGroupID(groupID *int64) int64 {
return *groupID return *groupID
} }
// stickySessionRateLimitThreshold 定义清除粘性会话的限流时间阈值。
// 当账号限流剩余时间超过此阈值时,清除粘性会话以便切换到其他账号。
// 低于此阈值时保持粘性会话,等待短暂限流结束。
const stickySessionRateLimitThreshold = 10 * time.Second
// shouldClearStickySession 检查账号是否处于不可调度状态,需要清理粘性会话绑定。 // shouldClearStickySession 检查账号是否处于不可调度状态,需要清理粘性会话绑定。
// 当账号状态为错误、禁用、不可调度、处于临时不可调度期间, // 当账号状态为错误、禁用、不可调度、处于临时不可调度期间,
// 或模型限流剩余时间超过 stickySessionRateLimitThreshold 时,返回 true。 // 或请求的模型处于限流状态时,返回 true。
// 这确保后续请求不会继续使用不可用的账号。 // 这确保后续请求不会继续使用不可用的账号。
// //
// shouldClearStickySession checks if an account is in an unschedulable state // shouldClearStickySession checks if an account is in an unschedulable state
// and the sticky session binding should be cleared. // and the sticky session binding should be cleared.
// Returns true when account status is error/disabled, schedulable is false, // Returns true when account status is error/disabled, schedulable is false,
// within temporary unschedulable period, or model rate limit remaining time // within temporary unschedulable period, or the requested model is rate-limited.
// exceeds stickySessionRateLimitThreshold.
// This ensures subsequent requests won't continue using unavailable accounts. // This ensures subsequent requests won't continue using unavailable accounts.
func shouldClearStickySession(account *Account, requestedModel string) bool { func shouldClearStickySession(account *Account, requestedModel string) bool {
if account == nil { if account == nil {
@@ -349,8 +316,8 @@ func shouldClearStickySession(account *Account, requestedModel string) bool {
if account.TempUnschedulableUntil != nil && time.Now().Before(*account.TempUnschedulableUntil) { if account.TempUnschedulableUntil != nil && time.Now().Before(*account.TempUnschedulableUntil) {
return true return true
} }
// 检查模型限流和 scope 限流,只在超过阈值时清除粘性会话 // 检查模型限流和 scope 限流,有限流即清除粘性会话
if remaining := account.GetRateLimitRemainingTimeWithContext(context.Background(), requestedModel); remaining > stickySessionRateLimitThreshold { if remaining := account.GetRateLimitRemainingTimeWithContext(context.Background(), requestedModel); remaining > 0 {
return true return true
} }
return false return false
@@ -413,6 +380,7 @@ type GatewayService struct {
userSubRepo UserSubscriptionRepository userSubRepo UserSubscriptionRepository
userGroupRateRepo UserGroupRateRepository userGroupRateRepo UserGroupRateRepository
cache GatewayCache cache GatewayCache
digestStore *DigestSessionStore
cfg *config.Config cfg *config.Config
schedulerSnapshot *SchedulerSnapshotService schedulerSnapshot *SchedulerSnapshotService
billingService *BillingService billingService *BillingService
@@ -446,6 +414,7 @@ func NewGatewayService(
deferredService *DeferredService, deferredService *DeferredService,
claudeTokenProvider *ClaudeTokenProvider, claudeTokenProvider *ClaudeTokenProvider,
sessionLimitCache SessionLimitCache, sessionLimitCache SessionLimitCache,
digestStore *DigestSessionStore,
) *GatewayService { ) *GatewayService {
return &GatewayService{ return &GatewayService{
accountRepo: accountRepo, accountRepo: accountRepo,
@@ -455,6 +424,7 @@ func NewGatewayService(
userSubRepo: userSubRepo, userSubRepo: userSubRepo,
userGroupRateRepo: userGroupRateRepo, userGroupRateRepo: userGroupRateRepo,
cache: cache, cache: cache,
digestStore: digestStore,
cfg: cfg, cfg: cfg,
schedulerSnapshot: schedulerSnapshot, schedulerSnapshot: schedulerSnapshot,
concurrencyService: concurrencyService, concurrencyService: concurrencyService,
@@ -488,23 +458,45 @@ func (s *GatewayService) GenerateSessionHash(parsed *ParsedRequest) string {
return s.hashContent(cacheableContent) return s.hashContent(cacheableContent)
} }
// 3. Fallback: 使用 system 内容 // 3. 最后 fallback: 使用 session上下文 + system + 所有消息的完整摘要串
var combined strings.Builder
// 混入请求上下文区分因子,避免不同用户相同消息产生相同 hash
if parsed.SessionContext != nil {
_, _ = combined.WriteString(parsed.SessionContext.ClientIP)
_, _ = combined.WriteString(":")
_, _ = combined.WriteString(parsed.SessionContext.UserAgent)
_, _ = combined.WriteString(":")
_, _ = combined.WriteString(strconv.FormatInt(parsed.SessionContext.APIKeyID, 10))
_, _ = combined.WriteString("|")
}
if parsed.System != nil { if parsed.System != nil {
systemText := s.extractTextFromSystem(parsed.System) systemText := s.extractTextFromSystem(parsed.System)
if systemText != "" { if systemText != "" {
return s.hashContent(systemText) _, _ = combined.WriteString(systemText)
} }
} }
for _, msg := range parsed.Messages {
// 4. 最后 fallback: 使用第一条消息 if m, ok := msg.(map[string]any); ok {
if len(parsed.Messages) > 0 { if content, exists := m["content"]; exists {
if firstMsg, ok := parsed.Messages[0].(map[string]any); ok { // Anthropic: messages[].content
msgText := s.extractTextFromContent(firstMsg["content"]) if msgText := s.extractTextFromContent(content); msgText != "" {
if msgText != "" { _, _ = combined.WriteString(msgText)
return s.hashContent(msgText) }
} else if parts, ok := m["parts"].([]any); ok {
// Gemini: contents[].parts[].text
for _, part := range parts {
if partMap, ok := part.(map[string]any); ok {
if text, ok := partMap["text"].(string); ok {
_, _ = combined.WriteString(text)
}
}
}
} }
} }
} }
if combined.Len() > 0 {
return s.hashContent(combined.String())
}
return "" return ""
} }
@@ -532,19 +524,37 @@ func (s *GatewayService) GetCachedSessionAccountID(ctx context.Context, groupID
// FindGeminiSession 查找 Gemini 会话(基于内容摘要链的 Fallback 匹配) // FindGeminiSession 查找 Gemini 会话(基于内容摘要链的 Fallback 匹配)
// 返回最长匹配的会话信息uuid, accountID // 返回最长匹配的会话信息uuid, accountID
func (s *GatewayService) FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) { func (s *GatewayService) FindGeminiSession(_ context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, matchedChain string, found bool) {
if digestChain == "" || s.cache == nil { if digestChain == "" || s.digestStore == nil {
return "", 0, false return "", 0, "", false
} }
return s.cache.FindGeminiSession(ctx, groupID, prefixHash, digestChain) return s.digestStore.Find(groupID, prefixHash, digestChain)
} }
// SaveGeminiSession 保存 Gemini 会话 // SaveGeminiSession 保存 Gemini 会话。oldDigestChain 为 Find 返回的 matchedChain用于删旧 key。
func (s *GatewayService) SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error { func (s *GatewayService) SaveGeminiSession(_ context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64, oldDigestChain string) error {
if digestChain == "" || s.cache == nil { if digestChain == "" || s.digestStore == nil {
return nil return nil
} }
return s.cache.SaveGeminiSession(ctx, groupID, prefixHash, digestChain, uuid, accountID) s.digestStore.Save(groupID, prefixHash, digestChain, uuid, accountID, oldDigestChain)
return nil
}
// FindAnthropicSession 查找 Anthropic 会话(基于内容摘要链的 Fallback 匹配)
func (s *GatewayService) FindAnthropicSession(_ context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, matchedChain string, found bool) {
if digestChain == "" || s.digestStore == nil {
return "", 0, "", false
}
return s.digestStore.Find(groupID, prefixHash, digestChain)
}
// SaveAnthropicSession 保存 Anthropic 会话
func (s *GatewayService) SaveAnthropicSession(_ context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64, oldDigestChain string) error {
if digestChain == "" || s.digestStore == nil {
return nil
}
s.digestStore.Save(groupID, prefixHash, digestChain, uuid, accountID, oldDigestChain)
return nil
} }
func (s *GatewayService) extractCacheableContent(parsed *ParsedRequest) string { func (s *GatewayService) extractCacheableContent(parsed *ParsedRequest) string {
@@ -629,8 +639,8 @@ func (s *GatewayService) extractTextFromContent(content any) string {
} }
func (s *GatewayService) hashContent(content string) string { func (s *GatewayService) hashContent(content string) string {
hash := sha256.Sum256([]byte(content)) h := xxhash.Sum64String(content)
return hex.EncodeToString(hash[:16]) // 32字符 return strconv.FormatUint(h, 36)
} }
// replaceModelInBody 替换请求体中的model字段 // replaceModelInBody 替换请求体中的model字段
@@ -989,13 +999,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
log.Printf("[ModelRoutingDebug] load-aware enabled: group_id=%v model=%s session=%s platform=%s", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), platform) log.Printf("[ModelRoutingDebug] load-aware enabled: group_id=%v model=%s session=%s platform=%s", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), platform)
} }
// Antigravity 模型系列检查(在账号选择前检查,确保所有代码路径都经过此检查)
if platform == PlatformAntigravity && groupID != nil && requestedModel != "" {
if err := s.checkAntigravityModelScope(ctx, *groupID, requestedModel); err != nil {
return nil, err
}
}
accounts, useMixed, err := s.listSchedulableAccounts(ctx, groupID, platform, hasForcePlatform) accounts, useMixed, err := s.listSchedulableAccounts(ctx, groupID, platform, hasForcePlatform)
if err != nil { if err != nil {
return nil, err return nil, err
@@ -1110,7 +1113,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
result.ReleaseFunc() // 释放槽位 result.ReleaseFunc() // 释放槽位
// 继续到负载感知选择 // 继续到负载感知选择
} else { } else {
_ = s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL)
if s.debugModelRoutingEnabled() { if s.debugModelRoutingEnabled() {
log.Printf("[ModelRoutingDebug] routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), stickyAccountID) log.Printf("[ModelRoutingDebug] routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), stickyAccountID)
} }
@@ -1190,6 +1192,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
return a.account.LastUsedAt.Before(*b.account.LastUsedAt) return a.account.LastUsedAt.Before(*b.account.LastUsedAt)
} }
}) })
shuffleWithinSortGroups(routingAvailable)
// 4. 尝试获取槽位 // 4. 尝试获取槽位
for _, item := range routingAvailable { for _, item := range routingAvailable {
@@ -1264,7 +1267,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
if !s.checkAndRegisterSession(ctx, account, sessionHash) { if !s.checkAndRegisterSession(ctx, account, sessionHash) {
result.ReleaseFunc() // 释放槽位,继续到 Layer 2 result.ReleaseFunc() // 释放槽位,继续到 Layer 2
} else { } else {
_ = s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL)
return &AccountSelectionResult{ return &AccountSelectionResult{
Account: account, Account: account,
Acquired: true, Acquired: true,
@@ -1344,10 +1346,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
return result, nil return result, nil
} }
} else { } else {
// Antigravity 平台:获取模型负载信息
var modelLoadMap map[int64]*ModelLoadInfo
isAntigravity := platform == PlatformAntigravity
var available []accountWithLoad var available []accountWithLoad
for _, acc := range candidates { for _, acc := range candidates {
loadInfo := loadMap[acc.ID] loadInfo := loadMap[acc.ID]
@@ -1362,109 +1360,44 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
} }
} }
// Antigravity 平台:按账号实际映射后的模型名获取模型负载(与 Forward 的统计保持一致) // 分层过滤选择:优先级 → 负载率 → LRU
if isAntigravity && requestedModel != "" && s.cache != nil && len(available) > 0 { for len(available) > 0 {
modelLoadMap = make(map[int64]*ModelLoadInfo, len(available)) // 1. 取优先级最小的集合
modelToAccountIDs := make(map[string][]int64) candidates := filterByMinPriority(available)
for _, item := range available { // 2. 取负载率最低的集合
mappedModel := mapAntigravityModel(item.account, requestedModel) candidates = filterByMinLoadRate(candidates)
if mappedModel == "" { // 3. LRU 选择最久未用的账号
continue selected := selectByLRU(candidates, preferOAuth)
} if selected == nil {
modelToAccountIDs[mappedModel] = append(modelToAccountIDs[mappedModel], item.account.ID) break
} }
for model, ids := range modelToAccountIDs {
batch, err := s.cache.GetModelLoadBatch(ctx, ids, model)
if err != nil {
continue
}
for id, info := range batch {
modelLoadMap[id] = info
}
}
if len(modelLoadMap) == 0 {
modelLoadMap = nil
}
}
// Antigravity 平台:优先级硬过滤 →(同优先级内)按调用次数选择(最少优先,新账号用平均值) result, err := s.tryAcquireAccountSlot(ctx, selected.account.ID, selected.account.Concurrency)
// 其他平台:分层过滤选择:优先级 → 负载率 → LRU if err == nil && result.Acquired {
if isAntigravity { // 会话数量限制检查
for len(available) > 0 { if !s.checkAndRegisterSession(ctx, selected.account, sessionHash) {
// 1. 取优先级最小的集合(硬过滤) result.ReleaseFunc() // 释放槽位,继续尝试下一个账号
candidates := filterByMinPriority(available) } else {
// 2. 同优先级内按调用次数选择(调用次数最少优先,新账号使用平均值) if sessionHash != "" && s.cache != nil {
selected := selectByCallCount(candidates, modelLoadMap, preferOAuth) _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.account.ID, stickySessionTTL)
if selected == nil {
break
}
result, err := s.tryAcquireAccountSlot(ctx, selected.account.ID, selected.account.Concurrency)
if err == nil && result.Acquired {
// 会话数量限制检查
if !s.checkAndRegisterSession(ctx, selected.account, sessionHash) {
result.ReleaseFunc() // 释放槽位,继续尝试下一个账号
} else {
if sessionHash != "" && s.cache != nil {
_ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.account.ID, stickySessionTTL)
}
return &AccountSelectionResult{
Account: selected.account,
Acquired: true,
ReleaseFunc: result.ReleaseFunc,
}, nil
} }
return &AccountSelectionResult{
Account: selected.account,
Acquired: true,
ReleaseFunc: result.ReleaseFunc,
}, nil
} }
// 移除已尝试的账号,重新选择
selectedID := selected.account.ID
newAvailable := make([]accountWithLoad, 0, len(available)-1)
for _, acc := range available {
if acc.account.ID != selectedID {
newAvailable = append(newAvailable, acc)
}
}
available = newAvailable
} }
} else {
for len(available) > 0 {
// 1. 取优先级最小的集合
candidates := filterByMinPriority(available)
// 2. 取负载率最低的集合
candidates = filterByMinLoadRate(candidates)
// 3. LRU 选择最久未用的账号
selected := selectByLRU(candidates, preferOAuth)
if selected == nil {
break
}
result, err := s.tryAcquireAccountSlot(ctx, selected.account.ID, selected.account.Concurrency) // 移除已尝试的账号,重新进行分层过滤
if err == nil && result.Acquired { selectedID := selected.account.ID
// 会话数量限制检查 newAvailable := make([]accountWithLoad, 0, len(available)-1)
if !s.checkAndRegisterSession(ctx, selected.account, sessionHash) { for _, acc := range available {
result.ReleaseFunc() // 释放槽位,继续尝试下一个账号 if acc.account.ID != selectedID {
} else { newAvailable = append(newAvailable, acc)
if sessionHash != "" && s.cache != nil {
_ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.account.ID, stickySessionTTL)
}
return &AccountSelectionResult{
Account: selected.account,
Acquired: true,
ReleaseFunc: result.ReleaseFunc,
}, nil
}
} }
// 移除已尝试的账号,重新进行分层过滤
selectedID := selected.account.ID
newAvailable := make([]accountWithLoad, 0, len(available)-1)
for _, acc := range available {
if acc.account.ID != selectedID {
newAvailable = append(newAvailable, acc)
}
}
available = newAvailable
} }
available = newAvailable
} }
} }
@@ -2000,87 +1933,79 @@ func sortAccountsByPriorityAndLastUsed(accounts []*Account, preferOAuth bool) {
return a.LastUsedAt.Before(*b.LastUsedAt) return a.LastUsedAt.Before(*b.LastUsedAt)
} }
}) })
shuffleWithinPriorityAndLastUsed(accounts)
} }
// selectByCallCount 从候选账号中选择调用次数最少的账号Antigravity 专用) // shuffleWithinSortGroups 对排序后的 accountWithLoad 切片,按 (Priority, LoadRate, LastUsedAt) 分组后组内随机打乱。
// 新账号CallCount=0使用平均调用次数作为虚拟值避免冷启动被猛调 // 防止并发请求读取同一快照时,确定性排序导致所有请求命中相同账号。
// 如果有多个账号具有相同的最小调用次数,则随机选择一个 func shuffleWithinSortGroups(accounts []accountWithLoad) {
func selectByCallCount(accounts []accountWithLoad, modelLoadMap map[int64]*ModelLoadInfo, preferOAuth bool) *accountWithLoad { if len(accounts) <= 1 {
if len(accounts) == 0 { return
return nil
} }
if len(accounts) == 1 { i := 0
return &accounts[0] for i < len(accounts) {
} j := i + 1
for j < len(accounts) && sameAccountWithLoadGroup(accounts[i], accounts[j]) {
// 如果没有负载信息,回退到 LRU j++
if modelLoadMap == nil {
return selectByLRU(accounts, preferOAuth)
}
// 1. 计算平均调用次数(用于新账号冷启动)
var totalCallCount int64
var countWithCalls int
for _, acc := range accounts {
if info := modelLoadMap[acc.account.ID]; info != nil && info.CallCount > 0 {
totalCallCount += info.CallCount
countWithCalls++
} }
} if j-i > 1 {
mathrand.Shuffle(j-i, func(a, b int) {
var avgCallCount int64 accounts[i+a], accounts[i+b] = accounts[i+b], accounts[i+a]
if countWithCalls > 0 { })
avgCallCount = totalCallCount / int64(countWithCalls)
}
// 2. 获取每个账号的有效调用次数
getEffectiveCallCount := func(acc accountWithLoad) int64 {
if acc.account == nil {
return 0
} }
info := modelLoadMap[acc.account.ID] i = j
if info == nil || info.CallCount == 0 {
return avgCallCount // 新账号使用平均值
}
return info.CallCount
} }
}
// 3. 找到最小调用次数 // sameAccountWithLoadGroup 判断两个 accountWithLoad 是否属于同一排序组
minCount := getEffectiveCallCount(accounts[0]) func sameAccountWithLoadGroup(a, b accountWithLoad) bool {
for _, acc := range accounts[1:] { if a.account.Priority != b.account.Priority {
if c := getEffectiveCallCount(acc); c < minCount { return false
minCount = c
}
} }
if a.loadInfo.LoadRate != b.loadInfo.LoadRate {
// 4. 收集所有具有最小调用次数的账号 return false
var candidateIdxs []int
for i, acc := range accounts {
if getEffectiveCallCount(acc) == minCount {
candidateIdxs = append(candidateIdxs, i)
}
} }
return sameLastUsedAt(a.account.LastUsedAt, b.account.LastUsedAt)
}
// 5. 如果只有一个候选,直接返回 // shuffleWithinPriorityAndLastUsed 对排序后的 []*Account 切片,按 (Priority, LastUsedAt) 分组后组内随机打乱。
if len(candidateIdxs) == 1 { func shuffleWithinPriorityAndLastUsed(accounts []*Account) {
return &accounts[candidateIdxs[0]] if len(accounts) <= 1 {
return
} }
i := 0
// 6. preferOAuth 处理 for i < len(accounts) {
if preferOAuth { j := i + 1
var oauthIdxs []int for j < len(accounts) && sameAccountGroup(accounts[i], accounts[j]) {
for _, idx := range candidateIdxs { j++
if accounts[idx].account.Type == AccountTypeOAuth {
oauthIdxs = append(oauthIdxs, idx)
}
} }
if len(oauthIdxs) > 0 { if j-i > 1 {
candidateIdxs = oauthIdxs mathrand.Shuffle(j-i, func(a, b int) {
accounts[i+a], accounts[i+b] = accounts[i+b], accounts[i+a]
})
} }
i = j
} }
}
// 7. 随机选择 // sameAccountGroup 判断两个 Account 是否属于同一排序组Priority + LastUsedAt
return &accounts[candidateIdxs[mathrand.Intn(len(candidateIdxs))]] func sameAccountGroup(a, b *Account) bool {
if a.Priority != b.Priority {
return false
}
return sameLastUsedAt(a.LastUsedAt, b.LastUsedAt)
}
// sameLastUsedAt 判断两个 LastUsedAt 是否相同(精度到秒)
func sameLastUsedAt(a, b *time.Time) bool {
switch {
case a == nil && b == nil:
return true
case a == nil || b == nil:
return false
default:
return a.Unix() == b.Unix()
}
} }
// sortCandidatesForFallback 根据配置选择排序策略 // sortCandidatesForFallback 根据配置选择排序策略
@@ -2135,13 +2060,6 @@ func shuffleWithinPriority(accounts []*Account) {
// selectAccountForModelWithPlatform 选择单平台账户(完全隔离) // selectAccountForModelWithPlatform 选择单平台账户(完全隔离)
func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, platform string) (*Account, error) { func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, platform string) (*Account, error) {
// 对 Antigravity 平台,检查请求的模型系列是否在分组支持范围内
if platform == PlatformAntigravity && groupID != nil && requestedModel != "" {
if err := s.checkAntigravityModelScope(ctx, *groupID, requestedModel); err != nil {
return nil, err
}
}
preferOAuth := platform == PlatformGemini preferOAuth := platform == PlatformGemini
routingAccountIDs := s.routingAccountIDsForRequest(ctx, groupID, requestedModel, platform) routingAccountIDs := s.routingAccountIDsForRequest(ctx, groupID, requestedModel, platform)
@@ -2169,9 +2087,6 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash) _ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
} }
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) { if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
}
if s.debugModelRoutingEnabled() { if s.debugModelRoutingEnabled() {
log.Printf("[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID) log.Printf("[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
} }
@@ -2272,9 +2187,6 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash) _ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
} }
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) { if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
}
return account, nil return account, nil
} }
} }
@@ -2383,9 +2295,6 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
} }
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) { if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) { if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
}
if s.debugModelRoutingEnabled() { if s.debugModelRoutingEnabled() {
log.Printf("[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID) log.Printf("[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
} }
@@ -2488,9 +2397,6 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
} }
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) { if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) { if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
}
return account, nil return account, nil
} }
} }
@@ -5159,27 +5065,6 @@ func (s *GatewayService) validateUpstreamBaseURL(raw string) (string, error) {
return normalized, nil return normalized, nil
} }
// checkAntigravityModelScope 检查 Antigravity 平台的模型系列是否在分组支持范围内
func (s *GatewayService) checkAntigravityModelScope(ctx context.Context, groupID int64, requestedModel string) error {
scope, ok := ResolveAntigravityQuotaScope(requestedModel)
if !ok {
return nil // 无法解析 scope跳过检查
}
group, err := s.resolveGroupByID(ctx, groupID)
if err != nil {
return nil // 查询失败时放行
}
if group == nil {
return nil // 分组不存在时放行
}
if !IsScopeSupported(group.SupportedModelScopes, scope) {
return ErrModelScopeNotSupported
}
return nil
}
// GetAvailableModels returns the list of models available for a group // GetAvailableModels returns the list of models available for a group
// It aggregates model_mapping keys from all schedulable accounts in the group // It aggregates model_mapping keys from all schedulable accounts in the group
func (s *GatewayService) GetAvailableModels(ctx context.Context, groupID *int64, platform string) []string { func (s *GatewayService) GetAvailableModels(ctx context.Context, groupID *int64, platform string) []string {

View File

@@ -14,7 +14,7 @@ func BenchmarkGenerateSessionHash_Metadata(b *testing.B) {
b.ReportAllocs() b.ReportAllocs()
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
parsed, err := ParseGatewayRequest(body) parsed, err := ParseGatewayRequest(body, "")
if err != nil { if err != nil {
b.Fatalf("解析请求失败: %v", err) b.Fatalf("解析请求失败: %v", err)
} }

View File

@@ -0,0 +1,384 @@
//go:build unit
package service
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// TestShouldFailoverGeminiUpstreamError — verifies the failover decision
// for the ErrorPolicyNone path (original logic preserved).
// ---------------------------------------------------------------------------
func TestShouldFailoverGeminiUpstreamError(t *testing.T) {
svc := &GeminiMessagesCompatService{}
tests := []struct {
name string
statusCode int
expected bool
}{
{"401_failover", 401, true},
{"403_failover", 403, true},
{"429_failover", 429, true},
{"529_failover", 529, true},
{"500_failover", 500, true},
{"502_failover", 502, true},
{"503_failover", 503, true},
{"400_no_failover", 400, false},
{"404_no_failover", 404, false},
{"422_no_failover", 422, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := svc.shouldFailoverGeminiUpstreamError(tt.statusCode)
require.Equal(t, tt.expected, got)
})
}
}
// ---------------------------------------------------------------------------
// TestCheckErrorPolicy_GeminiAccounts — verifies CheckErrorPolicy works
// correctly for Gemini platform accounts (API Key type).
// ---------------------------------------------------------------------------
func TestCheckErrorPolicy_GeminiAccounts(t *testing.T) {
tests := []struct {
name string
account *Account
statusCode int
body []byte
expected ErrorPolicyResult
}{
{
name: "gemini_apikey_custom_codes_hit",
account: &Account{
ID: 100,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(429), float64(500)},
},
},
statusCode: 429,
body: []byte(`{"error":"rate limited"}`),
expected: ErrorPolicyMatched,
},
{
name: "gemini_apikey_custom_codes_miss",
account: &Account{
ID: 101,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(429)},
},
},
statusCode: 500,
body: []byte(`{"error":"internal"}`),
expected: ErrorPolicySkipped,
},
{
name: "gemini_apikey_no_custom_codes_returns_none",
account: &Account{
ID: 102,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
},
statusCode: 500,
body: []byte(`{"error":"internal"}`),
expected: ErrorPolicyNone,
},
{
name: "gemini_apikey_temp_unschedulable_hit",
account: &Account{
ID: 103,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded"},
"duration_minutes": float64(10),
},
},
},
},
statusCode: 503,
body: []byte(`overloaded service`),
expected: ErrorPolicyTempUnscheduled,
},
{
name: "gemini_custom_codes_override_temp_unschedulable",
account: &Account{
ID: 104,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(503)},
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded"},
"duration_minutes": float64(10),
},
},
},
},
statusCode: 503,
body: []byte(`overloaded`),
expected: ErrorPolicyMatched, // custom codes take precedence
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &errorPolicyRepoStub{}
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
result := svc.CheckErrorPolicy(context.Background(), tt.account, tt.statusCode, tt.body)
require.Equal(t, tt.expected, result)
})
}
}
// ---------------------------------------------------------------------------
// TestGeminiErrorPolicyIntegration — verifies the Gemini error handling
// paths produce the correct behavior for each ErrorPolicyResult.
//
// These tests simulate the inline error policy switch in handleClaudeCompat
// and forwardNativeGemini by calling the same methods in the same order.
// ---------------------------------------------------------------------------
func TestGeminiErrorPolicyIntegration(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
account *Account
statusCode int
respBody []byte
expectFailover bool // expect UpstreamFailoverError
expectHandleError bool // expect handleGeminiUpstreamError to be called
expectShouldFailover bool // for None path, whether shouldFailover triggers
}{
{
name: "custom_codes_matched_429_failover",
account: &Account{
ID: 200,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(429)},
},
},
statusCode: 429,
respBody: []byte(`{"error":"rate limited"}`),
expectFailover: true,
expectHandleError: true,
},
{
name: "custom_codes_skipped_500_no_failover",
account: &Account{
ID: 201,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(429)},
},
},
statusCode: 500,
respBody: []byte(`{"error":"internal"}`),
expectFailover: false,
expectHandleError: false,
},
{
name: "temp_unschedulable_matched_failover",
account: &Account{
ID: 202,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded"},
"duration_minutes": float64(10),
},
},
},
},
statusCode: 503,
respBody: []byte(`overloaded`),
expectFailover: true,
expectHandleError: true,
},
{
name: "no_policy_429_failover_via_shouldFailover",
account: &Account{
ID: 203,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
},
statusCode: 429,
respBody: []byte(`{"error":"rate limited"}`),
expectFailover: true,
expectHandleError: true,
expectShouldFailover: true,
},
{
name: "no_policy_400_no_failover",
account: &Account{
ID: 204,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
},
statusCode: 400,
respBody: []byte(`{"error":"bad request"}`),
expectFailover: false,
expectHandleError: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &geminiErrorPolicyRepo{}
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
svc := &GeminiMessagesCompatService{
accountRepo: repo,
rateLimitService: rlSvc,
}
writer := httptest.NewRecorder()
c, _ := gin.CreateTestContext(writer)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
// Simulate the Claude compat error handling path (same logic as native).
// This mirrors the inline switch in handleClaudeCompat.
var handleErrorCalled bool
var gotFailover bool
ctx := context.Background()
statusCode := tt.statusCode
respBody := tt.respBody
account := tt.account
headers := http.Header{}
if svc.rateLimitService != nil {
switch svc.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, respBody) {
case ErrorPolicySkipped:
// Skipped → return error directly (no handleGeminiUpstreamError, no failover)
gotFailover = false
handleErrorCalled = false
goto verify
case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
svc.handleGeminiUpstreamError(ctx, account, statusCode, headers, respBody)
handleErrorCalled = true
gotFailover = true
goto verify
}
}
// ErrorPolicyNone → original logic
svc.handleGeminiUpstreamError(ctx, account, statusCode, headers, respBody)
handleErrorCalled = true
if svc.shouldFailoverGeminiUpstreamError(statusCode) {
gotFailover = true
}
verify:
require.Equal(t, tt.expectFailover, gotFailover, "failover mismatch")
require.Equal(t, tt.expectHandleError, handleErrorCalled, "handleGeminiUpstreamError call mismatch")
if tt.expectShouldFailover {
require.True(t, svc.shouldFailoverGeminiUpstreamError(statusCode),
"shouldFailoverGeminiUpstreamError should return true for status %d", statusCode)
}
})
}
}
// ---------------------------------------------------------------------------
// TestGeminiErrorPolicy_NilRateLimitService — verifies nil safety
// ---------------------------------------------------------------------------
func TestGeminiErrorPolicy_NilRateLimitService(t *testing.T) {
svc := &GeminiMessagesCompatService{
rateLimitService: nil,
}
// When rateLimitService is nil, error policy is skipped → falls through to
// shouldFailoverGeminiUpstreamError (original logic).
// Verify this doesn't panic and follows expected behavior.
ctx := context.Background()
account := &Account{
ID: 300,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(429)},
},
}
// The nil check should prevent CheckErrorPolicy from being called
if svc.rateLimitService != nil {
t.Fatal("rateLimitService should be nil for this test")
}
// shouldFailoverGeminiUpstreamError still works
require.True(t, svc.shouldFailoverGeminiUpstreamError(429))
require.False(t, svc.shouldFailoverGeminiUpstreamError(400))
// handleGeminiUpstreamError should not panic with nil rateLimitService
require.NotPanics(t, func() {
svc.handleGeminiUpstreamError(ctx, account, 500, http.Header{}, []byte(`error`))
})
}
// ---------------------------------------------------------------------------
// geminiErrorPolicyRepo — minimal AccountRepository stub for Gemini error
// policy tests. Embeds mockAccountRepoForGemini and adds tracking.
// ---------------------------------------------------------------------------
type geminiErrorPolicyRepo struct {
mockAccountRepoForGemini
setErrorCalls int
setRateLimitedCalls int
setTempCalls int
}
func (r *geminiErrorPolicyRepo) SetError(_ context.Context, _ int64, _ string) error {
r.setErrorCalls++
return nil
}
func (r *geminiErrorPolicyRepo) SetRateLimited(_ context.Context, _ int64, _ time.Time) error {
r.setRateLimitedCalls++
return nil
}
func (r *geminiErrorPolicyRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, _ string) error {
r.setTempCalls++
return nil
}

View File

@@ -560,10 +560,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
return nil, "", errors.New("gemini api_key not configured") return nil, "", errors.New("gemini api_key not configured")
} }
baseURL := strings.TrimSpace(account.GetCredential("base_url")) baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
if baseURL == "" {
baseURL = geminicli.AIStudioBaseURL
}
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
if err != nil { if err != nil {
return nil, "", err return nil, "", err
@@ -640,10 +637,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
return upstreamReq, "x-request-id", nil return upstreamReq, "x-request-id", nil
} else { } else {
// Mode 2: AI Studio API with OAuth (like API key mode, but using Bearer token) // Mode 2: AI Studio API with OAuth (like API key mode, but using Bearer token)
baseURL := strings.TrimSpace(account.GetCredential("base_url")) baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
if baseURL == "" {
baseURL = geminicli.AIStudioBaseURL
}
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
if err != nil { if err != nil {
return nil, "", err return nil, "", err
@@ -837,38 +831,47 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
if resp.StatusCode >= 400 { if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
tempMatched := false // 统一错误策略:自定义错误码 + 临时不可调度
if s.rateLimitService != nil { if s.rateLimitService != nil {
tempMatched = s.rateLimitService.HandleTempUnschedulable(ctx, account, resp.StatusCode, respBody) switch s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody) {
} case ErrorPolicySkipped:
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) upstreamReqID := resp.Header.Get(requestIDHeader)
if tempMatched { if upstreamReqID == "" {
upstreamReqID := resp.Header.Get(requestIDHeader) upstreamReqID = resp.Header.Get("x-goog-request-id")
if upstreamReqID == "" {
upstreamReqID = resp.Header.Get("x-goog-request-id")
}
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
upstreamDetail := ""
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
if maxBytes <= 0 {
maxBytes = 2048
} }
upstreamDetail = truncateString(string(respBody), maxBytes) return nil, s.writeGeminiMappedError(c, account, resp.StatusCode, upstreamReqID, respBody)
case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
upstreamReqID := resp.Header.Get(requestIDHeader)
if upstreamReqID == "" {
upstreamReqID = resp.Header.Get("x-goog-request-id")
}
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
upstreamDetail := ""
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
if maxBytes <= 0 {
maxBytes = 2048
}
upstreamDetail = truncateString(string(respBody), maxBytes)
}
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
UpstreamRequestID: upstreamReqID,
Kind: "failover",
Message: upstreamMsg,
Detail: upstreamDetail,
})
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
} }
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
UpstreamRequestID: upstreamReqID,
Kind: "failover",
Message: upstreamMsg,
Detail: upstreamDetail,
})
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
} }
// ErrorPolicyNone → 原有逻辑
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) { if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) {
upstreamReqID := resp.Header.Get(requestIDHeader) upstreamReqID := resp.Header.Get(requestIDHeader)
if upstreamReqID == "" { if upstreamReqID == "" {
@@ -1026,10 +1029,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
return nil, "", errors.New("gemini api_key not configured") return nil, "", errors.New("gemini api_key not configured")
} }
baseURL := strings.TrimSpace(account.GetCredential("base_url")) baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
if baseURL == "" {
baseURL = geminicli.AIStudioBaseURL
}
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
if err != nil { if err != nil {
return nil, "", err return nil, "", err
@@ -1097,10 +1097,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
return upstreamReq, "x-request-id", nil return upstreamReq, "x-request-id", nil
} else { } else {
// Mode 2: AI Studio API with OAuth (like API key mode, but using Bearer token) // Mode 2: AI Studio API with OAuth (like API key mode, but using Bearer token)
baseURL := strings.TrimSpace(account.GetCredential("base_url")) baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
if baseURL == "" {
baseURL = geminicli.AIStudioBaseURL
}
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
if err != nil { if err != nil {
return nil, "", err return nil, "", err
@@ -1261,14 +1258,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
if resp.StatusCode >= 400 { if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
tempMatched := false
if s.rateLimitService != nil {
tempMatched = s.rateLimitService.HandleTempUnschedulable(ctx, account, resp.StatusCode, respBody)
}
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
// Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens. // Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens.
// This avoids Gemini SDKs failing hard during preflight token counting. // This avoids Gemini SDKs failing hard during preflight token counting.
// Checked before error policy so it always works regardless of custom error codes.
if action == "countTokens" && isOAuth && isGeminiInsufficientScope(resp.Header, respBody) { if action == "countTokens" && isOAuth && isGeminiInsufficientScope(resp.Header, respBody) {
estimated := estimateGeminiCountTokens(body) estimated := estimateGeminiCountTokens(body)
c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated}) c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
@@ -1282,30 +1274,46 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
}, nil }, nil
} }
if tempMatched { // 统一错误策略:自定义错误码 + 临时不可调度
evBody := unwrapIfNeeded(isOAuth, respBody) if s.rateLimitService != nil {
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody)) switch s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody) {
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) case ErrorPolicySkipped:
upstreamDetail := "" respBody = unwrapIfNeeded(isOAuth, respBody)
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { contentType := resp.Header.Get("Content-Type")
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes if contentType == "" {
if maxBytes <= 0 { contentType = "application/json"
maxBytes = 2048
} }
upstreamDetail = truncateString(string(evBody), maxBytes) c.Data(resp.StatusCode, contentType, respBody)
return nil, fmt.Errorf("gemini upstream error: %d (skipped by error policy)", resp.StatusCode)
case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
evBody := unwrapIfNeeded(isOAuth, respBody)
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
upstreamDetail := ""
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
if maxBytes <= 0 {
maxBytes = 2048
}
upstreamDetail = truncateString(string(evBody), maxBytes)
}
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
UpstreamRequestID: requestID,
Kind: "failover",
Message: upstreamMsg,
Detail: upstreamDetail,
})
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
} }
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
UpstreamRequestID: requestID,
Kind: "failover",
Message: upstreamMsg,
Detail: upstreamDetail,
})
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
} }
// ErrorPolicyNone → 原有逻辑
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) { if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) {
evBody := unwrapIfNeeded(isOAuth, respBody) evBody := unwrapIfNeeded(isOAuth, respBody)
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody)) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody))
@@ -2420,10 +2428,7 @@ func (s *GeminiMessagesCompatService) ForwardAIStudioGET(ctx context.Context, ac
return nil, errors.New("invalid path") return nil, errors.New("invalid path")
} }
baseURL := strings.TrimSpace(account.GetCredential("base_url")) baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
if baseURL == "" {
baseURL = geminicli.AIStudioBaseURL
}
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
if err != nil { if err != nil {
return nil, err return nil, err

View File

@@ -66,6 +66,9 @@ func (m *mockAccountRepoForGemini) Create(ctx context.Context, account *Account)
func (m *mockAccountRepoForGemini) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) { func (m *mockAccountRepoForGemini) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) {
return nil, nil return nil, nil
} }
func (m *mockAccountRepoForGemini) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
return nil, nil
}
func (m *mockAccountRepoForGemini) Update(ctx context.Context, account *Account) error { return nil } func (m *mockAccountRepoForGemini) Update(ctx context.Context, account *Account) error { return nil }
func (m *mockAccountRepoForGemini) Delete(ctx context.Context, id int64) error { return nil } func (m *mockAccountRepoForGemini) Delete(ctx context.Context, id int64) error { return nil }
func (m *mockAccountRepoForGemini) List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) { func (m *mockAccountRepoForGemini) List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) {
@@ -133,9 +136,6 @@ func (m *mockAccountRepoForGemini) ListSchedulableByGroupIDAndPlatforms(ctx cont
func (m *mockAccountRepoForGemini) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { func (m *mockAccountRepoForGemini) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
return nil return nil
} }
func (m *mockAccountRepoForGemini) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error {
return nil
}
func (m *mockAccountRepoForGemini) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error { func (m *mockAccountRepoForGemini) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error {
return nil return nil
} }
@@ -226,6 +226,10 @@ func (m *mockGroupRepoForGemini) GetAccountIDsByGroupIDs(ctx context.Context, gr
return nil, nil return nil, nil
} }
func (m *mockGroupRepoForGemini) UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
return nil
}
var _ GroupRepository = (*mockGroupRepoForGemini)(nil) var _ GroupRepository = (*mockGroupRepoForGemini)(nil)
// mockGatewayCacheForGemini Gemini 测试用的 cache mock // mockGatewayCacheForGemini Gemini 测试用的 cache mock
@@ -265,22 +269,6 @@ func (m *mockGatewayCacheForGemini) DeleteSessionAccountID(ctx context.Context,
return nil return nil
} }
func (m *mockGatewayCacheForGemini) IncrModelCallCount(ctx context.Context, accountID int64, model string) (int64, error) {
return 0, nil
}
func (m *mockGatewayCacheForGemini) GetModelLoadBatch(ctx context.Context, accountIDs []int64, model string) (map[int64]*ModelLoadInfo, error) {
return nil, nil
}
func (m *mockGatewayCacheForGemini) FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
return "", 0, false
}
func (m *mockGatewayCacheForGemini) SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
return nil
}
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择 // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择
func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) { func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) {
ctx := context.Background() ctx := context.Background()

View File

@@ -6,26 +6,11 @@ import (
"encoding/json" "encoding/json"
"strconv" "strconv"
"strings" "strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
"github.com/cespare/xxhash/v2" "github.com/cespare/xxhash/v2"
) )
// Gemini 会话 ID Fallback 相关常量
const (
// geminiSessionTTLSeconds Gemini 会话缓存 TTL5 分钟)
geminiSessionTTLSeconds = 300
// geminiSessionKeyPrefix Gemini 会话 Redis key 前缀
geminiSessionKeyPrefix = "gemini:sess:"
)
// GeminiSessionTTL 返回 Gemini 会话缓存 TTL
func GeminiSessionTTL() time.Duration {
return geminiSessionTTLSeconds * time.Second
}
// shortHash 使用 XXHash64 + Base36 生成短 hash16 字符) // shortHash 使用 XXHash64 + Base36 生成短 hash16 字符)
// XXHash64 比 SHA256 快约 10 倍Base36 比 Hex 短约 20% // XXHash64 比 SHA256 快约 10 倍Base36 比 Hex 短约 20%
func shortHash(data []byte) string { func shortHash(data []byte) string {
@@ -79,35 +64,6 @@ func GenerateGeminiPrefixHash(userID, apiKeyID int64, ip, userAgent, platform, m
return base64.RawURLEncoding.EncodeToString(hash[:12]) return base64.RawURLEncoding.EncodeToString(hash[:12])
} }
// BuildGeminiSessionKey 构建 Gemini 会话 Redis key
// 格式: gemini:sess:{groupID}:{prefixHash}:{digestChain}
func BuildGeminiSessionKey(groupID int64, prefixHash, digestChain string) string {
return geminiSessionKeyPrefix + strconv.FormatInt(groupID, 10) + ":" + prefixHash + ":" + digestChain
}
// GenerateDigestChainPrefixes 生成摘要链的所有前缀(从长到短)
// 用于 MGET 批量查询最长匹配
func GenerateDigestChainPrefixes(chain string) []string {
if chain == "" {
return nil
}
var prefixes []string
c := chain
for c != "" {
prefixes = append(prefixes, c)
// 找到最后一个 "-" 的位置
if i := strings.LastIndex(c, "-"); i > 0 {
c = c[:i]
} else {
break
}
}
return prefixes
}
// ParseGeminiSessionValue 解析 Gemini 会话缓存值 // ParseGeminiSessionValue 解析 Gemini 会话缓存值
// 格式: {uuid}:{accountID} // 格式: {uuid}:{accountID}
func ParseGeminiSessionValue(value string) (uuid string, accountID int64, ok bool) { func ParseGeminiSessionValue(value string) (uuid string, accountID int64, ok bool) {
@@ -139,15 +95,6 @@ func FormatGeminiSessionValue(uuid string, accountID int64) string {
// geminiDigestSessionKeyPrefix Gemini 摘要 fallback 会话 key 前缀 // geminiDigestSessionKeyPrefix Gemini 摘要 fallback 会话 key 前缀
const geminiDigestSessionKeyPrefix = "gemini:digest:" const geminiDigestSessionKeyPrefix = "gemini:digest:"
// geminiTrieKeyPrefix Gemini Trie 会话 key 前缀
const geminiTrieKeyPrefix = "gemini:trie:"
// BuildGeminiTrieKey 构建 Gemini Trie Redis key
// 格式: gemini:trie:{groupID}:{prefixHash}
func BuildGeminiTrieKey(groupID int64, prefixHash string) string {
return geminiTrieKeyPrefix + strconv.FormatInt(groupID, 10) + ":" + prefixHash
}
// GenerateGeminiDigestSessionKey 生成 Gemini 摘要 fallback 的 sessionKey // GenerateGeminiDigestSessionKey 生成 Gemini 摘要 fallback 的 sessionKey
// 组合 prefixHash 前 8 位 + uuid 前 8 位,确保不同会话产生不同的 sessionKey // 组合 prefixHash 前 8 位 + uuid 前 8 位,确保不同会话产生不同的 sessionKey
// 用于在 SelectAccountWithLoadAwareness 中保持粘性会话 // 用于在 SelectAccountWithLoadAwareness 中保持粘性会话

View File

@@ -1,41 +1,14 @@
package service package service
import ( import (
"context"
"testing" "testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
) )
// mockGeminiSessionCache 模拟 Redis 缓存
type mockGeminiSessionCache struct {
sessions map[string]string // key -> value
}
func newMockGeminiSessionCache() *mockGeminiSessionCache {
return &mockGeminiSessionCache{sessions: make(map[string]string)}
}
func (m *mockGeminiSessionCache) Save(groupID int64, prefixHash, digestChain, uuid string, accountID int64) {
key := BuildGeminiSessionKey(groupID, prefixHash, digestChain)
value := FormatGeminiSessionValue(uuid, accountID)
m.sessions[key] = value
}
func (m *mockGeminiSessionCache) Find(groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
prefixes := GenerateDigestChainPrefixes(digestChain)
for _, p := range prefixes {
key := BuildGeminiSessionKey(groupID, prefixHash, p)
if val, ok := m.sessions[key]; ok {
return ParseGeminiSessionValue(val)
}
}
return "", 0, false
}
// TestGeminiSessionContinuousConversation 测试连续会话的摘要链匹配 // TestGeminiSessionContinuousConversation 测试连续会话的摘要链匹配
func TestGeminiSessionContinuousConversation(t *testing.T) { func TestGeminiSessionContinuousConversation(t *testing.T) {
cache := newMockGeminiSessionCache() store := NewDigestSessionStore()
groupID := int64(1) groupID := int64(1)
prefixHash := "test_prefix_hash" prefixHash := "test_prefix_hash"
sessionUUID := "session-uuid-12345" sessionUUID := "session-uuid-12345"
@@ -54,13 +27,13 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
t.Logf("Round 1 chain: %s", chain1) t.Logf("Round 1 chain: %s", chain1)
// 第一轮:没有找到会话,创建新会话 // 第一轮:没有找到会话,创建新会话
_, _, found := cache.Find(groupID, prefixHash, chain1) _, _, _, found := store.Find(groupID, prefixHash, chain1)
if found { if found {
t.Error("Round 1: should not find existing session") t.Error("Round 1: should not find existing session")
} }
// 保存第一轮会话 // 保存第一轮会话(首轮无旧 chain
cache.Save(groupID, prefixHash, chain1, sessionUUID, accountID) store.Save(groupID, prefixHash, chain1, sessionUUID, accountID, "")
// 模拟第二轮对话(用户继续对话) // 模拟第二轮对话(用户继续对话)
req2 := &antigravity.GeminiRequest{ req2 := &antigravity.GeminiRequest{
@@ -77,7 +50,7 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
t.Logf("Round 2 chain: %s", chain2) t.Logf("Round 2 chain: %s", chain2)
// 第二轮:应该能找到会话(通过前缀匹配) // 第二轮:应该能找到会话(通过前缀匹配)
foundUUID, foundAccID, found := cache.Find(groupID, prefixHash, chain2) foundUUID, foundAccID, matchedChain, found := store.Find(groupID, prefixHash, chain2)
if !found { if !found {
t.Error("Round 2: should find session via prefix matching") t.Error("Round 2: should find session via prefix matching")
} }
@@ -88,8 +61,8 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
t.Errorf("Round 2: expected accountID %d, got %d", accountID, foundAccID) t.Errorf("Round 2: expected accountID %d, got %d", accountID, foundAccID)
} }
// 保存第二轮会话 // 保存第二轮会话,传入 Find 返回的 matchedChain 以删旧 key
cache.Save(groupID, prefixHash, chain2, sessionUUID, accountID) store.Save(groupID, prefixHash, chain2, sessionUUID, accountID, matchedChain)
// 模拟第三轮对话 // 模拟第三轮对话
req3 := &antigravity.GeminiRequest{ req3 := &antigravity.GeminiRequest{
@@ -108,7 +81,7 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
t.Logf("Round 3 chain: %s", chain3) t.Logf("Round 3 chain: %s", chain3)
// 第三轮:应该能找到会话(通过第二轮的前缀匹配) // 第三轮:应该能找到会话(通过第二轮的前缀匹配)
foundUUID, foundAccID, found = cache.Find(groupID, prefixHash, chain3) foundUUID, foundAccID, _, found = store.Find(groupID, prefixHash, chain3)
if !found { if !found {
t.Error("Round 3: should find session via prefix matching") t.Error("Round 3: should find session via prefix matching")
} }
@@ -118,13 +91,11 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
if foundAccID != accountID { if foundAccID != accountID {
t.Errorf("Round 3: expected accountID %d, got %d", accountID, foundAccID) t.Errorf("Round 3: expected accountID %d, got %d", accountID, foundAccID)
} }
t.Log("✓ Continuous conversation session matching works correctly!")
} }
// TestGeminiSessionDifferentConversations 测试不同会话不会错误匹配 // TestGeminiSessionDifferentConversations 测试不同会话不会错误匹配
func TestGeminiSessionDifferentConversations(t *testing.T) { func TestGeminiSessionDifferentConversations(t *testing.T) {
cache := newMockGeminiSessionCache() store := NewDigestSessionStore()
groupID := int64(1) groupID := int64(1)
prefixHash := "test_prefix_hash" prefixHash := "test_prefix_hash"
@@ -135,7 +106,7 @@ func TestGeminiSessionDifferentConversations(t *testing.T) {
}, },
} }
chain1 := BuildGeminiDigestChain(req1) chain1 := BuildGeminiDigestChain(req1)
cache.Save(groupID, prefixHash, chain1, "session-1", 100) store.Save(groupID, prefixHash, chain1, "session-1", 100, "")
// 第二个完全不同的会话 // 第二个完全不同的会话
req2 := &antigravity.GeminiRequest{ req2 := &antigravity.GeminiRequest{
@@ -146,61 +117,29 @@ func TestGeminiSessionDifferentConversations(t *testing.T) {
chain2 := BuildGeminiDigestChain(req2) chain2 := BuildGeminiDigestChain(req2)
// 不同会话不应该匹配 // 不同会话不应该匹配
_, _, found := cache.Find(groupID, prefixHash, chain2) _, _, _, found := store.Find(groupID, prefixHash, chain2)
if found { if found {
t.Error("Different conversations should not match") t.Error("Different conversations should not match")
} }
t.Log("✓ Different conversations are correctly isolated!")
} }
// TestGeminiSessionPrefixMatchingOrder 测试前缀匹配的优先级(最长匹配优先) // TestGeminiSessionPrefixMatchingOrder 测试前缀匹配的优先级(最长匹配优先)
func TestGeminiSessionPrefixMatchingOrder(t *testing.T) { func TestGeminiSessionPrefixMatchingOrder(t *testing.T) {
cache := newMockGeminiSessionCache() store := NewDigestSessionStore()
groupID := int64(1) groupID := int64(1)
prefixHash := "test_prefix_hash" prefixHash := "test_prefix_hash"
// 创建一个三轮对话
req := &antigravity.GeminiRequest{
SystemInstruction: &antigravity.GeminiContent{
Parts: []antigravity.GeminiPart{{Text: "System prompt"}},
},
Contents: []antigravity.GeminiContent{
{Role: "user", Parts: []antigravity.GeminiPart{{Text: "Q1"}}},
{Role: "model", Parts: []antigravity.GeminiPart{{Text: "A1"}}},
{Role: "user", Parts: []antigravity.GeminiPart{{Text: "Q2"}}},
},
}
fullChain := BuildGeminiDigestChain(req)
prefixes := GenerateDigestChainPrefixes(fullChain)
t.Logf("Full chain: %s", fullChain)
t.Logf("Prefixes (longest first): %v", prefixes)
// 验证前缀生成顺序(从长到短)
if len(prefixes) != 4 {
t.Errorf("Expected 4 prefixes, got %d", len(prefixes))
}
// 保存不同轮次的会话到不同账号 // 保存不同轮次的会话到不同账号
// 第一轮(最短前缀)-> 账号 1 store.Save(groupID, prefixHash, "s:sys-u:q1", "session-round1", 1, "")
cache.Save(groupID, prefixHash, prefixes[3], "session-round1", 1) store.Save(groupID, prefixHash, "s:sys-u:q1-m:a1", "session-round2", 2, "")
// 第二轮 -> 账号 2 store.Save(groupID, prefixHash, "s:sys-u:q1-m:a1-u:q2", "session-round3", 3, "")
cache.Save(groupID, prefixHash, prefixes[2], "session-round2", 2)
// 第三轮(最长前缀,完整链)-> 账号 3
cache.Save(groupID, prefixHash, prefixes[0], "session-round3", 3)
// 查找应该返回最长匹配(账号 3 // 查找更长的链,应该返回最长匹配(账号 3
_, accID, found := cache.Find(groupID, prefixHash, fullChain) _, accID, _, found := store.Find(groupID, prefixHash, "s:sys-u:q1-m:a1-u:q2-m:a2")
if !found { if !found {
t.Error("Should find session") t.Error("Should find session")
} }
if accID != 3 { if accID != 3 {
t.Errorf("Should match longest prefix (account 3), got account %d", accID) t.Errorf("Should match longest prefix (account 3), got account %d", accID)
} }
t.Log("✓ Longest prefix matching works correctly!")
} }
// 确保 context 包被使用(避免未使用的导入警告)
var _ = context.Background

View File

@@ -152,61 +152,6 @@ func TestGenerateGeminiPrefixHash(t *testing.T) {
} }
} }
func TestGenerateDigestChainPrefixes(t *testing.T) {
tests := []struct {
name string
chain string
want []string
wantLen int
}{
{
name: "empty",
chain: "",
wantLen: 0,
},
{
name: "single part",
chain: "u:abc123",
want: []string{"u:abc123"},
wantLen: 1,
},
{
name: "two parts",
chain: "s:xyz-u:abc",
want: []string{"s:xyz-u:abc", "s:xyz"},
wantLen: 2,
},
{
name: "four parts",
chain: "s:a-u:b-m:c-u:d",
want: []string{"s:a-u:b-m:c-u:d", "s:a-u:b-m:c", "s:a-u:b", "s:a"},
wantLen: 4,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := GenerateDigestChainPrefixes(tt.chain)
if len(result) != tt.wantLen {
t.Errorf("expected %d prefixes, got %d: %v", tt.wantLen, len(result), result)
}
if tt.want != nil {
for i, want := range tt.want {
if i >= len(result) {
t.Errorf("missing prefix at index %d", i)
continue
}
if result[i] != want {
t.Errorf("prefix[%d]: expected %s, got %s", i, want, result[i])
}
}
}
})
}
}
func TestParseGeminiSessionValue(t *testing.T) { func TestParseGeminiSessionValue(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
@@ -442,40 +387,3 @@ func TestGenerateGeminiDigestSessionKey(t *testing.T) {
} }
}) })
} }
func TestBuildGeminiTrieKey(t *testing.T) {
tests := []struct {
name string
groupID int64
prefixHash string
want string
}{
{
name: "normal",
groupID: 123,
prefixHash: "abcdef12",
want: "gemini:trie:123:abcdef12",
},
{
name: "zero group",
groupID: 0,
prefixHash: "xyz",
want: "gemini:trie:0:xyz",
},
{
name: "empty prefix",
groupID: 1,
prefixHash: "",
want: "gemini:trie:1:",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := BuildGeminiTrieKey(tt.groupID, tt.prefixHash)
if got != tt.want {
t.Errorf("BuildGeminiTrieKey(%d, %q) = %q, want %q", tt.groupID, tt.prefixHash, got, tt.want)
}
})
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -45,6 +45,9 @@ type Group struct {
// 可选值: claude, gemini_text, gemini_image // 可选值: claude, gemini_text, gemini_image
SupportedModelScopes []string SupportedModelScopes []string
// 分组排序
SortOrder int
CreatedAt time.Time CreatedAt time.Time
UpdatedAt time.Time UpdatedAt time.Time

View File

@@ -33,6 +33,14 @@ type GroupRepository interface {
GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []int64) ([]int64, error) GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []int64) ([]int64, error)
// BindAccountsToGroup 将多个账号绑定到指定分组 // BindAccountsToGroup 将多个账号绑定到指定分组
BindAccountsToGroup(ctx context.Context, groupID int64, accountIDs []int64) error BindAccountsToGroup(ctx context.Context, groupID int64, accountIDs []int64) error
// UpdateSortOrders 批量更新分组排序
UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error
}
// GroupSortOrderUpdate 分组排序更新
type GroupSortOrderUpdate struct {
ID int64 `json:"id"`
SortOrder int `json:"sort_order"`
} }
// CreateGroupRequest 创建分组请求 // CreateGroupRequest 创建分组请求

View File

@@ -318,110 +318,6 @@ func TestGetModelRateLimitRemainingTime(t *testing.T) {
} }
} }
func TestGetQuotaScopeRateLimitRemainingTime(t *testing.T) {
now := time.Now()
future10m := now.Add(10 * time.Minute).Format(time.RFC3339)
past := now.Add(-10 * time.Minute).Format(time.RFC3339)
tests := []struct {
name string
account *Account
requestedModel string
minExpected time.Duration
maxExpected time.Duration
}{
{
name: "nil account",
account: nil,
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
{
name: "non-antigravity platform",
account: &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{
antigravityQuotaScopesKey: map[string]any{
"claude": map[string]any{
"rate_limit_reset_at": future10m,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
{
name: "claude scope rate limited",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
antigravityQuotaScopesKey: map[string]any{
"claude": map[string]any{
"rate_limit_reset_at": future10m,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 9 * time.Minute,
maxExpected: 11 * time.Minute,
},
{
name: "gemini_text scope rate limited",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
antigravityQuotaScopesKey: map[string]any{
"gemini_text": map[string]any{
"rate_limit_reset_at": future10m,
},
},
},
},
requestedModel: "gemini-3-flash",
minExpected: 9 * time.Minute,
maxExpected: 11 * time.Minute,
},
{
name: "expired scope rate limit",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
antigravityQuotaScopesKey: map[string]any{
"claude": map[string]any{
"rate_limit_reset_at": past,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
{
name: "unsupported model",
account: &Account{
Platform: PlatformAntigravity,
},
requestedModel: "gpt-4",
minExpected: 0,
maxExpected: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.account.GetQuotaScopeRateLimitRemainingTime(tt.requestedModel)
if result < tt.minExpected || result > tt.maxExpected {
t.Errorf("GetQuotaScopeRateLimitRemainingTime() = %v, want between %v and %v", result, tt.minExpected, tt.maxExpected)
}
})
}
}
func TestGetRateLimitRemainingTime(t *testing.T) { func TestGetRateLimitRemainingTime(t *testing.T) {
now := time.Now() now := time.Now()
future15m := now.Add(15 * time.Minute).Format(time.RFC3339) future15m := now.Add(15 * time.Minute).Format(time.RFC3339)
@@ -442,45 +338,19 @@ func TestGetRateLimitRemainingTime(t *testing.T) {
maxExpected: 0, maxExpected: 0,
}, },
{ {
name: "model remaining > scope remaining - returns model", name: "model rate limited - 15 minutes",
account: &Account{ account: &Account{
Platform: PlatformAntigravity, Platform: PlatformAntigravity,
Extra: map[string]any{ Extra: map[string]any{
modelRateLimitsKey: map[string]any{ modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{ "claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": future15m, // 15 分钟 "rate_limit_reset_at": future15m,
},
},
antigravityQuotaScopesKey: map[string]any{
"claude": map[string]any{
"rate_limit_reset_at": future5m, // 5 分钟
}, },
}, },
}, },
}, },
requestedModel: "claude-sonnet-4-5", requestedModel: "claude-sonnet-4-5",
minExpected: 14 * time.Minute, // 应返回较大的 15 分钟 minExpected: 14 * time.Minute,
maxExpected: 16 * time.Minute,
},
{
name: "scope remaining > model remaining - returns scope",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": future5m, // 5 分钟
},
},
antigravityQuotaScopesKey: map[string]any{
"claude": map[string]any{
"rate_limit_reset_at": future15m, // 15 分钟
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 14 * time.Minute, // 应返回较大的 15 分钟
maxExpected: 16 * time.Minute, maxExpected: 16 * time.Minute,
}, },
{ {
@@ -499,22 +369,6 @@ func TestGetRateLimitRemainingTime(t *testing.T) {
minExpected: 4 * time.Minute, minExpected: 4 * time.Minute,
maxExpected: 6 * time.Minute, maxExpected: 6 * time.Minute,
}, },
{
name: "only scope rate limited",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
antigravityQuotaScopesKey: map[string]any{
"claude": map[string]any{
"rate_limit_reset_at": future5m,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 4 * time.Minute,
maxExpected: 6 * time.Minute,
},
{ {
name: "neither rate limited", name: "neither rate limited",
account: &Account{ account: &Account{

View File

@@ -582,10 +582,6 @@ func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Contex
} }
} }
} else { } else {
type accountWithLoad struct {
account *Account
loadInfo *AccountLoadInfo
}
var available []accountWithLoad var available []accountWithLoad
for _, acc := range candidates { for _, acc := range candidates {
loadInfo := loadMap[acc.ID] loadInfo := loadMap[acc.ID]
@@ -620,6 +616,7 @@ func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Contex
return a.account.LastUsedAt.Before(*b.account.LastUsedAt) return a.account.LastUsedAt.Before(*b.account.LastUsedAt)
} }
}) })
shuffleWithinSortGroups(available)
for _, item := range available { for _, item := range available {
result, err := s.tryAcquireAccountSlot(ctx, item.account.ID, item.account.Concurrency) result, err := s.tryAcquireAccountSlot(ctx, item.account.ID, item.account.Concurrency)

View File

@@ -205,22 +205,6 @@ func (c *stubGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID i
return nil return nil
} }
func (c *stubGatewayCache) IncrModelCallCount(ctx context.Context, accountID int64, model string) (int64, error) {
return 0, nil
}
func (c *stubGatewayCache) GetModelLoadBatch(ctx context.Context, accountIDs []int64, model string) (map[int64]*ModelLoadInfo, error) {
return nil, nil
}
func (c *stubGatewayCache) FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
return "", 0, false
}
func (c *stubGatewayCache) SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
return nil
}
func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) { func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) {
now := time.Now() now := time.Now()
resetAt := now.Add(10 * time.Minute) resetAt := now.Add(10 * time.Minute)

View File

@@ -66,7 +66,6 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
} }
isAvailable := acc.Status == StatusActive && acc.Schedulable && !isRateLimited && !isOverloaded && !isTempUnsched isAvailable := acc.Status == StatusActive && acc.Schedulable && !isRateLimited && !isOverloaded && !isTempUnsched
scopeRateLimits := acc.GetAntigravityScopeRateLimits()
if acc.Platform != "" { if acc.Platform != "" {
if _, ok := platform[acc.Platform]; !ok { if _, ok := platform[acc.Platform]; !ok {
@@ -85,14 +84,6 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
if hasError { if hasError {
p.ErrorCount++ p.ErrorCount++
} }
if len(scopeRateLimits) > 0 {
if p.ScopeRateLimitCount == nil {
p.ScopeRateLimitCount = make(map[string]int64)
}
for scope := range scopeRateLimits {
p.ScopeRateLimitCount[scope]++
}
}
} }
for _, grp := range acc.Groups { for _, grp := range acc.Groups {
@@ -117,14 +108,6 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
if hasError { if hasError {
g.ErrorCount++ g.ErrorCount++
} }
if len(scopeRateLimits) > 0 {
if g.ScopeRateLimitCount == nil {
g.ScopeRateLimitCount = make(map[string]int64)
}
for scope := range scopeRateLimits {
g.ScopeRateLimitCount[scope]++
}
}
} }
displayGroupID := int64(0) displayGroupID := int64(0)
@@ -157,9 +140,6 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
item.RateLimitRemainingSec = &remainingSec item.RateLimitRemainingSec = &remainingSec
} }
} }
if len(scopeRateLimits) > 0 {
item.ScopeRateLimits = scopeRateLimits
}
if isOverloaded && acc.OverloadUntil != nil { if isOverloaded && acc.OverloadUntil != nil {
item.OverloadUntil = acc.OverloadUntil item.OverloadUntil = acc.OverloadUntil
remainingSec := int64(time.Until(*acc.OverloadUntil).Seconds()) remainingSec := int64(time.Until(*acc.OverloadUntil).Seconds())

View File

@@ -50,24 +50,22 @@ type UserConcurrencyInfo struct {
// PlatformAvailability aggregates account availability by platform. // PlatformAvailability aggregates account availability by platform.
type PlatformAvailability struct { type PlatformAvailability struct {
Platform string `json:"platform"` Platform string `json:"platform"`
TotalAccounts int64 `json:"total_accounts"` TotalAccounts int64 `json:"total_accounts"`
AvailableCount int64 `json:"available_count"` AvailableCount int64 `json:"available_count"`
RateLimitCount int64 `json:"rate_limit_count"` RateLimitCount int64 `json:"rate_limit_count"`
ScopeRateLimitCount map[string]int64 `json:"scope_rate_limit_count,omitempty"` ErrorCount int64 `json:"error_count"`
ErrorCount int64 `json:"error_count"`
} }
// GroupAvailability aggregates account availability by group. // GroupAvailability aggregates account availability by group.
type GroupAvailability struct { type GroupAvailability struct {
GroupID int64 `json:"group_id"` GroupID int64 `json:"group_id"`
GroupName string `json:"group_name"` GroupName string `json:"group_name"`
Platform string `json:"platform"` Platform string `json:"platform"`
TotalAccounts int64 `json:"total_accounts"` TotalAccounts int64 `json:"total_accounts"`
AvailableCount int64 `json:"available_count"` AvailableCount int64 `json:"available_count"`
RateLimitCount int64 `json:"rate_limit_count"` RateLimitCount int64 `json:"rate_limit_count"`
ScopeRateLimitCount map[string]int64 `json:"scope_rate_limit_count,omitempty"` ErrorCount int64 `json:"error_count"`
ErrorCount int64 `json:"error_count"`
} }
// AccountAvailability represents current availability for a single account. // AccountAvailability represents current availability for a single account.
@@ -85,11 +83,10 @@ type AccountAvailability struct {
IsOverloaded bool `json:"is_overloaded"` IsOverloaded bool `json:"is_overloaded"`
HasError bool `json:"has_error"` HasError bool `json:"has_error"`
RateLimitResetAt *time.Time `json:"rate_limit_reset_at"` RateLimitResetAt *time.Time `json:"rate_limit_reset_at"`
RateLimitRemainingSec *int64 `json:"rate_limit_remaining_sec"` RateLimitRemainingSec *int64 `json:"rate_limit_remaining_sec"`
ScopeRateLimits map[string]int64 `json:"scope_rate_limits,omitempty"` OverloadUntil *time.Time `json:"overload_until"`
OverloadUntil *time.Time `json:"overload_until"` OverloadRemainingSec *int64 `json:"overload_remaining_sec"`
OverloadRemainingSec *int64 `json:"overload_remaining_sec"` ErrorMessage string `json:"error_message"`
ErrorMessage string `json:"error_message"` TempUnschedulableUntil *time.Time `json:"temp_unschedulable_until,omitempty"`
TempUnschedulableUntil *time.Time `json:"temp_unschedulable_until,omitempty"`
} }

View File

@@ -12,6 +12,7 @@ import (
"strings" "strings"
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
@@ -528,7 +529,7 @@ func (s *OpsService) selectAccountForRetry(ctx context.Context, reqType opsRetry
func extractRetryModelAndStream(reqType opsRetryRequestType, errorLog *OpsErrorLogDetail, body []byte) (model string, stream bool, err error) { func extractRetryModelAndStream(reqType opsRetryRequestType, errorLog *OpsErrorLogDetail, body []byte) (model string, stream bool, err error) {
switch reqType { switch reqType {
case opsRetryTypeMessages: case opsRetryTypeMessages:
parsed, parseErr := ParseGatewayRequest(body) parsed, parseErr := ParseGatewayRequest(body, domain.PlatformAnthropic)
if parseErr != nil { if parseErr != nil {
return "", false, fmt.Errorf("failed to parse messages request body: %w", parseErr) return "", false, fmt.Errorf("failed to parse messages request body: %w", parseErr)
} }
@@ -596,7 +597,7 @@ func (s *OpsService) executeWithAccount(ctx context.Context, reqType opsRetryReq
if s.gatewayService == nil { if s.gatewayService == nil {
return &opsRetryExecution{status: opsRetryStatusFailed, errorMessage: "gateway service not available"} return &opsRetryExecution{status: opsRetryStatusFailed, errorMessage: "gateway service not available"}
} }
parsedReq, parseErr := ParseGatewayRequest(body) parsedReq, parseErr := ParseGatewayRequest(body, domain.PlatformAnthropic)
if parseErr != nil { if parseErr != nil {
return &opsRetryExecution{status: opsRetryStatusFailed, errorMessage: "failed to parse request body"} return &opsRetryExecution{status: opsRetryStatusFailed, errorMessage: "failed to parse request body"}
} }

View File

@@ -62,6 +62,32 @@ func (s *RateLimitService) SetTokenCacheInvalidator(invalidator TokenCacheInvali
s.tokenCacheInvalidator = invalidator s.tokenCacheInvalidator = invalidator
} }
// ErrorPolicyResult 表示错误策略检查的结果
type ErrorPolicyResult int
const (
ErrorPolicyNone ErrorPolicyResult = iota // 未命中任何策略,继续默认逻辑
ErrorPolicySkipped // 自定义错误码开启但未命中,跳过处理
ErrorPolicyMatched // 自定义错误码命中,应停止调度
ErrorPolicyTempUnscheduled // 临时不可调度规则命中
)
// CheckErrorPolicy 检查自定义错误码和临时不可调度规则。
// 自定义错误码开启时覆盖后续所有逻辑(包括临时不可调度)。
func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Account, statusCode int, responseBody []byte) ErrorPolicyResult {
if account.IsCustomErrorCodesEnabled() {
if account.ShouldHandleErrorCode(statusCode) {
return ErrorPolicyMatched
}
slog.Info("account_error_code_skipped", "account_id", account.ID, "status_code", statusCode)
return ErrorPolicySkipped
}
if s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
return ErrorPolicyTempUnscheduled
}
return ErrorPolicyNone
}
// HandleUpstreamError 处理上游错误响应,标记账号状态 // HandleUpstreamError 处理上游错误响应,标记账号状态
// 返回是否应该停止该账号的调度 // 返回是否应该停止该账号的调度
func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) (shouldDisable bool) { func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) (shouldDisable bool) {

View File

@@ -0,0 +1,318 @@
//go:build unit
package service
import (
"testing"
"time"
"github.com/stretchr/testify/require"
)
// ============ shuffleWithinSortGroups 测试 ============
func TestShuffleWithinSortGroups_Empty(t *testing.T) {
shuffleWithinSortGroups(nil)
shuffleWithinSortGroups([]accountWithLoad{})
}
func TestShuffleWithinSortGroups_SingleElement(t *testing.T) {
accounts := []accountWithLoad{
{account: &Account{ID: 1, Priority: 1}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
}
shuffleWithinSortGroups(accounts)
require.Equal(t, int64(1), accounts[0].account.ID)
}
func TestShuffleWithinSortGroups_DifferentGroups_OrderPreserved(t *testing.T) {
now := time.Now()
earlier := now.Add(-1 * time.Hour)
accounts := []accountWithLoad{
{account: &Account{ID: 1, Priority: 1, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
{account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 20}},
{account: &Account{ID: 3, Priority: 2, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
}
// 每个元素都属于不同组Priority 或 LoadRate 或 LastUsedAt 不同),顺序不变
for i := 0; i < 20; i++ {
cpy := make([]accountWithLoad, len(accounts))
copy(cpy, accounts)
shuffleWithinSortGroups(cpy)
require.Equal(t, int64(1), cpy[0].account.ID)
require.Equal(t, int64(2), cpy[1].account.ID)
require.Equal(t, int64(3), cpy[2].account.ID)
}
}
func TestShuffleWithinSortGroups_SameGroup_Shuffled(t *testing.T) {
now := time.Now()
// 同一秒的时间戳视为同一组
sameSecond := time.Unix(now.Unix(), 0)
sameSecond2 := time.Unix(now.Unix(), 500_000_000) // 同一秒但不同纳秒
accounts := []accountWithLoad{
{account: &Account{ID: 1, Priority: 1, LastUsedAt: &sameSecond}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
{account: &Account{ID: 2, Priority: 1, LastUsedAt: &sameSecond2}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
{account: &Account{ID: 3, Priority: 1, LastUsedAt: &sameSecond}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
}
// 多次执行,验证所有 ID 都出现在第一个位置(说明确实被打乱了)
seen := map[int64]bool{}
for i := 0; i < 100; i++ {
cpy := make([]accountWithLoad, len(accounts))
copy(cpy, accounts)
shuffleWithinSortGroups(cpy)
seen[cpy[0].account.ID] = true
// 无论怎么打乱,所有 ID 都应在候选中
ids := map[int64]bool{}
for _, a := range cpy {
ids[a.account.ID] = true
}
require.True(t, ids[1] && ids[2] && ids[3])
}
// 至少 2 个不同的 ID 出现在首位(随机性验证)
require.GreaterOrEqual(t, len(seen), 2, "shuffle should produce different orderings")
}
func TestShuffleWithinSortGroups_NilLastUsedAt_SameGroup(t *testing.T) {
accounts := []accountWithLoad{
{account: &Account{ID: 1, Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}},
{account: &Account{ID: 2, Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}},
{account: &Account{ID: 3, Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}},
}
seen := map[int64]bool{}
for i := 0; i < 100; i++ {
cpy := make([]accountWithLoad, len(accounts))
copy(cpy, accounts)
shuffleWithinSortGroups(cpy)
seen[cpy[0].account.ID] = true
}
require.GreaterOrEqual(t, len(seen), 2, "nil LastUsedAt accounts should be shuffled")
}
func TestShuffleWithinSortGroups_MixedGroups(t *testing.T) {
now := time.Now()
earlier := now.Add(-1 * time.Hour)
sameAsNow := time.Unix(now.Unix(), 0)
// 组1: Priority=1, LoadRate=10, LastUsedAt=earlier (ID 1) — 单元素组
// 组2: Priority=1, LoadRate=20, LastUsedAt=now (ID 2, 3) — 双元素组
// 组3: Priority=2, LoadRate=10, LastUsedAt=earlier (ID 4) — 单元素组
accounts := []accountWithLoad{
{account: &Account{ID: 1, Priority: 1, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
{account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 20}},
{account: &Account{ID: 3, Priority: 1, LastUsedAt: &sameAsNow}, loadInfo: &AccountLoadInfo{LoadRate: 20}},
{account: &Account{ID: 4, Priority: 2, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
}
for i := 0; i < 20; i++ {
cpy := make([]accountWithLoad, len(accounts))
copy(cpy, accounts)
shuffleWithinSortGroups(cpy)
// 组间顺序不变
require.Equal(t, int64(1), cpy[0].account.ID, "group 1 position fixed")
require.Equal(t, int64(4), cpy[3].account.ID, "group 3 position fixed")
// 组2 内部可以打乱,但仍在位置 1 和 2
mid := map[int64]bool{cpy[1].account.ID: true, cpy[2].account.ID: true}
require.True(t, mid[2] && mid[3], "group 2 elements should stay in positions 1-2")
}
}
// ============ shuffleWithinPriorityAndLastUsed 测试 ============
func TestShuffleWithinPriorityAndLastUsed_Empty(t *testing.T) {
shuffleWithinPriorityAndLastUsed(nil)
shuffleWithinPriorityAndLastUsed([]*Account{})
}
func TestShuffleWithinPriorityAndLastUsed_SingleElement(t *testing.T) {
accounts := []*Account{{ID: 1, Priority: 1}}
shuffleWithinPriorityAndLastUsed(accounts)
require.Equal(t, int64(1), accounts[0].ID)
}
func TestShuffleWithinPriorityAndLastUsed_SameGroup_Shuffled(t *testing.T) {
accounts := []*Account{
{ID: 1, Priority: 1, LastUsedAt: nil},
{ID: 2, Priority: 1, LastUsedAt: nil},
{ID: 3, Priority: 1, LastUsedAt: nil},
}
seen := map[int64]bool{}
for i := 0; i < 100; i++ {
cpy := make([]*Account, len(accounts))
copy(cpy, accounts)
shuffleWithinPriorityAndLastUsed(cpy)
seen[cpy[0].ID] = true
}
require.GreaterOrEqual(t, len(seen), 2, "same group should be shuffled")
}
func TestShuffleWithinPriorityAndLastUsed_DifferentPriority_OrderPreserved(t *testing.T) {
accounts := []*Account{
{ID: 1, Priority: 1, LastUsedAt: nil},
{ID: 2, Priority: 2, LastUsedAt: nil},
{ID: 3, Priority: 3, LastUsedAt: nil},
}
for i := 0; i < 20; i++ {
cpy := make([]*Account, len(accounts))
copy(cpy, accounts)
shuffleWithinPriorityAndLastUsed(cpy)
require.Equal(t, int64(1), cpy[0].ID)
require.Equal(t, int64(2), cpy[1].ID)
require.Equal(t, int64(3), cpy[2].ID)
}
}
func TestShuffleWithinPriorityAndLastUsed_DifferentLastUsedAt_OrderPreserved(t *testing.T) {
now := time.Now()
earlier := now.Add(-1 * time.Hour)
accounts := []*Account{
{ID: 1, Priority: 1, LastUsedAt: nil},
{ID: 2, Priority: 1, LastUsedAt: &earlier},
{ID: 3, Priority: 1, LastUsedAt: &now},
}
for i := 0; i < 20; i++ {
cpy := make([]*Account, len(accounts))
copy(cpy, accounts)
shuffleWithinPriorityAndLastUsed(cpy)
require.Equal(t, int64(1), cpy[0].ID)
require.Equal(t, int64(2), cpy[1].ID)
require.Equal(t, int64(3), cpy[2].ID)
}
}
// ============ sameLastUsedAt 测试 ============
func TestSameLastUsedAt(t *testing.T) {
now := time.Now()
sameSecond := time.Unix(now.Unix(), 0)
sameSecondDiffNano := time.Unix(now.Unix(), 999_999_999)
differentSecond := now.Add(1 * time.Second)
t.Run("both nil", func(t *testing.T) {
require.True(t, sameLastUsedAt(nil, nil))
})
t.Run("one nil one not", func(t *testing.T) {
require.False(t, sameLastUsedAt(nil, &now))
require.False(t, sameLastUsedAt(&now, nil))
})
t.Run("same second different nanoseconds", func(t *testing.T) {
require.True(t, sameLastUsedAt(&sameSecond, &sameSecondDiffNano))
})
t.Run("different seconds", func(t *testing.T) {
require.False(t, sameLastUsedAt(&now, &differentSecond))
})
t.Run("exact same time", func(t *testing.T) {
require.True(t, sameLastUsedAt(&now, &now))
})
}
// ============ sameAccountWithLoadGroup 测试 ============
func TestSameAccountWithLoadGroup(t *testing.T) {
now := time.Now()
sameSecond := time.Unix(now.Unix(), 0)
t.Run("same group", func(t *testing.T) {
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
b := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &sameSecond}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
require.True(t, sameAccountWithLoadGroup(a, b))
})
t.Run("different priority", func(t *testing.T) {
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
b := accountWithLoad{account: &Account{Priority: 2, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
require.False(t, sameAccountWithLoadGroup(a, b))
})
t.Run("different load rate", func(t *testing.T) {
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
b := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 20}}
require.False(t, sameAccountWithLoadGroup(a, b))
})
t.Run("different last used at", func(t *testing.T) {
later := now.Add(1 * time.Second)
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
b := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &later}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
require.False(t, sameAccountWithLoadGroup(a, b))
})
t.Run("both nil LastUsedAt", func(t *testing.T) {
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}}
b := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}}
require.True(t, sameAccountWithLoadGroup(a, b))
})
}
// ============ sameAccountGroup 测试 ============
func TestSameAccountGroup(t *testing.T) {
now := time.Now()
t.Run("same group", func(t *testing.T) {
a := &Account{Priority: 1, LastUsedAt: nil}
b := &Account{Priority: 1, LastUsedAt: nil}
require.True(t, sameAccountGroup(a, b))
})
t.Run("different priority", func(t *testing.T) {
a := &Account{Priority: 1, LastUsedAt: nil}
b := &Account{Priority: 2, LastUsedAt: nil}
require.False(t, sameAccountGroup(a, b))
})
t.Run("different LastUsedAt", func(t *testing.T) {
later := now.Add(1 * time.Second)
a := &Account{Priority: 1, LastUsedAt: &now}
b := &Account{Priority: 1, LastUsedAt: &later}
require.False(t, sameAccountGroup(a, b))
})
}
// ============ sortAccountsByPriorityAndLastUsed 集成随机化测试 ============
func TestSortAccountsByPriorityAndLastUsed_WithShuffle(t *testing.T) {
t.Run("same priority and nil LastUsedAt are shuffled", func(t *testing.T) {
accounts := []*Account{
{ID: 1, Priority: 1, LastUsedAt: nil},
{ID: 2, Priority: 1, LastUsedAt: nil},
{ID: 3, Priority: 1, LastUsedAt: nil},
}
seen := map[int64]bool{}
for i := 0; i < 100; i++ {
cpy := make([]*Account, len(accounts))
copy(cpy, accounts)
sortAccountsByPriorityAndLastUsed(cpy, false)
seen[cpy[0].ID] = true
}
require.GreaterOrEqual(t, len(seen), 2, "identical sort keys should produce different orderings after shuffle")
})
t.Run("different priorities still sorted correctly", func(t *testing.T) {
now := time.Now()
accounts := []*Account{
{ID: 3, Priority: 3, LastUsedAt: &now},
{ID: 1, Priority: 1, LastUsedAt: &now},
{ID: 2, Priority: 2, LastUsedAt: &now},
}
sortAccountsByPriorityAndLastUsed(accounts, false)
require.Equal(t, int64(1), accounts[0].ID)
require.Equal(t, int64(2), accounts[1].ID)
require.Equal(t, int64(3), accounts[2].ID)
})
}

View File

@@ -23,8 +23,7 @@ import (
// - 临时不可调度且未过期:清理 // - 临时不可调度且未过期:清理
// - 临时不可调度已过期:不清理 // - 临时不可调度已过期:不清理
// - 正常可调度状态:不清理 // - 正常可调度状态:不清理
// - 模型限流超过阈值:清理 // - 模型限流(任意时长):清理
// - 模型限流未超过阈值:不清理
// //
// TestShouldClearStickySession tests the sticky session clearing logic. // TestShouldClearStickySession tests the sticky session clearing logic.
// Verifies correct behavior for various account states including: // Verifies correct behavior for various account states including:
@@ -35,9 +34,9 @@ func TestShouldClearStickySession(t *testing.T) {
future := now.Add(1 * time.Hour) future := now.Add(1 * time.Hour)
past := now.Add(-1 * time.Hour) past := now.Add(-1 * time.Hour)
// 短限流时间(低于阈值,不应清除粘性会话) // 短限流时间(有限流即清除粘性会话)
shortRateLimitReset := now.Add(5 * time.Second).Format(time.RFC3339) shortRateLimitReset := now.Add(5 * time.Second).Format(time.RFC3339)
// 长限流时间(超过阈值,应清除粘性会话) // 长限流时间(有限流即清除粘性会话)
longRateLimitReset := now.Add(30 * time.Second).Format(time.RFC3339) longRateLimitReset := now.Add(30 * time.Second).Format(time.RFC3339)
tests := []struct { tests := []struct {
@@ -53,7 +52,7 @@ func TestShouldClearStickySession(t *testing.T) {
{name: "temp unschedulable", account: &Account{Status: StatusActive, Schedulable: true, TempUnschedulableUntil: &future}, requestedModel: "", want: true}, {name: "temp unschedulable", account: &Account{Status: StatusActive, Schedulable: true, TempUnschedulableUntil: &future}, requestedModel: "", want: true},
{name: "temp unschedulable expired", account: &Account{Status: StatusActive, Schedulable: true, TempUnschedulableUntil: &past}, requestedModel: "", want: false}, {name: "temp unschedulable expired", account: &Account{Status: StatusActive, Schedulable: true, TempUnschedulableUntil: &past}, requestedModel: "", want: false},
{name: "active schedulable", account: &Account{Status: StatusActive, Schedulable: true}, requestedModel: "", want: false}, {name: "active schedulable", account: &Account{Status: StatusActive, Schedulable: true}, requestedModel: "", want: false},
// 模型限流测试 // 模型限流测试:有限流即清除
{ {
name: "model rate limited short duration", name: "model rate limited short duration",
account: &Account{ account: &Account{
@@ -68,7 +67,7 @@ func TestShouldClearStickySession(t *testing.T) {
}, },
}, },
requestedModel: "claude-sonnet-4", requestedModel: "claude-sonnet-4",
want: false, // 低于阈值,不清除 want: true, // 有限流即清除
}, },
{ {
name: "model rate limited long duration", name: "model rate limited long duration",
@@ -84,7 +83,7 @@ func TestShouldClearStickySession(t *testing.T) {
}, },
}, },
requestedModel: "claude-sonnet-4", requestedModel: "claude-sonnet-4",
want: true, // 超过阈值,清除 want: true, // 有限流即清除
}, },
{ {
name: "model rate limited different model", name: "model rate limited different model",

View File

@@ -275,4 +275,5 @@ var ProviderSet = wire.NewSet(
NewUsageCache, NewUsageCache,
NewTotpService, NewTotpService,
NewErrorPassthroughService, NewErrorPassthroughService,
NewDigestSessionStore,
) )

View File

@@ -0,0 +1,8 @@
-- Add sort_order field to groups table for custom ordering
ALTER TABLE groups ADD COLUMN IF NOT EXISTS sort_order INT NOT NULL DEFAULT 0;
-- Initialize existing groups with sort_order based on their ID
UPDATE groups SET sort_order = id WHERE sort_order = 0;
-- Create index for efficient sorting
CREATE INDEX IF NOT EXISTS idx_groups_sort_order ON groups(sort_order);

View File

@@ -0,0 +1,11 @@
-- Migrate upstream accounts to apikey type
-- Background: upstream type is no longer needed. Antigravity platform APIKey accounts
-- with base_url pointing to an upstream sub2api instance can reuse the standard
-- APIKey forwarding path. GetBaseURL()/GetGeminiBaseURL() automatically appends
-- /antigravity for Antigravity platform APIKey accounts.
UPDATE accounts
SET type = 'apikey'
WHERE type = 'upstream'
AND platform = 'antigravity'
AND deleted_at IS NULL;

View File

@@ -27,6 +27,7 @@
"qrcode": "^1.5.4", "qrcode": "^1.5.4",
"vue": "^3.4.0", "vue": "^3.4.0",
"vue-chartjs": "^5.3.0", "vue-chartjs": "^5.3.0",
"vue-draggable-plus": "^0.6.1",
"vue-i18n": "^9.14.5", "vue-i18n": "^9.14.5",
"vue-router": "^4.2.5", "vue-router": "^4.2.5",
"xlsx": "^0.18.5" "xlsx": "^0.18.5"

View File

@@ -44,6 +44,9 @@ importers:
vue-chartjs: vue-chartjs:
specifier: ^5.3.0 specifier: ^5.3.0
version: 5.3.3(chart.js@4.5.1)(vue@3.5.26(typescript@5.6.3)) version: 5.3.3(chart.js@4.5.1)(vue@3.5.26(typescript@5.6.3))
vue-draggable-plus:
specifier: ^0.6.1
version: 0.6.1(@types/sortablejs@1.15.9)
vue-i18n: vue-i18n:
specifier: ^9.14.5 specifier: ^9.14.5
version: 9.14.5(vue@3.5.26(typescript@5.6.3)) version: 9.14.5(vue@3.5.26(typescript@5.6.3))
@@ -1254,67 +1257,56 @@ packages:
resolution: {integrity: sha512-EHMUcDwhtdRGlXZsGSIuXSYwD5kOT9NVnx9sqzYiwAc91wfYOE1g1djOEDseZJKKqtHAHGwnGPQu3kytmfaXLQ==} resolution: {integrity: sha512-EHMUcDwhtdRGlXZsGSIuXSYwD5kOT9NVnx9sqzYiwAc91wfYOE1g1djOEDseZJKKqtHAHGwnGPQu3kytmfaXLQ==}
cpu: [arm] cpu: [arm]
os: [linux] os: [linux]
libc: [glibc]
'@rollup/rollup-linux-arm-musleabihf@4.54.0': '@rollup/rollup-linux-arm-musleabihf@4.54.0':
resolution: {integrity: sha512-+pBrqEjaakN2ySv5RVrj/qLytYhPKEUwk+e3SFU5jTLHIcAtqh2rLrd/OkbNuHJpsBgxsD8ccJt5ga/SeG0JmA==} resolution: {integrity: sha512-+pBrqEjaakN2ySv5RVrj/qLytYhPKEUwk+e3SFU5jTLHIcAtqh2rLrd/OkbNuHJpsBgxsD8ccJt5ga/SeG0JmA==}
cpu: [arm] cpu: [arm]
os: [linux] os: [linux]
libc: [musl]
'@rollup/rollup-linux-arm64-gnu@4.54.0': '@rollup/rollup-linux-arm64-gnu@4.54.0':
resolution: {integrity: sha512-NSqc7rE9wuUaRBsBp5ckQ5CVz5aIRKCwsoa6WMF7G01sX3/qHUw/z4pv+D+ahL1EIKy6Enpcnz1RY8pf7bjwng==} resolution: {integrity: sha512-NSqc7rE9wuUaRBsBp5ckQ5CVz5aIRKCwsoa6WMF7G01sX3/qHUw/z4pv+D+ahL1EIKy6Enpcnz1RY8pf7bjwng==}
cpu: [arm64] cpu: [arm64]
os: [linux] os: [linux]
libc: [glibc]
'@rollup/rollup-linux-arm64-musl@4.54.0': '@rollup/rollup-linux-arm64-musl@4.54.0':
resolution: {integrity: sha512-gr5vDbg3Bakga5kbdpqx81m2n9IX8M6gIMlQQIXiLTNeQW6CucvuInJ91EuCJ/JYvc+rcLLsDFcfAD1K7fMofg==} resolution: {integrity: sha512-gr5vDbg3Bakga5kbdpqx81m2n9IX8M6gIMlQQIXiLTNeQW6CucvuInJ91EuCJ/JYvc+rcLLsDFcfAD1K7fMofg==}
cpu: [arm64] cpu: [arm64]
os: [linux] os: [linux]
libc: [musl]
'@rollup/rollup-linux-loong64-gnu@4.54.0': '@rollup/rollup-linux-loong64-gnu@4.54.0':
resolution: {integrity: sha512-gsrtB1NA3ZYj2vq0Rzkylo9ylCtW/PhpLEivlgWe0bpgtX5+9j9EZa0wtZiCjgu6zmSeZWyI/e2YRX1URozpIw==} resolution: {integrity: sha512-gsrtB1NA3ZYj2vq0Rzkylo9ylCtW/PhpLEivlgWe0bpgtX5+9j9EZa0wtZiCjgu6zmSeZWyI/e2YRX1URozpIw==}
cpu: [loong64] cpu: [loong64]
os: [linux] os: [linux]
libc: [glibc]
'@rollup/rollup-linux-ppc64-gnu@4.54.0': '@rollup/rollup-linux-ppc64-gnu@4.54.0':
resolution: {integrity: sha512-y3qNOfTBStmFNq+t4s7Tmc9hW2ENtPg8FeUD/VShI7rKxNW7O4fFeaYbMsd3tpFlIg1Q8IapFgy7Q9i2BqeBvA==} resolution: {integrity: sha512-y3qNOfTBStmFNq+t4s7Tmc9hW2ENtPg8FeUD/VShI7rKxNW7O4fFeaYbMsd3tpFlIg1Q8IapFgy7Q9i2BqeBvA==}
cpu: [ppc64] cpu: [ppc64]
os: [linux] os: [linux]
libc: [glibc]
'@rollup/rollup-linux-riscv64-gnu@4.54.0': '@rollup/rollup-linux-riscv64-gnu@4.54.0':
resolution: {integrity: sha512-89sepv7h2lIVPsFma8iwmccN7Yjjtgz0Rj/Ou6fEqg3HDhpCa+Et+YSufy27i6b0Wav69Qv4WBNl3Rs6pwhebQ==} resolution: {integrity: sha512-89sepv7h2lIVPsFma8iwmccN7Yjjtgz0Rj/Ou6fEqg3HDhpCa+Et+YSufy27i6b0Wav69Qv4WBNl3Rs6pwhebQ==}
cpu: [riscv64] cpu: [riscv64]
os: [linux] os: [linux]
libc: [glibc]
'@rollup/rollup-linux-riscv64-musl@4.54.0': '@rollup/rollup-linux-riscv64-musl@4.54.0':
resolution: {integrity: sha512-ZcU77ieh0M2Q8Ur7D5X7KvK+UxbXeDHwiOt/CPSBTI1fBmeDMivW0dPkdqkT4rOgDjrDDBUed9x4EgraIKoR2A==} resolution: {integrity: sha512-ZcU77ieh0M2Q8Ur7D5X7KvK+UxbXeDHwiOt/CPSBTI1fBmeDMivW0dPkdqkT4rOgDjrDDBUed9x4EgraIKoR2A==}
cpu: [riscv64] cpu: [riscv64]
os: [linux] os: [linux]
libc: [musl]
'@rollup/rollup-linux-s390x-gnu@4.54.0': '@rollup/rollup-linux-s390x-gnu@4.54.0':
resolution: {integrity: sha512-2AdWy5RdDF5+4YfG/YesGDDtbyJlC9LHmL6rZw6FurBJ5n4vFGupsOBGfwMRjBYH7qRQowT8D/U4LoSvVwOhSQ==} resolution: {integrity: sha512-2AdWy5RdDF5+4YfG/YesGDDtbyJlC9LHmL6rZw6FurBJ5n4vFGupsOBGfwMRjBYH7qRQowT8D/U4LoSvVwOhSQ==}
cpu: [s390x] cpu: [s390x]
os: [linux] os: [linux]
libc: [glibc]
'@rollup/rollup-linux-x64-gnu@4.54.0': '@rollup/rollup-linux-x64-gnu@4.54.0':
resolution: {integrity: sha512-WGt5J8Ij/rvyqpFexxk3ffKqqbLf9AqrTBbWDk7ApGUzaIs6V+s2s84kAxklFwmMF/vBNGrVdYgbblCOFFezMQ==} resolution: {integrity: sha512-WGt5J8Ij/rvyqpFexxk3ffKqqbLf9AqrTBbWDk7ApGUzaIs6V+s2s84kAxklFwmMF/vBNGrVdYgbblCOFFezMQ==}
cpu: [x64] cpu: [x64]
os: [linux] os: [linux]
libc: [glibc]
'@rollup/rollup-linux-x64-musl@4.54.0': '@rollup/rollup-linux-x64-musl@4.54.0':
resolution: {integrity: sha512-JzQmb38ATzHjxlPHuTH6tE7ojnMKM2kYNzt44LO/jJi8BpceEC8QuXYA908n8r3CNuG/B3BV8VR3Hi1rYtmPiw==} resolution: {integrity: sha512-JzQmb38ATzHjxlPHuTH6tE7ojnMKM2kYNzt44LO/jJi8BpceEC8QuXYA908n8r3CNuG/B3BV8VR3Hi1rYtmPiw==}
cpu: [x64] cpu: [x64]
os: [linux] os: [linux]
libc: [musl]
'@rollup/rollup-openharmony-arm64@4.54.0': '@rollup/rollup-openharmony-arm64@4.54.0':
resolution: {integrity: sha512-huT3fd0iC7jigGh7n3q/+lfPcXxBi+om/Rs3yiFxjvSxbSB6aohDFXbWvlspaqjeOh+hx7DDHS+5Es5qRkWkZg==} resolution: {integrity: sha512-huT3fd0iC7jigGh7n3q/+lfPcXxBi+om/Rs3yiFxjvSxbSB6aohDFXbWvlspaqjeOh+hx7DDHS+5Es5qRkWkZg==}
@@ -1515,6 +1507,9 @@ packages:
'@types/react@19.2.7': '@types/react@19.2.7':
resolution: {integrity: sha512-MWtvHrGZLFttgeEj28VXHxpmwYbor/ATPYbBfSFZEIRK0ecCFLl2Qo55z52Hss+UV9CRN7trSeq1zbgx7YDWWg==} resolution: {integrity: sha512-MWtvHrGZLFttgeEj28VXHxpmwYbor/ATPYbBfSFZEIRK0ecCFLl2Qo55z52Hss+UV9CRN7trSeq1zbgx7YDWWg==}
'@types/sortablejs@1.15.9':
resolution: {integrity: sha512-7HP+rZGE2p886PKV9c9OJzLBI6BBJu1O7lJGYnPyG3fS4/duUCcngkNCjsLwIMV+WMqANe3tt4irrXHSIe68OQ==}
'@types/trusted-types@2.0.7': '@types/trusted-types@2.0.7':
resolution: {integrity: sha512-ScaPdn1dQczgbl0QFTeTOmVHFULt394XJgOQNoyVhZ6r2vLnMLJfBPd53SB52T/3G36VI1/g2MZaX0cwDuXsfw==} resolution: {integrity: sha512-ScaPdn1dQczgbl0QFTeTOmVHFULt394XJgOQNoyVhZ6r2vLnMLJfBPd53SB52T/3G36VI1/g2MZaX0cwDuXsfw==}
@@ -4298,6 +4293,15 @@ packages:
'@vue/composition-api': '@vue/composition-api':
optional: true optional: true
vue-draggable-plus@0.6.1:
resolution: {integrity: sha512-FbtQ/fuoixiOfTZzG3yoPl4JAo9HJXRHmBQZFB9x2NYCh6pq0TomHf7g5MUmpaDYv+LU2n6BPq2YN9sBO+FbIg==}
peerDependencies:
'@types/sortablejs': ^1.15.0
'@vue/composition-api': '*'
peerDependenciesMeta:
'@vue/composition-api':
optional: true
vue-eslint-parser@9.4.3: vue-eslint-parser@9.4.3:
resolution: {integrity: sha512-2rYRLWlIpaiN8xbPiDyXZXRgLGOtWxERV7ND5fFAv5qo1D2N9Fu9MNajBNc6o13lZ+24DAWCkQCvj4klgmcITg==} resolution: {integrity: sha512-2rYRLWlIpaiN8xbPiDyXZXRgLGOtWxERV7ND5fFAv5qo1D2N9Fu9MNajBNc6o13lZ+24DAWCkQCvj4klgmcITg==}
engines: {node: ^14.17.0 || >=16.0.0} engines: {node: ^14.17.0 || >=16.0.0}
@@ -5958,6 +5962,8 @@ snapshots:
dependencies: dependencies:
csstype: 3.2.3 csstype: 3.2.3
'@types/sortablejs@1.15.9': {}
'@types/trusted-types@2.0.7': {} '@types/trusted-types@2.0.7': {}
'@types/unist@2.0.11': {} '@types/unist@2.0.11': {}
@@ -9401,6 +9407,10 @@ snapshots:
dependencies: dependencies:
vue: 3.5.26(typescript@5.6.3) vue: 3.5.26(typescript@5.6.3)
vue-draggable-plus@0.6.1(@types/sortablejs@1.15.9):
dependencies:
'@types/sortablejs': 1.15.9
vue-eslint-parser@9.4.3(eslint@8.57.1): vue-eslint-parser@9.4.3(eslint@8.57.1):
dependencies: dependencies:
debug: 4.4.3 debug: 4.4.3

View File

@@ -327,11 +327,34 @@ export async function getAvailableModels(id: number): Promise<ClaudeModel[]> {
return data return data
} }
export interface CRSPreviewAccount {
crs_account_id: string
kind: string
name: string
platform: string
type: string
}
export interface PreviewFromCRSResult {
new_accounts: CRSPreviewAccount[]
existing_accounts: CRSPreviewAccount[]
}
export async function previewFromCrs(params: {
base_url: string
username: string
password: string
}): Promise<PreviewFromCRSResult> {
const { data } = await apiClient.post<PreviewFromCRSResult>('/admin/accounts/sync/crs/preview', params)
return data
}
export async function syncFromCrs(params: { export async function syncFromCrs(params: {
base_url: string base_url: string
username: string username: string
password: string password: string
sync_proxies?: boolean sync_proxies?: boolean
selected_account_ids?: string[]
}): Promise<{ }): Promise<{
created: number created: number
updated: number updated: number
@@ -345,7 +368,19 @@ export async function syncFromCrs(params: {
error?: string error?: string
}> }>
}> { }> {
const { data } = await apiClient.post('/admin/accounts/sync/crs', params) const { data } = await apiClient.post<{
created: number
updated: number
skipped: number
failed: number
items: Array<{
crs_account_id: string
kind: string
name: string
action: string
error?: string
}>
}>('/admin/accounts/sync/crs', params)
return data return data
} }
@@ -398,6 +433,26 @@ export async function getAntigravityDefaultModelMapping(): Promise<Record<string
return data return data
} }
/**
* Refresh OpenAI token using refresh token
* @param refreshToken - The refresh token
* @param proxyId - Optional proxy ID
* @returns Token information including access_token, email, etc.
*/
export async function refreshOpenAIToken(
refreshToken: string,
proxyId?: number | null
): Promise<Record<string, unknown>> {
const payload: { refresh_token: string; proxy_id?: number } = {
refresh_token: refreshToken
}
if (proxyId) {
payload.proxy_id = proxyId
}
const { data } = await apiClient.post<Record<string, unknown>>('/admin/openai/refresh-token', payload)
return data
}
export const accountsAPI = { export const accountsAPI = {
list, list,
getById, getById,
@@ -418,9 +473,11 @@ export const accountsAPI = {
getAvailableModels, getAvailableModels,
generateAuthUrl, generateAuthUrl,
exchangeCode, exchangeCode,
refreshOpenAIToken,
batchCreate, batchCreate,
batchUpdateCredentials, batchUpdateCredentials,
bulkUpdate, bulkUpdate,
previewFromCrs,
syncFromCrs, syncFromCrs,
exportData, exportData,
importData, importData,

View File

@@ -153,6 +153,20 @@ export async function getGroupApiKeys(
return data return data
} }
/**
* Update group sort orders
* @param updates - Array of { id, sort_order } objects
* @returns Success confirmation
*/
export async function updateSortOrder(
updates: Array<{ id: number; sort_order: number }>
): Promise<{ message: string }> {
const { data } = await apiClient.put<{ message: string }>('/admin/groups/sort-order', {
updates
})
return data
}
export const groupsAPI = { export const groupsAPI = {
list, list,
getAll, getAll,
@@ -163,7 +177,8 @@ export const groupsAPI = {
delete: deleteGroup, delete: deleteGroup,
toggleStatus, toggleStatus,
getStats, getStats,
getGroupApiKeys getGroupApiKeys,
updateSortOrder
} }
export default groupsAPI export default groupsAPI

View File

@@ -376,7 +376,6 @@ export interface PlatformAvailability {
total_accounts: number total_accounts: number
available_count: number available_count: number
rate_limit_count: number rate_limit_count: number
scope_rate_limit_count?: Record<string, number>
error_count: number error_count: number
} }
@@ -387,7 +386,6 @@ export interface GroupAvailability {
total_accounts: number total_accounts: number
available_count: number available_count: number
rate_limit_count: number rate_limit_count: number
scope_rate_limit_count?: Record<string, number>
error_count: number error_count: number
} }
@@ -402,7 +400,6 @@ export interface AccountAvailability {
is_rate_limited: boolean is_rate_limited: boolean
rate_limit_reset_at?: string rate_limit_reset_at?: string
rate_limit_remaining_sec?: number rate_limit_remaining_sec?: number
scope_rate_limits?: Record<string, number>
is_overloaded: boolean is_overloaded: boolean
overload_until?: string overload_until?: string
overload_remaining_sec?: number overload_remaining_sec?: number

View File

@@ -76,26 +76,6 @@
</div> </div>
</div> </div>
<!-- Scope Rate Limit Indicators (Antigravity) -->
<template v-if="activeScopeRateLimits.length > 0">
<div v-for="item in activeScopeRateLimits" :key="item.scope" class="group relative">
<span
class="inline-flex items-center gap-1 rounded bg-orange-100 px-1.5 py-0.5 text-xs font-medium text-orange-700 dark:bg-orange-900/30 dark:text-orange-400"
>
<Icon name="exclamationTriangle" size="xs" :stroke-width="2" />
{{ formatScopeName(item.scope) }}
</span>
<!-- Tooltip -->
<div
class="pointer-events-none absolute bottom-full left-1/2 z-50 mb-2 -translate-x-1/2 whitespace-nowrap rounded bg-gray-900 px-2 py-1 text-xs text-white opacity-0 transition-opacity group-hover:opacity-100 dark:bg-gray-700"
>
{{ t('admin.accounts.status.scopeRateLimitedUntil', { scope: formatScopeName(item.scope), time: formatTime(item.reset_at) }) }}
<div
class="absolute left-1/2 top-full -translate-x-1/2 border-4 border-transparent border-t-gray-900 dark:border-t-gray-700" ></div>
</div>
</div>
</template>
<!-- Model Rate Limit Indicators (Antigravity OAuth Smart Retry) --> <!-- Model Rate Limit Indicators (Antigravity OAuth Smart Retry) -->
<template v-if="activeModelRateLimits.length > 0"> <template v-if="activeModelRateLimits.length > 0">
<div v-for="item in activeModelRateLimits" :key="item.model" class="group relative"> <div v-for="item in activeModelRateLimits" :key="item.model" class="group relative">
@@ -160,15 +140,6 @@ const isRateLimited = computed(() => {
return new Date(props.account.rate_limit_reset_at) > new Date() return new Date(props.account.rate_limit_reset_at) > new Date()
}) })
// Computed: active scope rate limits (Antigravity)
const activeScopeRateLimits = computed(() => {
const scopeLimits = props.account.scope_rate_limits
if (!scopeLimits) return []
const now = new Date()
return Object.entries(scopeLimits)
.filter(([, info]) => new Date(info.reset_at) > now)
.map(([scope, info]) => ({ scope, reset_at: info.reset_at }))
})
// Computed: active model rate limits (Antigravity OAuth Smart Retry) // Computed: active model rate limits (Antigravity OAuth Smart Retry)
const activeModelRateLimits = computed(() => { const activeModelRateLimits = computed(() => {

View File

@@ -1038,10 +1038,7 @@
</div> </div>
<!-- Custom Error Codes Section --> <!-- Custom Error Codes Section -->
<div <div class="border-t border-gray-200 pt-4 dark:border-dark-600">
v-if="form.platform !== 'gemini'"
class="border-t border-gray-200 pt-4 dark:border-dark-600"
>
<div class="mb-3 flex items-center justify-between"> <div class="mb-3 flex items-center justify-between">
<div> <div>
<label class="input-label mb-0">{{ t('admin.accounts.customErrorCodes') }}</label> <label class="input-label mb-0">{{ t('admin.accounts.customErrorCodes') }}</label>
@@ -1650,10 +1647,12 @@
:show-proxy-warning="form.platform !== 'openai' && !!form.proxy_id" :show-proxy-warning="form.platform !== 'openai' && !!form.proxy_id"
:allow-multiple="form.platform === 'anthropic'" :allow-multiple="form.platform === 'anthropic'"
:show-cookie-option="form.platform === 'anthropic'" :show-cookie-option="form.platform === 'anthropic'"
:show-refresh-token-option="form.platform === 'openai'"
:platform="form.platform" :platform="form.platform"
:show-project-id="geminiOAuthType === 'code_assist'" :show-project-id="geminiOAuthType === 'code_assist'"
@generate-url="handleGenerateUrl" @generate-url="handleGenerateUrl"
@cookie-auth="handleCookieAuth" @cookie-auth="handleCookieAuth"
@validate-refresh-token="handleOpenAIValidateRT"
/> />
</div> </div>
@@ -2010,6 +2009,7 @@ interface OAuthFlowExposed {
oauthState: string oauthState: string
projectId: string projectId: string
sessionKey: string sessionKey: string
refreshToken: string
inputMethod: AuthInputMethod inputMethod: AuthInputMethod
reset: () => void reset: () => void
} }
@@ -2289,9 +2289,9 @@ watch(
watch( watch(
[accountCategory, addMethod, antigravityAccountType], [accountCategory, addMethod, antigravityAccountType],
([category, method, agType]) => { ([category, method, agType]) => {
// Antigravity upstream 类型 // Antigravity upstream 类型(实际创建为 apikey
if (form.platform === 'antigravity' && agType === 'upstream') { if (form.platform === 'antigravity' && agType === 'upstream') {
form.type = 'upstream' form.type = 'apikey'
return return
} }
if (category === 'oauth-based') { if (category === 'oauth-based') {
@@ -2714,7 +2714,8 @@ const handleSubmit = async () => {
submitting.value = true submitting.value = true
try { try {
await createAccountAndFinish(form.platform, 'upstream', credentials) const extra = mixedScheduling.value ? { mixed_scheduling: true } : undefined
await createAccountAndFinish(form.platform, 'apikey', credentials, extra)
} catch (error: any) { } catch (error: any) {
appStore.showError(error.response?.data?.detail || t('admin.accounts.failedToCreate')) appStore.showError(error.response?.data?.detail || t('admin.accounts.failedToCreate'))
} finally { } finally {
@@ -2860,6 +2861,95 @@ const handleOpenAIExchange = async (authCode: string) => {
} }
} }
// OpenAI 手动 RT 批量验证和创建
const handleOpenAIValidateRT = async (refreshTokenInput: string) => {
if (!refreshTokenInput.trim()) return
// Parse multiple refresh tokens (one per line)
const refreshTokens = refreshTokenInput
.split('\n')
.map((rt) => rt.trim())
.filter((rt) => rt)
if (refreshTokens.length === 0) {
openaiOAuth.error.value = t('admin.accounts.oauth.openai.pleaseEnterRefreshToken')
return
}
openaiOAuth.loading.value = true
openaiOAuth.error.value = ''
let successCount = 0
let failedCount = 0
const errors: string[] = []
try {
for (let i = 0; i < refreshTokens.length; i++) {
try {
const tokenInfo = await openaiOAuth.validateRefreshToken(
refreshTokens[i],
form.proxy_id
)
if (!tokenInfo) {
failedCount++
errors.push(`#${i + 1}: ${openaiOAuth.error.value || 'Validation failed'}`)
openaiOAuth.error.value = ''
continue
}
const credentials = openaiOAuth.buildCredentials(tokenInfo)
const extra = openaiOAuth.buildExtraInfo(tokenInfo)
// Generate account name with index for batch
const accountName = refreshTokens.length > 1 ? `${form.name} #${i + 1}` : form.name
await adminAPI.accounts.create({
name: accountName,
notes: form.notes,
platform: 'openai',
type: 'oauth',
credentials,
extra,
proxy_id: form.proxy_id,
concurrency: form.concurrency,
priority: form.priority,
rate_multiplier: form.rate_multiplier,
group_ids: form.group_ids,
expires_at: form.expires_at,
auto_pause_on_expired: autoPauseOnExpired.value
})
successCount++
} catch (error: any) {
failedCount++
const errMsg = error.response?.data?.detail || error.message || 'Unknown error'
errors.push(`#${i + 1}: ${errMsg}`)
}
}
// Show results
if (successCount > 0 && failedCount === 0) {
appStore.showSuccess(
refreshTokens.length > 1
? t('admin.accounts.oauth.batchSuccess', { count: successCount })
: t('admin.accounts.accountCreated')
)
emit('created')
handleClose()
} else if (successCount > 0 && failedCount > 0) {
appStore.showWarning(
t('admin.accounts.oauth.batchPartialSuccess', { success: successCount, failed: failedCount })
)
openaiOAuth.error.value = errors.join('\n')
emit('created')
} else {
openaiOAuth.error.value = errors.join('\n')
appStore.showError(t('admin.accounts.oauth.batchFailed'))
}
} finally {
openaiOAuth.loading.value = false
}
}
// Gemini OAuth 授权码兑换 // Gemini OAuth 授权码兑换
const handleGeminiExchange = async (authCode: string) => { const handleGeminiExchange = async (authCode: string) => {
if (!authCode.trim() || !geminiOAuth.sessionId.value) return if (!authCode.trim() || !geminiOAuth.sessionId.value) return

View File

@@ -364,6 +364,30 @@
</div> </div>
</div> </div>
<!-- Upstream fields (only for upstream type) -->
<div v-if="account.type === 'upstream'" class="space-y-4">
<div>
<label class="input-label">{{ t('admin.accounts.upstream.baseUrl') }}</label>
<input
v-model="editBaseUrl"
type="text"
class="input"
placeholder="https://s.konstants.xyz"
/>
<p class="input-hint">{{ t('admin.accounts.upstream.baseUrlHint') }}</p>
</div>
<div>
<label class="input-label">{{ t('admin.accounts.upstream.apiKey') }}</label>
<input
v-model="editApiKey"
type="password"
class="input font-mono"
placeholder="sk-..."
/>
<p class="input-hint">{{ t('admin.accounts.leaveEmptyToKeep') }}</p>
</div>
</div>
<!-- Antigravity model restriction (applies to all antigravity types) --> <!-- Antigravity model restriction (applies to all antigravity types) -->
<!-- Antigravity 只支持模型映射模式不支持白名单模式 --> <!-- Antigravity 只支持模型映射模式不支持白名单模式 -->
<div v-if="account.platform === 'antigravity'" class="border-t border-gray-200 pt-4 dark:border-dark-600"> <div v-if="account.platform === 'antigravity'" class="border-t border-gray-200 pt-4 dark:border-dark-600">
@@ -1244,6 +1268,9 @@ watch(
} else { } else {
selectedErrorCodes.value = [] selectedErrorCodes.value = []
} }
} else if (newAccount.type === 'upstream' && newAccount.credentials) {
const credentials = newAccount.credentials as Record<string, unknown>
editBaseUrl.value = (credentials.base_url as string) || ''
} else { } else {
const platformDefaultUrl = const platformDefaultUrl =
newAccount.platform === 'openai' newAccount.platform === 'openai'
@@ -1584,6 +1611,22 @@ const handleSubmit = async () => {
return return
} }
updatePayload.credentials = newCredentials
} else if (props.account.type === 'upstream') {
const currentCredentials = (props.account.credentials as Record<string, unknown>) || {}
const newCredentials: Record<string, unknown> = { ...currentCredentials }
newCredentials.base_url = editBaseUrl.value.trim()
if (editApiKey.value.trim()) {
newCredentials.api_key = editApiKey.value.trim()
}
if (!applyTempUnschedConfig(newCredentials)) {
submitting.value = false
return
}
updatePayload.credentials = newCredentials updatePayload.credentials = newCredentials
} else { } else {
// For oauth/setup-token types, only update intercept_warmup_requests if changed // For oauth/setup-token types, only update intercept_warmup_requests if changed

View File

@@ -10,11 +10,11 @@
<h4 class="mb-3 font-semibold text-blue-900 dark:text-blue-200">{{ oauthTitle }}</h4> <h4 class="mb-3 font-semibold text-blue-900 dark:text-blue-200">{{ oauthTitle }}</h4>
<!-- Auth Method Selection --> <!-- Auth Method Selection -->
<div v-if="showCookieOption" class="mb-4"> <div v-if="showMethodSelection" class="mb-4">
<label class="mb-2 block text-sm font-medium text-blue-800 dark:text-blue-300"> <label class="mb-2 block text-sm font-medium text-blue-800 dark:text-blue-300">
{{ methodLabel }} {{ methodLabel }}
</label> </label>
<div class="flex gap-4"> <div class="flex flex-wrap gap-4">
<label class="flex cursor-pointer items-center gap-2"> <label class="flex cursor-pointer items-center gap-2">
<input <input
v-model="inputMethod" v-model="inputMethod"
@@ -26,7 +26,7 @@
t('admin.accounts.oauth.manualAuth') t('admin.accounts.oauth.manualAuth')
}}</span> }}</span>
</label> </label>
<label class="flex cursor-pointer items-center gap-2"> <label v-if="showCookieOption" class="flex cursor-pointer items-center gap-2">
<input <input
v-model="inputMethod" v-model="inputMethod"
type="radio" type="radio"
@@ -37,6 +37,101 @@
t('admin.accounts.oauth.cookieAutoAuth') t('admin.accounts.oauth.cookieAutoAuth')
}}</span> }}</span>
</label> </label>
<label v-if="showRefreshTokenOption" class="flex cursor-pointer items-center gap-2">
<input
v-model="inputMethod"
type="radio"
value="refresh_token"
class="text-blue-600 focus:ring-blue-500"
/>
<span class="text-sm text-blue-900 dark:text-blue-200">{{
t('admin.accounts.oauth.openai.refreshTokenAuth')
}}</span>
</label>
</div>
</div>
<!-- Refresh Token Input (OpenAI only) -->
<div v-if="inputMethod === 'refresh_token'" class="space-y-4">
<div
class="rounded-lg border border-blue-300 bg-white/80 p-4 dark:border-blue-600 dark:bg-gray-800/80"
>
<p class="mb-3 text-sm text-blue-700 dark:text-blue-300">
{{ t('admin.accounts.oauth.openai.refreshTokenDesc') }}
</p>
<!-- Refresh Token Input -->
<div class="mb-4">
<label
class="mb-2 flex items-center gap-2 text-sm font-semibold text-gray-700 dark:text-gray-300"
>
<Icon name="key" size="sm" class="text-blue-500" />
Refresh Token
<span
v-if="parsedRefreshTokenCount > 1"
class="rounded-full bg-blue-500 px-2 py-0.5 text-xs text-white"
>
{{ t('admin.accounts.oauth.keysCount', { count: parsedRefreshTokenCount }) }}
</span>
</label>
<textarea
v-model="refreshTokenInput"
rows="3"
class="input w-full resize-y font-mono text-sm"
:placeholder="t('admin.accounts.oauth.openai.refreshTokenPlaceholder')"
></textarea>
<p
v-if="parsedRefreshTokenCount > 1"
class="mt-1 text-xs text-blue-600 dark:text-blue-400"
>
{{ t('admin.accounts.oauth.batchCreateAccounts', { count: parsedRefreshTokenCount }) }}
</p>
</div>
<!-- Error Message -->
<div
v-if="error"
class="mb-4 rounded-lg border border-red-200 bg-red-50 p-3 dark:border-red-700 dark:bg-red-900/30"
>
<p class="whitespace-pre-line text-sm text-red-600 dark:text-red-400">
{{ error }}
</p>
</div>
<!-- Validate Button -->
<button
type="button"
class="btn btn-primary w-full"
:disabled="loading || !refreshTokenInput.trim()"
@click="handleValidateRefreshToken"
>
<svg
v-if="loading"
class="-ml-1 mr-2 h-4 w-4 animate-spin"
fill="none"
viewBox="0 0 24 24"
>
<circle
class="opacity-25"
cx="12"
cy="12"
r="10"
stroke="currentColor"
stroke-width="4"
></circle>
<path
class="opacity-75"
fill="currentColor"
d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"
></path>
</svg>
<Icon v-else name="sparkles" size="sm" class="mr-2" />
{{
loading
? t('admin.accounts.oauth.openai.validating')
: t('admin.accounts.oauth.openai.validateAndCreate')
}}
</button>
</div> </div>
</div> </div>
@@ -173,7 +268,7 @@
</div> </div>
<!-- Manual Authorization Flow --> <!-- Manual Authorization Flow -->
<div v-else class="space-y-4"> <div v-if="inputMethod === 'manual'" class="space-y-4">
<p class="mb-4 text-sm text-blue-800 dark:text-blue-300"> <p class="mb-4 text-sm text-blue-800 dark:text-blue-300">
{{ oauthFollowSteps }} {{ oauthFollowSteps }}
</p> </p>
@@ -428,6 +523,7 @@ interface Props {
allowMultiple?: boolean allowMultiple?: boolean
methodLabel?: string methodLabel?: string
showCookieOption?: boolean // Whether to show cookie auto-auth option showCookieOption?: boolean // Whether to show cookie auto-auth option
showRefreshTokenOption?: boolean // Whether to show refresh token input option (OpenAI only)
platform?: 'anthropic' | 'openai' | 'gemini' | 'antigravity' // Platform type for different UI/text platform?: 'anthropic' | 'openai' | 'gemini' | 'antigravity' // Platform type for different UI/text
showProjectId?: boolean // New prop to control project ID visibility showProjectId?: boolean // New prop to control project ID visibility
} }
@@ -442,6 +538,7 @@ const props = withDefaults(defineProps<Props>(), {
allowMultiple: false, allowMultiple: false,
methodLabel: 'Authorization Method', methodLabel: 'Authorization Method',
showCookieOption: true, showCookieOption: true,
showRefreshTokenOption: false,
platform: 'anthropic', platform: 'anthropic',
showProjectId: true showProjectId: true
}) })
@@ -450,6 +547,7 @@ const emit = defineEmits<{
'generate-url': [] 'generate-url': []
'exchange-code': [code: string] 'exchange-code': [code: string]
'cookie-auth': [sessionKey: string] 'cookie-auth': [sessionKey: string]
'validate-refresh-token': [refreshToken: string]
'update:inputMethod': [method: AuthInputMethod] 'update:inputMethod': [method: AuthInputMethod]
}>() }>()
@@ -487,10 +585,14 @@ const oauthImportantNotice = computed(() => {
const inputMethod = ref<AuthInputMethod>(props.showCookieOption ? 'manual' : 'manual') const inputMethod = ref<AuthInputMethod>(props.showCookieOption ? 'manual' : 'manual')
const authCodeInput = ref('') const authCodeInput = ref('')
const sessionKeyInput = ref('') const sessionKeyInput = ref('')
const refreshTokenInput = ref('')
const showHelpDialog = ref(false) const showHelpDialog = ref(false)
const oauthState = ref('') const oauthState = ref('')
const projectId = ref('') const projectId = ref('')
// Computed: show method selection when either cookie or refresh token option is enabled
const showMethodSelection = computed(() => props.showCookieOption || props.showRefreshTokenOption)
// Clipboard // Clipboard
const { copied, copyToClipboard } = useClipboard() const { copied, copyToClipboard } = useClipboard()
@@ -502,6 +604,14 @@ const parsedKeyCount = computed(() => {
.filter((k) => k).length .filter((k) => k).length
}) })
// Computed: count of refresh tokens entered
const parsedRefreshTokenCount = computed(() => {
return refreshTokenInput.value
.split('\n')
.map((rt) => rt.trim())
.filter((rt) => rt).length
})
// Watchers // Watchers
watch(inputMethod, (newVal) => { watch(inputMethod, (newVal) => {
emit('update:inputMethod', newVal) emit('update:inputMethod', newVal)
@@ -563,18 +673,26 @@ const handleCookieAuth = () => {
} }
} }
const handleValidateRefreshToken = () => {
if (refreshTokenInput.value.trim()) {
emit('validate-refresh-token', refreshTokenInput.value.trim())
}
}
// Expose methods and state // Expose methods and state
defineExpose({ defineExpose({
authCode: authCodeInput, authCode: authCodeInput,
oauthState, oauthState,
projectId, projectId,
sessionKey: sessionKeyInput, sessionKey: sessionKeyInput,
refreshToken: refreshTokenInput,
inputMethod, inputMethod,
reset: () => { reset: () => {
authCodeInput.value = '' authCodeInput.value = ''
oauthState.value = '' oauthState.value = ''
projectId.value = '' projectId.value = ''
sessionKeyInput.value = '' sessionKeyInput.value = ''
refreshTokenInput.value = ''
inputMethod.value = 'manual' inputMethod.value = 'manual'
showHelpDialog.value = false showHelpDialog.value = false
} }

View File

@@ -6,15 +6,20 @@
close-on-click-outside close-on-click-outside
@close="handleClose" @close="handleClose"
> >
<form id="sync-from-crs-form" class="space-y-4" @submit.prevent="handleSync"> <!-- Step 1: Input credentials -->
<form
v-if="currentStep === 'input'"
id="sync-from-crs-form"
class="space-y-4"
@submit.prevent="handlePreview"
>
<div class="text-sm text-gray-600 dark:text-dark-300"> <div class="text-sm text-gray-600 dark:text-dark-300">
{{ t('admin.accounts.syncFromCrsDesc') }} {{ t('admin.accounts.syncFromCrsDesc') }}
</div> </div>
<div <div
class="rounded-lg bg-gray-50 p-3 text-xs text-gray-500 dark:bg-dark-700/60 dark:text-dark-400" class="rounded-lg bg-gray-50 p-3 text-xs text-gray-500 dark:bg-dark-700/60 dark:text-dark-400"
> >
已有账号仅同步 CRS {{ t('admin.accounts.crsUpdateBehaviorNote') }}
返回的字段缺失字段保持原值凭据按键合并不会清空未下发的键未勾选"同步代理"时保留原有代理
</div> </div>
<div <div
class="rounded-lg border border-amber-200 bg-amber-50 p-3 text-xs text-amber-600 dark:border-amber-800 dark:bg-amber-900/20 dark:text-amber-400" class="rounded-lg border border-amber-200 bg-amber-50 p-3 text-xs text-amber-600 dark:border-amber-800 dark:bg-amber-900/20 dark:text-amber-400"
@@ -24,26 +29,30 @@
<div class="grid grid-cols-1 gap-4"> <div class="grid grid-cols-1 gap-4">
<div> <div>
<label class="input-label">{{ t('admin.accounts.crsBaseUrl') }}</label> <label for="crs-base-url" class="input-label">{{ t('admin.accounts.crsBaseUrl') }}</label>
<input <input
id="crs-base-url"
v-model="form.base_url" v-model="form.base_url"
type="text" type="text"
class="input" class="input"
required
:placeholder="t('admin.accounts.crsBaseUrlPlaceholder')" :placeholder="t('admin.accounts.crsBaseUrlPlaceholder')"
/> />
</div> </div>
<div class="grid grid-cols-1 gap-4 sm:grid-cols-2"> <div class="grid grid-cols-1 gap-4 sm:grid-cols-2">
<div> <div>
<label class="input-label">{{ t('admin.accounts.crsUsername') }}</label> <label for="crs-username" class="input-label">{{ t('admin.accounts.crsUsername') }}</label>
<input v-model="form.username" type="text" class="input" autocomplete="username" /> <input id="crs-username" v-model="form.username" type="text" class="input" required autocomplete="username" />
</div> </div>
<div> <div>
<label class="input-label">{{ t('admin.accounts.crsPassword') }}</label> <label for="crs-password" class="input-label">{{ t('admin.accounts.crsPassword') }}</label>
<input <input
id="crs-password"
v-model="form.password" v-model="form.password"
type="password" type="password"
class="input" class="input"
required
autocomplete="current-password" autocomplete="current-password"
/> />
</div> </div>
@@ -58,9 +67,101 @@
{{ t('admin.accounts.syncProxies') }} {{ t('admin.accounts.syncProxies') }}
</label> </label>
</div> </div>
</form>
<!-- Step 2: Preview & select -->
<div v-else-if="currentStep === 'preview' && previewResult" class="space-y-4">
<!-- Existing accounts (read-only info) -->
<div
v-if="previewResult.existing_accounts.length"
class="rounded-lg bg-gray-50 p-3 dark:bg-dark-700/60"
>
<div class="mb-2 text-sm font-medium text-gray-700 dark:text-dark-300">
{{ t('admin.accounts.crsExistingAccounts') }}
<span class="ml-1 text-xs text-gray-400">({{ previewResult.existing_accounts.length }})</span>
</div>
<div class="max-h-32 overflow-auto text-xs text-gray-500 dark:text-dark-400">
<div
v-for="acc in previewResult.existing_accounts"
:key="acc.crs_account_id"
class="flex items-center gap-2 py-0.5"
>
<span
class="inline-block rounded bg-blue-100 px-1.5 py-0.5 text-[10px] font-medium text-blue-700 dark:bg-blue-900/30 dark:text-blue-400"
>{{ acc.platform }} / {{ acc.type }}</span>
<span class="truncate">{{ acc.name }}</span>
</div>
</div>
</div>
<!-- New accounts (selectable) -->
<div v-if="previewResult.new_accounts.length">
<div class="mb-2 flex items-center justify-between">
<div class="text-sm font-medium text-gray-900 dark:text-white">
{{ t('admin.accounts.crsNewAccounts') }}
<span class="ml-1 text-xs text-gray-400">({{ previewResult.new_accounts.length }})</span>
</div>
<div class="flex gap-2">
<button
type="button"
class="text-xs text-blue-600 hover:text-blue-700 dark:text-blue-400"
@click="selectAll"
>{{ t('admin.accounts.crsSelectAll') }}</button>
<button
type="button"
class="text-xs text-gray-500 hover:text-gray-600 dark:text-gray-400"
@click="selectNone"
>{{ t('admin.accounts.crsSelectNone') }}</button>
</div>
</div>
<div
class="max-h-48 overflow-auto rounded-lg border border-gray-200 p-2 dark:border-dark-600"
>
<label
v-for="acc in previewResult.new_accounts"
:key="acc.crs_account_id"
class="flex cursor-pointer items-center gap-2 rounded px-2 py-1.5 hover:bg-gray-50 dark:hover:bg-dark-700/40"
>
<input
type="checkbox"
:checked="selectedIds.has(acc.crs_account_id)"
class="rounded border-gray-300 dark:border-dark-600"
@change="toggleSelect(acc.crs_account_id)"
/>
<span
class="inline-block rounded bg-green-100 px-1.5 py-0.5 text-[10px] font-medium text-green-700 dark:bg-green-900/30 dark:text-green-400"
>{{ acc.platform }} / {{ acc.type }}</span>
<span class="truncate text-sm text-gray-700 dark:text-dark-300">{{ acc.name }}</span>
</label>
</div>
<div class="mt-1 text-xs text-gray-400">
{{ t('admin.accounts.crsSelectedCount', { count: selectedIds.size }) }}
</div>
</div>
<!-- Sync options summary -->
<div class="flex items-center gap-2 text-xs text-gray-500 dark:text-dark-400">
<span>{{ t('admin.accounts.syncProxies') }}:</span>
<span :class="form.sync_proxies ? 'text-green-600 dark:text-green-400' : 'text-gray-400 dark:text-dark-500'">
{{ form.sync_proxies ? t('common.yes') : t('common.no') }}
</span>
</div>
<!-- No new accounts -->
<div
v-if="!previewResult.new_accounts.length"
class="rounded-lg bg-gray-50 p-4 text-center text-sm text-gray-500 dark:bg-dark-700/60 dark:text-dark-400"
>
{{ t('admin.accounts.crsNoNewAccounts') }}
<span v-if="previewResult.existing_accounts.length">
{{ t('admin.accounts.crsWillUpdate', { count: previewResult.existing_accounts.length }) }}
</span>
</div>
</div>
<!-- Step 3: Result -->
<div v-else-if="currentStep === 'result' && result" class="space-y-4">
<div <div
v-if="result"
class="space-y-2 rounded-xl border border-gray-200 p-4 dark:border-dark-700" class="space-y-2 rounded-xl border border-gray-200 p-4 dark:border-dark-700"
> >
<div class="text-sm font-medium text-gray-900 dark:text-white"> <div class="text-sm font-medium text-gray-900 dark:text-white">
@@ -84,21 +185,56 @@
</div> </div>
</div> </div>
</div> </div>
</form> </div>
<template #footer> <template #footer>
<div class="flex justify-end gap-3"> <div class="flex justify-end gap-3">
<button class="btn btn-secondary" type="button" :disabled="syncing" @click="handleClose"> <!-- Step 1: Input -->
{{ t('common.cancel') }} <template v-if="currentStep === 'input'">
</button> <button
<button class="btn btn-secondary"
class="btn btn-primary" type="button"
type="submit" :disabled="previewing"
form="sync-from-crs-form" @click="handleClose"
:disabled="syncing" >
> {{ t('common.cancel') }}
{{ syncing ? t('admin.accounts.syncing') : t('admin.accounts.syncNow') }} </button>
</button> <button
class="btn btn-primary"
type="submit"
form="sync-from-crs-form"
:disabled="previewing"
>
{{ previewing ? t('admin.accounts.crsPreviewing') : t('admin.accounts.crsPreview') }}
</button>
</template>
<!-- Step 2: Preview -->
<template v-else-if="currentStep === 'preview'">
<button
class="btn btn-secondary"
type="button"
:disabled="syncing"
@click="handleBack"
>
{{ t('admin.accounts.crsBack') }}
</button>
<button
class="btn btn-primary"
type="button"
:disabled="syncing || hasNewButNoneSelected"
@click="handleSync"
>
{{ syncing ? t('admin.accounts.syncing') : t('admin.accounts.syncNow') }}
</button>
</template>
<!-- Step 3: Result -->
<template v-else-if="currentStep === 'result'">
<button class="btn btn-secondary" type="button" @click="handleClose">
{{ t('common.close') }}
</button>
</template>
</div> </div>
</template> </template>
</BaseDialog> </BaseDialog>
@@ -110,6 +246,7 @@ import { useI18n } from 'vue-i18n'
import BaseDialog from '@/components/common/BaseDialog.vue' import BaseDialog from '@/components/common/BaseDialog.vue'
import { useAppStore } from '@/stores/app' import { useAppStore } from '@/stores/app'
import { adminAPI } from '@/api/admin' import { adminAPI } from '@/api/admin'
import type { PreviewFromCRSResult } from '@/api/admin/accounts'
interface Props { interface Props {
show: boolean show: boolean
@@ -126,7 +263,12 @@ const emit = defineEmits<Emits>()
const { t } = useI18n() const { t } = useI18n()
const appStore = useAppStore() const appStore = useAppStore()
type Step = 'input' | 'preview' | 'result'
const currentStep = ref<Step>('input')
const previewing = ref(false)
const syncing = ref(false) const syncing = ref(false)
const previewResult = ref<PreviewFromCRSResult | null>(null)
const selectedIds = ref(new Set<string>())
const result = ref<Awaited<ReturnType<typeof adminAPI.accounts.syncFromCrs>> | null>(null) const result = ref<Awaited<ReturnType<typeof adminAPI.accounts.syncFromCrs>> | null>(null)
const form = reactive({ const form = reactive({
@@ -136,28 +278,90 @@ const form = reactive({
sync_proxies: true sync_proxies: true
}) })
const hasNewButNoneSelected = computed(() => {
if (!previewResult.value) return false
return previewResult.value.new_accounts.length > 0 && selectedIds.value.size === 0
})
const errorItems = computed(() => { const errorItems = computed(() => {
if (!result.value?.items) return [] if (!result.value?.items) return []
return result.value.items.filter((i) => i.action === 'failed' || i.action === 'skipped') return result.value.items.filter(
(i) => i.action === 'failed' || (i.action === 'skipped' && i.error !== 'not selected')
)
}) })
watch( watch(
() => props.show, () => props.show,
(open) => { (open) => {
if (open) { if (open) {
currentStep.value = 'input'
previewResult.value = null
selectedIds.value = new Set()
result.value = null result.value = null
form.base_url = ''
form.username = ''
form.password = ''
form.sync_proxies = true
} }
} }
) )
const handleClose = () => { const handleClose = () => {
// 防止在同步进行中关闭对话框 if (syncing.value || previewing.value) {
if (syncing.value) {
return return
} }
emit('close') emit('close')
} }
const handleBack = () => {
currentStep.value = 'input'
previewResult.value = null
selectedIds.value = new Set()
}
const selectAll = () => {
if (!previewResult.value) return
selectedIds.value = new Set(previewResult.value.new_accounts.map((a) => a.crs_account_id))
}
const selectNone = () => {
selectedIds.value = new Set()
}
const toggleSelect = (id: string) => {
const s = new Set(selectedIds.value)
if (s.has(id)) {
s.delete(id)
} else {
s.add(id)
}
selectedIds.value = s
}
const handlePreview = async () => {
if (!form.base_url.trim() || !form.username.trim() || !form.password.trim()) {
appStore.showError(t('admin.accounts.syncMissingFields'))
return
}
previewing.value = true
try {
const res = await adminAPI.accounts.previewFromCrs({
base_url: form.base_url.trim(),
username: form.username.trim(),
password: form.password
})
previewResult.value = res
// Auto-select all new accounts
selectedIds.value = new Set(res.new_accounts.map((a) => a.crs_account_id))
currentStep.value = 'preview'
} catch (error: any) {
appStore.showError(error?.message || t('admin.accounts.crsPreviewFailed'))
} finally {
previewing.value = false
}
}
const handleSync = async () => { const handleSync = async () => {
if (!form.base_url.trim() || !form.username.trim() || !form.password.trim()) { if (!form.base_url.trim() || !form.username.trim() || !form.password.trim()) {
appStore.showError(t('admin.accounts.syncMissingFields')) appStore.showError(t('admin.accounts.syncMissingFields'))
@@ -170,16 +374,18 @@ const handleSync = async () => {
base_url: form.base_url.trim(), base_url: form.base_url.trim(),
username: form.username.trim(), username: form.username.trim(),
password: form.password, password: form.password,
sync_proxies: form.sync_proxies sync_proxies: form.sync_proxies,
selected_account_ids: [...selectedIds.value]
}) })
result.value = res result.value = res
currentStep.value = 'result'
if (res.failed > 0) { if (res.failed > 0) {
appStore.showError(t('admin.accounts.syncCompletedWithErrors', res)) appStore.showError(t('admin.accounts.syncCompletedWithErrors', res))
} else { } else {
appStore.showSuccess(t('admin.accounts.syncCompleted', res)) appStore.showSuccess(t('admin.accounts.syncCompleted', res))
emit('synced')
} }
emit('synced')
} catch (error: any) { } catch (error: any) {
appStore.showError(error?.message || t('admin.accounts.syncFailed')) appStore.showError(error?.message || t('admin.accounts.syncFailed'))
} finally { } finally {

View File

@@ -22,6 +22,7 @@
/> />
<GroupBadge <GroupBadge
:name="group.name" :name="group.name"
:platform="group.platform"
:subscription-type="group.subscription_type" :subscription-type="group.subscription_type"
:rate-multiplier="group.rate_multiplier" :rate-multiplier="group.rate_multiplier"
class="min-w-0 flex-1" class="min-w-0 flex-1"

View File

@@ -58,6 +58,7 @@ const icons = {
arrowLeft: 'M10.5 19.5L3 12m0 0l7.5-7.5M3 12h18', arrowLeft: 'M10.5 19.5L3 12m0 0l7.5-7.5M3 12h18',
arrowUp: 'M5 10l7-7m0 0l7 7m-7-7v18', arrowUp: 'M5 10l7-7m0 0l7 7m-7-7v18',
arrowDown: 'M19 14l-7 7m0 0l-7-7m7 7V3', arrowDown: 'M19 14l-7 7m0 0l-7-7m7 7V3',
arrowsUpDown: 'M3 7.5L7.5 3m0 0L12 7.5M7.5 3v13.5m13.5 0L16.5 21m0 0L12 16.5m4.5 4.5V7.5',
chevronUp: 'M5 15l7-7 7 7', chevronUp: 'M5 15l7-7 7 7',
externalLink: 'M10 6H6a2 2 0 00-2 2v10a2 2 0 002 2h10a2 2 0 002-2v-4M14 4h6m0 0v6m0-6L10 14', externalLink: 'M10 6H6a2 2 0 00-2 2v10a2 2 0 002 2h10a2 2 0 002-2v-4M14 4h6m0 0v6m0-6L10 14',

View File

@@ -0,0 +1,43 @@
<template>
<div class="flex items-center">
<span
:class="[
'inline-flex items-center gap-1 rounded-md px-2 py-0.5 text-xs font-medium',
statusClass
]"
>
<!-- Four-square grid icon -->
<svg class="h-3 w-3" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
<path stroke-linecap="round" stroke-linejoin="round" d="M3.75 6A2.25 2.25 0 016 3.75h2.25A2.25 2.25 0 0110.5 6v2.25a2.25 2.25 0 01-2.25 2.25H6a2.25 2.25 0 01-2.25-2.25V6zM3.75 15.75A2.25 2.25 0 016 13.5h2.25a2.25 2.25 0 012.25 2.25V18a2.25 2.25 0 01-2.25 2.25H6A2.25 2.25 0 013.75 18v-2.25zM13.5 6a2.25 2.25 0 012.25-2.25H18A2.25 2.25 0 0120.25 6v2.25A2.25 2.25 0 0118 10.5h-2.25a2.25 2.25 0 01-2.25-2.25V6zM13.5 15.75a2.25 2.25 0 012.25-2.25H18a2.25 2.25 0 012.25 2.25V18A2.25 2.25 0 0118 20.25h-2.25A2.25 2.25 0 0113.5 18v-2.25z" />
</svg>
<span class="font-mono">{{ current }}</span>
<span class="text-gray-400 dark:text-gray-500">/</span>
<span class="font-mono">{{ max }}</span>
</span>
</div>
</template>
<script setup lang="ts">
import { computed } from 'vue'
const props = defineProps<{
current: number
max: number
}>()
// Status color based on usage
const statusClass = computed(() => {
const { current, max } = props
// Full: red
if (current >= max && max > 0) {
return 'bg-red-100 text-red-700 dark:bg-red-900/30 dark:text-red-400'
}
// In use: yellow
if (current > 0) {
return 'bg-yellow-100 text-yellow-700 dark:bg-yellow-900/30 dark:text-yellow-400'
}
// Idle: gray
return 'bg-gray-100 text-gray-600 dark:bg-gray-800 dark:text-gray-400'
})
</script>

View File

@@ -3,7 +3,7 @@ import { useAppStore } from '@/stores/app'
import { adminAPI } from '@/api/admin' import { adminAPI } from '@/api/admin'
export type AddMethod = 'oauth' | 'setup-token' export type AddMethod = 'oauth' | 'setup-token'
export type AuthInputMethod = 'manual' | 'cookie' export type AuthInputMethod = 'manual' | 'cookie' | 'refresh_token'
export interface OAuthState { export interface OAuthState {
authUrl: string authUrl: string

View File

@@ -105,6 +105,32 @@ export function useOpenAIOAuth() {
} }
} }
// Validate refresh token and get full token info
const validateRefreshToken = async (
refreshToken: string,
proxyId?: number | null
): Promise<OpenAITokenInfo | null> => {
if (!refreshToken.trim()) {
error.value = 'Missing refresh token'
return null
}
loading.value = true
error.value = ''
try {
// Use dedicated refresh-token endpoint
const tokenInfo = await adminAPI.accounts.refreshOpenAIToken(refreshToken.trim(), proxyId)
return tokenInfo as OpenAITokenInfo
} catch (err: any) {
error.value = err.response?.data?.detail || 'Failed to validate refresh token'
appStore.showError(error.value)
return null
} finally {
loading.value = false
}
}
// Build credentials for OpenAI OAuth account // Build credentials for OpenAI OAuth account
const buildCredentials = (tokenInfo: OpenAITokenInfo): Record<string, unknown> => { const buildCredentials = (tokenInfo: OpenAITokenInfo): Record<string, unknown> => {
const creds: Record<string, unknown> = { const creds: Record<string, unknown> = {
@@ -152,6 +178,7 @@ export function useOpenAIOAuth() {
resetState, resetState,
generateAuthUrl, generateAuthUrl,
exchangeAuthCode, exchangeAuthCode,
validateRefreshToken,
buildCredentials, buildCredentials,
buildExtraInfo buildExtraInfo
} }

View File

@@ -1042,6 +1042,10 @@ export default {
createGroup: 'Create Group', createGroup: 'Create Group',
editGroup: 'Edit Group', editGroup: 'Edit Group',
deleteGroup: 'Delete Group', deleteGroup: 'Delete Group',
sortOrder: 'Sort',
sortOrderHint: 'Drag groups to adjust display order, groups at the top will be displayed first',
sortOrderUpdated: 'Sort order updated',
failedToUpdateSortOrder: 'Failed to update sort order',
allPlatforms: 'All Platforms', allPlatforms: 'All Platforms',
allStatus: 'All Status', allStatus: 'All Status',
allGroups: 'All Groups', allGroups: 'All Groups',
@@ -1305,10 +1309,23 @@ export default {
syncResult: 'Sync Result', syncResult: 'Sync Result',
syncResultSummary: 'Created {created}, updated {updated}, skipped {skipped}, failed {failed}', syncResultSummary: 'Created {created}, updated {updated}, skipped {skipped}, failed {failed}',
syncErrors: 'Errors / Skipped Details', syncErrors: 'Errors / Skipped Details',
syncCompleted: 'Sync completed: created {created}, updated {updated}', syncCompleted: 'Sync completed: created {created}, updated {updated}, skipped {skipped}',
syncCompletedWithErrors: syncCompletedWithErrors:
'Sync completed with errors: failed {failed} (created {created}, updated {updated})', 'Sync completed with errors: failed {failed} (created {created}, updated {updated}, skipped {skipped})',
syncFailed: 'Sync failed', syncFailed: 'Sync failed',
crsPreview: 'Preview',
crsPreviewing: 'Previewing...',
crsPreviewFailed: 'Preview failed',
crsExistingAccounts: 'Existing accounts (will be updated)',
crsNewAccounts: 'New accounts (select to sync)',
crsSelectAll: 'Select all',
crsSelectNone: 'Select none',
crsNoNewAccounts: 'All CRS accounts are already synced.',
crsWillUpdate: 'Will update {count} existing accounts.',
crsSelectedCount: '{count} new accounts selected',
crsUpdateBehaviorNote:
'Existing accounts only sync fields returned by CRS; missing fields keep their current values. Credentials are merged by key — keys not returned by CRS are preserved. Proxies are kept when "Sync proxies" is unchecked.',
crsBack: 'Back',
editAccount: 'Edit Account', editAccount: 'Edit Account',
deleteAccount: 'Delete Account', deleteAccount: 'Delete Account',
searchAccounts: 'Search accounts...', searchAccounts: 'Search accounts...',
@@ -1356,7 +1373,6 @@ export default {
overloaded: 'Overloaded', overloaded: 'Overloaded',
tempUnschedulable: 'Temp Unschedulable', tempUnschedulable: 'Temp Unschedulable',
rateLimitedUntil: 'Rate limited until {time}', rateLimitedUntil: 'Rate limited until {time}',
scopeRateLimitedUntil: '{scope} rate limited until {time}',
modelRateLimitedUntil: '{model} rate limited until {time}', modelRateLimitedUntil: '{model} rate limited until {time}',
overloadedUntil: 'Overloaded until {time}', overloadedUntil: 'Overloaded until {time}',
viewTempUnschedDetails: 'View temp unschedulable details' viewTempUnschedDetails: 'View temp unschedulable details'
@@ -1662,6 +1678,9 @@ export default {
cookieAuthFailed: 'Cookie authorization failed', cookieAuthFailed: 'Cookie authorization failed',
keyAuthFailed: 'Key {index}: {error}', keyAuthFailed: 'Key {index}: {error}',
successCreated: 'Successfully created {count} account(s)', successCreated: 'Successfully created {count} account(s)',
batchSuccess: 'Successfully created {count} account(s)',
batchPartialSuccess: 'Partial success: {success} succeeded, {failed} failed',
batchFailed: 'Batch creation failed',
// OpenAI specific // OpenAI specific
openai: { openai: {
title: 'OpenAI Account Authorization', title: 'OpenAI Account Authorization',
@@ -1680,7 +1699,14 @@ export default {
authCodePlaceholder: authCodePlaceholder:
'Option 1: Copy the complete URL\n(http://localhost:xxx/auth/callback?code=...)\nOption 2: Copy only the code parameter value', 'Option 1: Copy the complete URL\n(http://localhost:xxx/auth/callback?code=...)\nOption 2: Copy only the code parameter value',
authCodeHint: authCodeHint:
'You can copy the entire URL or just the code parameter value, the system will auto-detect' 'You can copy the entire URL or just the code parameter value, the system will auto-detect',
// Refresh Token auth
refreshTokenAuth: 'Manual RT Input',
refreshTokenDesc: 'Enter your existing OpenAI Refresh Token(s). Supports batch input (one per line). The system will automatically validate and create accounts.',
refreshTokenPlaceholder: 'Paste your OpenAI Refresh Token...\nSupports multiple, one per line',
validating: 'Validating...',
validateAndCreate: 'Validate & Create Account',
pleaseEnterRefreshToken: 'Please enter Refresh Token'
}, },
// Gemini specific // Gemini specific
gemini: { gemini: {
@@ -3049,7 +3075,6 @@ export default {
empty: 'No data', empty: 'No data',
queued: 'Queue {count}', queued: 'Queue {count}',
rateLimited: 'Rate-limited {count}', rateLimited: 'Rate-limited {count}',
scopeRateLimitedTooltip: '{scope} rate-limited ({count} accounts)',
errorAccounts: 'Errors {count}', errorAccounts: 'Errors {count}',
loadFailed: 'Failed to load concurrency data' loadFailed: 'Failed to load concurrency data'
}, },

View File

@@ -1099,6 +1099,10 @@ export default {
createGroup: '创建分组', createGroup: '创建分组',
editGroup: '编辑分组', editGroup: '编辑分组',
deleteGroup: '删除分组', deleteGroup: '删除分组',
sortOrder: '排序',
sortOrderHint: '拖拽分组调整显示顺序,排在前面的分组会优先显示',
sortOrderUpdated: '排序已更新',
failedToUpdateSortOrder: '更新排序失败',
deleteConfirm: "确定要删除分组 '{name}' 吗?所有关联的 API 密钥将不再属于任何分组。", deleteConfirm: "确定要删除分组 '{name}' 吗?所有关联的 API 密钥将不再属于任何分组。",
deleteConfirmSubscription: deleteConfirmSubscription:
"确定要删除订阅分组 '{name}' 吗?此操作会让所有绑定此订阅的用户的 API Key 失效,并删除所有相关的订阅记录。此操作无法撤销。", "确定要删除订阅分组 '{name}' 吗?此操作会让所有绑定此订阅的用户的 API Key 失效,并删除所有相关的订阅记录。此操作无法撤销。",
@@ -1393,9 +1397,22 @@ export default {
syncResult: '同步结果', syncResult: '同步结果',
syncResultSummary: '创建 {created},更新 {updated},跳过 {skipped},失败 {failed}', syncResultSummary: '创建 {created},更新 {updated},跳过 {skipped},失败 {failed}',
syncErrors: '错误/跳过详情', syncErrors: '错误/跳过详情',
syncCompleted: '同步完成:创建 {created},更新 {updated}', syncCompleted: '同步完成:创建 {created},更新 {updated},跳过 {skipped}',
syncCompletedWithErrors: '同步完成但有错误:失败 {failed}(创建 {created},更新 {updated}', syncCompletedWithErrors: '同步完成但有错误:失败 {failed}(创建 {created},更新 {updated},跳过 {skipped}',
syncFailed: '同步失败', syncFailed: '同步失败',
crsPreview: '预览',
crsPreviewing: '预览中...',
crsPreviewFailed: '预览失败',
crsExistingAccounts: '将自动更新的已有账号',
crsNewAccounts: '新账号(可选择)',
crsSelectAll: '全选',
crsSelectNone: '全不选',
crsNoNewAccounts: '所有 CRS 账号均已同步。',
crsWillUpdate: '将更新 {count} 个已有账号。',
crsSelectedCount: '已选择 {count} 个新账号',
crsUpdateBehaviorNote:
'已有账号仅同步 CRS 返回的字段,缺失字段保持原值;凭据按键合并,不会清空未下发的键;未勾选"同步代理"时保留原有代理。',
crsBack: '返回',
editAccount: '编辑账号', editAccount: '编辑账号',
deleteAccount: '删除账号', deleteAccount: '删除账号',
deleteConfirmMessage: "确定要删除账号 '{name}' 吗?", deleteConfirmMessage: "确定要删除账号 '{name}' 吗?",
@@ -1492,7 +1509,6 @@ export default {
overloaded: '过载中', overloaded: '过载中',
tempUnschedulable: '临时不可调度', tempUnschedulable: '临时不可调度',
rateLimitedUntil: '限流中,重置时间:{time}', rateLimitedUntil: '限流中,重置时间:{time}',
scopeRateLimitedUntil: '{scope} 限流中,重置时间:{time}',
modelRateLimitedUntil: '{model} 限流至 {time}', modelRateLimitedUntil: '{model} 限流至 {time}',
overloadedUntil: '负载过重,重置时间:{time}', overloadedUntil: '负载过重,重置时间:{time}',
viewTempUnschedDetails: '查看临时不可调度详情' viewTempUnschedDetails: '查看临时不可调度详情'
@@ -1804,6 +1820,9 @@ export default {
cookieAuthFailed: 'Cookie 授权失败', cookieAuthFailed: 'Cookie 授权失败',
keyAuthFailed: '密钥 {index}: {error}', keyAuthFailed: '密钥 {index}: {error}',
successCreated: '成功创建 {count} 个账号', successCreated: '成功创建 {count} 个账号',
batchSuccess: '成功创建 {count} 个账号',
batchPartialSuccess: '部分成功:{success} 个成功,{failed} 个失败',
batchFailed: '批量创建失败',
// OpenAI specific // OpenAI specific
openai: { openai: {
title: 'OpenAI 账户授权', title: 'OpenAI 账户授权',
@@ -1820,7 +1839,14 @@ export default {
authCode: '授权链接或 Code', authCode: '授权链接或 Code',
authCodePlaceholder: authCodePlaceholder:
'方式1复制完整的链接\n(http://localhost:xxx/auth/callback?code=...)\n方式2仅复制 code 参数的值', '方式1复制完整的链接\n(http://localhost:xxx/auth/callback?code=...)\n方式2仅复制 code 参数的值',
authCodeHint: '您可以直接复制整个链接或仅复制 code 参数值,系统会自动识别' authCodeHint: '您可以直接复制整个链接或仅复制 code 参数值,系统会自动识别',
// Refresh Token auth
refreshTokenAuth: '手动输入 RT',
refreshTokenDesc: '输入您已有的 OpenAI Refresh Token支持批量输入每行一个系统将自动验证并创建账号。',
refreshTokenPlaceholder: '粘贴您的 OpenAI Refresh Token...\n支持多个每行一个',
validating: '验证中...',
validateAndCreate: '验证并创建账号',
pleaseEnterRefreshToken: '请输入 Refresh Token'
}, },
// Gemini specific // Gemini specific
gemini: { gemini: {
@@ -3222,7 +3248,6 @@ export default {
empty: '暂无数据', empty: '暂无数据',
queued: '队列 {count}', queued: '队列 {count}',
rateLimited: '限流 {count}', rateLimited: '限流 {count}',
scopeRateLimitedTooltip: '{scope} 限流中 ({count} 个账号)',
errorAccounts: '异常 {count}', errorAccounts: '异常 {count}',
loadFailed: '加载并发数据失败' loadFailed: '加载并发数据失败'
}, },

View File

@@ -43,6 +43,8 @@ export interface AdminUser extends User {
notes: string notes: string
// 用户专属分组倍率配置 (group_id -> rate_multiplier) // 用户专属分组倍率配置 (group_id -> rate_multiplier)
group_rates?: Record<number, number> group_rates?: Record<number, number>
// 当前并发数(仅管理员列表接口返回)
current_concurrency?: number
} }
export interface LoginRequest { export interface LoginRequest {
@@ -377,6 +379,9 @@ export interface AdminGroup extends Group {
// 分组下账号数量(仅管理员可见) // 分组下账号数量(仅管理员可见)
account_count?: number account_count?: number
// 分组排序
sort_order: number
} }
export interface ApiKey { export interface ApiKey {
@@ -589,9 +594,6 @@ export interface Account {
temp_unschedulable_until: string | null temp_unschedulable_until: string | null
temp_unschedulable_reason: string | null temp_unschedulable_reason: string | null
// Antigravity scope 级限流状态
scope_rate_limits?: Record<string, { reset_at: string; remaining_sec: number }>
// Session window fields (5-hour window) // Session window fields (5-hour window)
session_window_start: string | null session_window_start: string | null
session_window_end: string | null session_window_end: string | null

View File

@@ -1,26 +1,10 @@
<template> <template>
<AppLayout> <AppLayout>
<TablePageLayout> <TablePageLayout>
<template #actions>
<div class="flex justify-end gap-3">
<button
@click="loadAnnouncements"
:disabled="loading"
class="btn btn-secondary"
:title="t('common.refresh')"
>
<Icon name="refresh" size="md" :class="loading ? 'animate-spin' : ''" />
</button>
<button @click="openCreateDialog" class="btn btn-primary">
<Icon name="plus" size="md" class="mr-1" />
{{ t('admin.announcements.createAnnouncement') }}
</button>
</div>
</template>
<template #filters> <template #filters>
<div class="flex flex-col gap-4 sm:flex-row sm:items-center sm:justify-between"> <div class="flex flex-wrap items-center gap-3">
<div class="max-w-md flex-1"> <!-- Left: Search + Filters -->
<div class="flex-1 sm:max-w-64">
<input <input
v-model="searchQuery" v-model="searchQuery"
type="text" type="text"
@@ -29,13 +13,27 @@
@input="handleSearch" @input="handleSearch"
/> />
</div> </div>
<div class="flex gap-2"> <Select
<Select v-model="filters.status"
v-model="filters.status" :options="statusFilterOptions"
:options="statusFilterOptions" class="w-40"
class="w-40" @change="handleStatusChange"
@change="handleStatusChange" />
/>
<!-- Right: Action buttons -->
<div class="flex flex-1 flex-wrap items-center justify-end gap-2">
<button
@click="loadAnnouncements"
:disabled="loading"
class="btn btn-secondary"
:title="t('common.refresh')"
>
<Icon name="refresh" size="md" :class="loading ? 'animate-spin' : ''" />
</button>
<button @click="openCreateDialog" class="btn btn-primary">
<Icon name="plus" size="md" class="mr-1" />
{{ t('admin.announcements.createAnnouncement') }}
</button>
</div> </div>
</div> </div>
</template> </template>

View File

@@ -52,6 +52,14 @@
> >
<Icon name="refresh" size="md" :class="loading ? 'animate-spin' : ''" /> <Icon name="refresh" size="md" :class="loading ? 'animate-spin' : ''" />
</button> </button>
<button
@click="openSortModal"
class="btn btn-secondary"
:title="t('admin.groups.sortOrder')"
>
<Icon name="arrowsUpDown" size="md" class="mr-2" />
{{ t('admin.groups.sortOrder') }}
</button>
<button <button
@click="showCreateModal = true" @click="showCreateModal = true"
class="btn btn-primary" class="btn btn-primary"
@@ -1455,6 +1463,92 @@
@confirm="confirmDelete" @confirm="confirmDelete"
@cancel="showDeleteDialog = false" @cancel="showDeleteDialog = false"
/> />
<!-- Sort Order Modal -->
<BaseDialog
:show="showSortModal"
:title="t('admin.groups.sortOrder')"
width="normal"
@close="closeSortModal"
>
<div class="space-y-4">
<p class="text-sm text-gray-500 dark:text-gray-400">
{{ t('admin.groups.sortOrderHint') }}
</p>
<VueDraggable
v-model="sortableGroups"
:animation="200"
class="space-y-2"
>
<div
v-for="group in sortableGroups"
:key="group.id"
class="flex cursor-grab items-center gap-3 rounded-lg border border-gray-200 bg-white p-3 transition-shadow hover:shadow-md active:cursor-grabbing dark:border-dark-600 dark:bg-dark-700"
>
<div class="text-gray-400">
<Icon name="menu" size="md" />
</div>
<div class="flex-1">
<div class="font-medium text-gray-900 dark:text-white">{{ group.name }}</div>
<div class="text-xs text-gray-500 dark:text-gray-400">
<span
:class="[
'inline-flex items-center gap-1 rounded-full px-2 py-0.5 text-xs font-medium',
group.platform === 'anthropic'
? 'bg-orange-100 text-orange-700 dark:bg-orange-900/30 dark:text-orange-400'
: group.platform === 'openai'
? 'bg-emerald-100 text-emerald-700 dark:bg-emerald-900/30 dark:text-emerald-400'
: group.platform === 'antigravity'
? 'bg-purple-100 text-purple-700 dark:bg-purple-900/30 dark:text-purple-400'
: 'bg-blue-100 text-blue-700 dark:bg-blue-900/30 dark:text-blue-400'
]"
>
{{ t('admin.groups.platforms.' + group.platform) }}
</span>
</div>
</div>
<div class="text-sm text-gray-400">
#{{ group.id }}
</div>
</div>
</VueDraggable>
</div>
<template #footer>
<div class="flex justify-end gap-3 pt-4">
<button @click="closeSortModal" type="button" class="btn btn-secondary">
{{ t('common.cancel') }}
</button>
<button
@click="saveSortOrder"
:disabled="sortSubmitting"
class="btn btn-primary"
>
<svg
v-if="sortSubmitting"
class="-ml-1 mr-2 h-4 w-4 animate-spin"
fill="none"
viewBox="0 0 24 24"
>
<circle
class="opacity-25"
cx="12"
cy="12"
r="10"
stroke="currentColor"
stroke-width="4"
></circle>
<path
class="opacity-75"
fill="currentColor"
d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"
></path>
</svg>
{{ sortSubmitting ? t('common.saving') : t('common.save') }}
</button>
</div>
</template>
</BaseDialog>
</AppLayout> </AppLayout>
</template> </template>
@@ -1476,6 +1570,7 @@ import EmptyState from '@/components/common/EmptyState.vue'
import Select from '@/components/common/Select.vue' import Select from '@/components/common/Select.vue'
import PlatformIcon from '@/components/common/PlatformIcon.vue' import PlatformIcon from '@/components/common/PlatformIcon.vue'
import Icon from '@/components/icons/Icon.vue' import Icon from '@/components/icons/Icon.vue'
import { VueDraggable } from 'vue-draggable-plus'
const { t } = useI18n() const { t } = useI18n()
const appStore = useAppStore() const appStore = useAppStore()
@@ -1640,9 +1735,12 @@ let abortController: AbortController | null = null
const showCreateModal = ref(false) const showCreateModal = ref(false)
const showEditModal = ref(false) const showEditModal = ref(false)
const showDeleteDialog = ref(false) const showDeleteDialog = ref(false)
const showSortModal = ref(false)
const submitting = ref(false) const submitting = ref(false)
const sortSubmitting = ref(false)
const editingGroup = ref<AdminGroup | null>(null) const editingGroup = ref<AdminGroup | null>(null)
const deletingGroup = ref<AdminGroup | null>(null) const deletingGroup = ref<AdminGroup | null>(null)
const sortableGroups = ref<AdminGroup[]>([])
const createForm = reactive({ const createForm = reactive({
name: '', name: '',
@@ -2101,6 +2199,46 @@ const handleClickOutside = (event: MouseEvent) => {
} }
} }
// 打开排序弹窗
const openSortModal = async () => {
try {
// 获取所有分组(不分页)
const allGroups = await adminAPI.groups.getAll()
// 按 sort_order 排序
sortableGroups.value = [...allGroups].sort((a, b) => a.sort_order - b.sort_order)
showSortModal.value = true
} catch (error) {
appStore.showError(t('admin.groups.failedToLoad'))
console.error('Error loading groups for sorting:', error)
}
}
// 关闭排序弹窗
const closeSortModal = () => {
showSortModal.value = false
sortableGroups.value = []
}
// 保存排序
const saveSortOrder = async () => {
sortSubmitting.value = true
try {
const updates = sortableGroups.value.map((g, index) => ({
id: g.id,
sort_order: index * 10
}))
await adminAPI.groups.updateSortOrder(updates)
appStore.showSuccess(t('admin.groups.sortOrderUpdated'))
closeSortModal()
loadGroups()
} catch (error: any) {
appStore.showError(error.response?.data?.detail || t('admin.groups.failedToUpdateSortOrder'))
console.error('Error updating sort order:', error)
} finally {
sortSubmitting.value = false
}
}
onMounted(() => { onMounted(() => {
loadGroups() loadGroups()
document.addEventListener('click', handleClickOutside) document.addEventListener('click', handleClickOutside)

View File

@@ -1,26 +1,10 @@
<template> <template>
<AppLayout> <AppLayout>
<TablePageLayout> <TablePageLayout>
<template #actions>
<div class="flex justify-end gap-3">
<button
@click="loadCodes"
:disabled="loading"
class="btn btn-secondary"
:title="t('common.refresh')"
>
<Icon name="refresh" size="md" :class="loading ? 'animate-spin' : ''" />
</button>
<button @click="showCreateDialog = true" class="btn btn-primary">
<Icon name="plus" size="md" class="mr-1" />
{{ t('admin.promo.createCode') }}
</button>
</div>
</template>
<template #filters> <template #filters>
<div class="flex flex-col gap-4 sm:flex-row sm:items-center sm:justify-between"> <div class="flex flex-wrap items-center gap-3">
<div class="max-w-md flex-1"> <!-- Left: Search + Filters -->
<div class="flex-1 sm:max-w-64">
<input <input
v-model="searchQuery" v-model="searchQuery"
type="text" type="text"
@@ -29,13 +13,27 @@
@input="handleSearch" @input="handleSearch"
/> />
</div> </div>
<div class="flex gap-2"> <Select
<Select v-model="filters.status"
v-model="filters.status" :options="filterStatusOptions"
:options="filterStatusOptions" class="w-36"
class="w-36" @change="loadCodes"
@change="loadCodes" />
/>
<!-- Right: Action buttons -->
<div class="flex flex-1 flex-wrap items-center justify-end gap-2">
<button
@click="loadCodes"
:disabled="loading"
class="btn btn-secondary"
:title="t('common.refresh')"
>
<Icon name="refresh" size="md" :class="loading ? 'animate-spin' : ''" />
</button>
<button @click="showCreateDialog = true" class="btn btn-primary">
<Icon name="plus" size="md" class="mr-1" />
{{ t('admin.promo.createCode') }}
</button>
</div> </div>
</div> </div>
</template> </template>

View File

@@ -2,9 +2,42 @@
<AppLayout> <AppLayout>
<TablePageLayout> <TablePageLayout>
<template #filters> <template #filters>
<div class="space-y-3"> <div class="flex flex-wrap items-center gap-3">
<!-- Row 1: Actions --> <!-- Left: Search + Filters -->
<div class="flex flex-wrap items-center gap-3"> <div class="relative w-full sm:w-64">
<Icon
name="search"
size="md"
class="absolute left-3 top-1/2 -translate-y-1/2 text-gray-400 dark:text-gray-500"
/>
<input
v-model="searchQuery"
type="text"
:placeholder="t('admin.proxies.searchProxies')"
class="input pl-10"
@input="handleSearch"
/>
</div>
<div class="w-full sm:w-40">
<Select
v-model="filters.protocol"
:options="protocolOptions"
:placeholder="t('admin.proxies.allProtocols')"
@change="loadProxies"
/>
</div>
<div class="w-full sm:w-36">
<Select
v-model="filters.status"
:options="statusOptions"
:placeholder="t('admin.proxies.allStatus')"
@change="loadProxies"
/>
</div>
<!-- Right: All action buttons -->
<div class="flex flex-1 flex-wrap items-center justify-end gap-2">
<button <button
@click="loadProxies" @click="loadProxies"
:disabled="loading" :disabled="loading"
@@ -42,41 +75,6 @@
{{ t('admin.proxies.createProxy') }} {{ t('admin.proxies.createProxy') }}
</button> </button>
</div> </div>
<!-- Row 2: Search + Filters -->
<div class="flex flex-wrap items-center gap-3">
<div class="relative w-full sm:w-64">
<Icon
name="search"
size="md"
class="absolute left-3 top-1/2 -translate-y-1/2 text-gray-400 dark:text-gray-500"
/>
<input
v-model="searchQuery"
type="text"
:placeholder="t('admin.proxies.searchProxies')"
class="input pl-10"
@input="handleSearch"
/>
</div>
<div class="w-full sm:w-40">
<Select
v-model="filters.protocol"
:options="protocolOptions"
:placeholder="t('admin.proxies.allProtocols')"
@change="loadProxies"
/>
</div>
<div class="w-full sm:w-36">
<Select
v-model="filters.status"
:options="statusOptions"
:placeholder="t('admin.proxies.allStatus')"
@change="loadProxies"
/>
</div>
</div>
</div> </div>
</template> </template>

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