mirror of
https://gitee.com/wanwujie/sub2api
synced 2026-04-06 00:10:21 +08:00
Compare commits
117 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fb883f0092 | ||
|
|
64b8219245 | ||
|
|
2004230b66 | ||
|
|
0026e871f0 | ||
|
|
19d0ee130d | ||
|
|
942c3e1529 | ||
|
|
caa8c47b68 | ||
|
|
c328b741cb | ||
|
|
e85b35c6bd | ||
|
|
6e21a52271 | ||
|
|
4bbf71b7da | ||
|
|
42e2c5061d | ||
|
|
380c43cb03 | ||
|
|
bc75edd800 | ||
|
|
9774339fef | ||
|
|
026740b5e5 | ||
|
|
21a04332ec | ||
|
|
eec8b4c91e | ||
|
|
ef22d6f628 | ||
|
|
58545efbd7 | ||
|
|
2bd288a677 | ||
|
|
234e98f1b3 | ||
|
|
b31bfd53ab | ||
|
|
23412965f8 | ||
|
|
6a55b153fc | ||
|
|
e847cfc8a0 | ||
|
|
337d9ad755 | ||
|
|
dd247e55e9 | ||
|
|
1ad29032d3 | ||
|
|
c01db6b180 | ||
|
|
32b4b139a4 | ||
|
|
31fef105c7 | ||
|
|
1f5ced7069 | ||
|
|
2a70870469 | ||
|
|
9e9811cbb3 | ||
|
|
a5d6035c28 | ||
|
|
ecfad788d9 | ||
|
|
cf1d0f23cc | ||
|
|
995adaeee4 | ||
|
|
e247be6ead | ||
|
|
30b95cf5ce | ||
|
|
25b8a22648 | ||
|
|
0084da9ca5 | ||
|
|
31d4c1d2fe | ||
|
|
08ce6de4db | ||
|
|
7d4b7deea9 | ||
|
|
b6b739431c | ||
|
|
ad15d9970c | ||
|
|
ff57c860e3 | ||
|
|
635d7e77e1 | ||
|
|
ba9eb684ed | ||
|
|
9594c9c83a | ||
|
|
ff06583c5d | ||
|
|
b0389ca4d2 | ||
|
|
1d085d982b | ||
|
|
fb9d087838 | ||
|
|
18c6686fed | ||
|
|
6648e6506c | ||
|
|
386f6da14d | ||
|
|
d895a2c469 | ||
|
|
5f2d81d154 | ||
|
|
4e3499c0d7 | ||
|
|
26cdb1805d | ||
|
|
506cb21cb1 | ||
|
|
fd51ff6970 | ||
|
|
295d71be0a | ||
|
|
9bbe468c91 | ||
|
|
fbdff4f34f | ||
|
|
0aa480283f | ||
|
|
cd9d31f5f2 | ||
|
|
cbfce49aa1 | ||
|
|
1d1da7362b | ||
|
|
a8c173f043 | ||
|
|
97ab649d16 | ||
|
|
d3e73f1260 | ||
|
|
f3da4b202e | ||
|
|
530f6ad81c | ||
|
|
3252c378aa | ||
|
|
b5ca6a654c | ||
|
|
94749b12ac | ||
|
|
523fa9f71e | ||
|
|
54636781ea | ||
|
|
5187db5ee5 | ||
|
|
0b9c4ae69e | ||
|
|
0d5a8a95c8 | ||
|
|
9cd97c9e1d | ||
|
|
d521191e87 | ||
|
|
fd78993b91 | ||
|
|
937b1fb05d | ||
|
|
80cce858cb | ||
|
|
7bdb0e6b12 | ||
|
|
0743652d92 | ||
|
|
96bec5c9b1 | ||
|
|
cfeb6b8b14 | ||
|
|
481310dea0 | ||
|
|
ea2821d11d | ||
|
|
7a0de1765f | ||
|
|
17c3cb2403 | ||
|
|
35b1bc3753 | ||
|
|
8d38788672 | ||
|
|
c615a4264d | ||
|
|
227d506c53 | ||
|
|
36a86e9ab4 | ||
|
|
f133b051dc | ||
|
|
7af1bdbf4c | ||
|
|
016d7ef645 | ||
|
|
f1e47291cd | ||
|
|
d7e9ae38e4 | ||
|
|
88be981afc | ||
|
|
3f92a43170 | ||
|
|
2101f1d1c8 | ||
|
|
f0f920e49f | ||
|
|
95583fce83 | ||
|
|
a413fa3b17 | ||
|
|
3a8dbf5a99 | ||
|
|
254f12543c | ||
|
|
cf8a64528c |
58
.github/workflows/release.yml
vendored
58
.github/workflows/release.yml
vendored
@@ -143,3 +143,61 @@ jobs:
|
||||
repository: ${{ secrets.DOCKERHUB_USERNAME }}/sub2api
|
||||
short-description: "Sub2API - AI API Gateway Platform"
|
||||
readme-filepath: ./deploy/DOCKER.md
|
||||
|
||||
# Send Telegram notification
|
||||
- name: Send Telegram Notification
|
||||
env:
|
||||
TELEGRAM_BOT_TOKEN: ${{ secrets.TELEGRAM_BOT_TOKEN }}
|
||||
TELEGRAM_CHAT_ID: ${{ secrets.TELEGRAM_CHAT_ID }}
|
||||
continue-on-error: true
|
||||
run: |
|
||||
# 检查必要的环境变量
|
||||
if [ -z "$TELEGRAM_BOT_TOKEN" ] || [ -z "$TELEGRAM_CHAT_ID" ]; then
|
||||
echo "Telegram credentials not configured, skipping notification"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
TAG_NAME=${GITHUB_REF#refs/tags/}
|
||||
VERSION=${TAG_NAME#v}
|
||||
REPO="${{ github.repository }}"
|
||||
DOCKER_IMAGE="${{ secrets.DOCKERHUB_USERNAME }}/sub2api"
|
||||
|
||||
# 获取 tag message 内容
|
||||
TAG_MESSAGE='${{ steps.tag_message.outputs.message }}'
|
||||
|
||||
# 限制消息长度(Telegram 消息限制 4096 字符,预留空间给头尾固定内容)
|
||||
if [ ${#TAG_MESSAGE} -gt 3500 ]; then
|
||||
TAG_MESSAGE="${TAG_MESSAGE:0:3500}..."
|
||||
fi
|
||||
|
||||
# 构建消息内容
|
||||
MESSAGE="🚀 *Sub2API 新版本发布!*"$'\n'$'\n'
|
||||
MESSAGE+="📦 版本号: \`${VERSION}\`"$'\n'$'\n'
|
||||
|
||||
# 添加更新内容
|
||||
if [ -n "$TAG_MESSAGE" ]; then
|
||||
MESSAGE+="${TAG_MESSAGE}"$'\n'$'\n'
|
||||
fi
|
||||
|
||||
MESSAGE+="🐳 *Docker 部署:*"$'\n'
|
||||
MESSAGE+="\`\`\`bash"$'\n'
|
||||
MESSAGE+="docker pull ${DOCKER_IMAGE}:${TAG_NAME}"$'\n'
|
||||
MESSAGE+="docker pull ${DOCKER_IMAGE}:latest"$'\n'
|
||||
MESSAGE+="\`\`\`"$'\n'$'\n'
|
||||
MESSAGE+="🔗 *相关链接:*"$'\n'
|
||||
MESSAGE+="• [GitHub Release](https://github.com/${REPO}/releases/tag/${TAG_NAME})"$'\n'
|
||||
MESSAGE+="• [Docker Hub](https://hub.docker.com/r/${DOCKER_IMAGE})"$'\n'$'\n'
|
||||
MESSAGE+="#Sub2API #Release #${TAG_NAME//./_}"
|
||||
|
||||
# 发送消息
|
||||
curl -s -X POST "https://api.telegram.org/bot${TELEGRAM_BOT_TOKEN}/sendMessage" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "$(jq -n \
|
||||
--arg chat_id "${TELEGRAM_CHAT_ID}" \
|
||||
--arg text "${MESSAGE}" \
|
||||
'{
|
||||
chat_id: $chat_id,
|
||||
text: $text,
|
||||
parse_mode: "Markdown",
|
||||
disable_web_page_preview: true
|
||||
}')"
|
||||
|
||||
1
.gitignore
vendored
1
.gitignore
vendored
@@ -32,6 +32,7 @@ frontend/node_modules/
|
||||
frontend/dist/
|
||||
*.local
|
||||
*.tsbuildinfo
|
||||
vite.config.d.ts
|
||||
|
||||
# 日志
|
||||
npm-debug.log*
|
||||
|
||||
@@ -9,7 +9,7 @@
|
||||
# -----------------------------------------------------------------------------
|
||||
# Stage 1: Frontend Builder
|
||||
# -----------------------------------------------------------------------------
|
||||
FROM node:20-alpine AS frontend-builder
|
||||
FROM node:24-alpine AS frontend-builder
|
||||
|
||||
WORKDIR /app/frontend
|
||||
|
||||
@@ -24,7 +24,7 @@ RUN npm run build
|
||||
# -----------------------------------------------------------------------------
|
||||
# Stage 2: Backend Builder
|
||||
# -----------------------------------------------------------------------------
|
||||
FROM golang:1.24-alpine AS backend-builder
|
||||
FROM golang:1.25-alpine AS backend-builder
|
||||
|
||||
# Build arguments for version info (set by CI)
|
||||
ARG VERSION=docker
|
||||
|
||||
26
README.md
26
README.md
@@ -283,6 +283,32 @@ npm run dev
|
||||
|
||||
---
|
||||
|
||||
## Antigravity Support
|
||||
|
||||
Sub2API supports [Antigravity](https://antigravity.so/) accounts. After authorization, dedicated endpoints are available for Claude and Gemini models.
|
||||
|
||||
### Dedicated Endpoints
|
||||
|
||||
| Endpoint | Model |
|
||||
|----------|-------|
|
||||
| `/antigravity/v1/messages` | Claude models |
|
||||
| `/antigravity/v1beta/` | Gemini models |
|
||||
|
||||
### Claude Code Configuration
|
||||
|
||||
```bash
|
||||
export ANTHROPIC_BASE_URL="http://localhost:8080/antigravity"
|
||||
export ANTHROPIC_AUTH_TOKEN="sk-xxx"
|
||||
```
|
||||
|
||||
### Hybrid Scheduling Mode
|
||||
|
||||
Antigravity accounts support optional **hybrid scheduling**. When enabled, the general endpoints `/v1/messages` and `/v1beta/` will also route requests to Antigravity accounts.
|
||||
|
||||
> **⚠️ Warning**: Anthropic Claude and Antigravity Claude **cannot be mixed within the same conversation context**. Use groups to isolate them properly.
|
||||
|
||||
---
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
|
||||
36
README_CN.md
36
README_CN.md
@@ -283,6 +283,42 @@ npm run dev
|
||||
|
||||
---
|
||||
|
||||
## 简易模式
|
||||
|
||||
简易模式适合个人开发者或内部团队快速使用,不依赖完整 SaaS 功能。
|
||||
|
||||
- 启用方式:设置环境变量 `RUN_MODE=simple`
|
||||
- 功能差异:隐藏 SaaS 相关功能,跳过计费流程
|
||||
- 安全注意事项:生产环境需同时设置 `SIMPLE_MODE_CONFIRM=true` 才允许启动
|
||||
|
||||
---
|
||||
|
||||
## Antigravity 使用说明
|
||||
|
||||
Sub2API 支持 [Antigravity](https://antigravity.so/) 账户,授权后可通过专用端点访问 Claude 和 Gemini 模型。
|
||||
|
||||
### 专用端点
|
||||
|
||||
| 端点 | 模型 |
|
||||
|------|------|
|
||||
| `/antigravity/v1/messages` | Claude 模型 |
|
||||
| `/antigravity/v1beta/` | Gemini 模型 |
|
||||
|
||||
### Claude Code 配置示例
|
||||
|
||||
```bash
|
||||
export ANTHROPIC_BASE_URL="http://localhost:8080/antigravity"
|
||||
export ANTHROPIC_AUTH_TOKEN="sk-xxx"
|
||||
```
|
||||
|
||||
### 混合调度模式
|
||||
|
||||
Antigravity 账户支持可选的**混合调度**功能。开启后,通用端点 `/v1/messages` 和 `/v1beta/` 也会调度该账户。
|
||||
|
||||
> **⚠️ 注意**:Anthropic Claude 和 Antigravity Claude **不能在同一上下文中混合使用**,请通过分组功能做好隔离。
|
||||
|
||||
---
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
|
||||
@@ -599,4 +599,4 @@ formatters:
|
||||
- pattern: 'interface{}'
|
||||
replacement: 'any'
|
||||
- pattern: 'a[b:len(a)]'
|
||||
replacement: 'a[b:]'
|
||||
replacement: 'a[b:]'
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
.PHONY: wire build build-embed test-unit test-integration test-cover-integration clean-coverage
|
||||
.PHONY: wire build build-embed test-unit test-integration test-e2e test-cover-integration clean-coverage
|
||||
|
||||
wire:
|
||||
@echo "生成 Wire 代码..."
|
||||
@@ -21,6 +21,10 @@ test-unit:
|
||||
test-integration:
|
||||
@go test -tags integration ./... -count=1 -race -parallel=8
|
||||
|
||||
test-e2e:
|
||||
@echo "运行 E2E 测试(需要本地服务器运行)..."
|
||||
@go test -tags e2e ./internal/integration/... -count=1 -v
|
||||
|
||||
test-cover-integration:
|
||||
@echo "运行集成测试并生成覆盖率报告..."
|
||||
@go test -tags=integration -cover -coverprofile=coverage.out -count=1 -race -parallel=8 ./...
|
||||
|
||||
@@ -107,6 +107,14 @@ func runSetupServer() {
|
||||
}
|
||||
|
||||
func runMainServer() {
|
||||
cfg, err := config.Load()
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to load config: %v", err)
|
||||
}
|
||||
if cfg.RunMode == config.RunModeSimple {
|
||||
log.Println("⚠️ WARNING: Running in SIMPLE mode - billing and quota checks are DISABLED")
|
||||
}
|
||||
|
||||
buildInfo := handler.BuildInfo{
|
||||
Version: Version,
|
||||
BuildType: BuildType,
|
||||
|
||||
@@ -29,26 +29,26 @@ type Application struct {
|
||||
|
||||
func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
wire.Build(
|
||||
// 基础设施层 ProviderSets
|
||||
// Infrastructure layer ProviderSets
|
||||
config.ProviderSet,
|
||||
infrastructure.ProviderSet,
|
||||
|
||||
// 业务层 ProviderSets
|
||||
// Business layer ProviderSets
|
||||
repository.ProviderSet,
|
||||
service.ProviderSet,
|
||||
middleware.ProviderSet,
|
||||
handler.ProviderSet,
|
||||
|
||||
// 服务器层 ProviderSet
|
||||
// Server layer ProviderSet
|
||||
server.ProviderSet,
|
||||
|
||||
// BuildInfo provider
|
||||
provideServiceBuildInfo,
|
||||
|
||||
// 清理函数提供者
|
||||
// Cleanup function provider
|
||||
provideCleanup,
|
||||
|
||||
// 应用程序结构体
|
||||
// Application struct
|
||||
wire.Struct(new(Application), "Server", "Cleanup"),
|
||||
)
|
||||
return nil, nil
|
||||
@@ -70,6 +70,8 @@ func provideCleanup(
|
||||
oauth *service.OAuthService,
|
||||
openaiOAuth *service.OpenAIOAuthService,
|
||||
geminiOAuth *service.GeminiOAuthService,
|
||||
antigravityOAuth *service.AntigravityOAuthService,
|
||||
antigravityQuota *service.AntigravityQuotaRefresher,
|
||||
) func() {
|
||||
return func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
@@ -104,6 +106,14 @@ func provideCleanup(
|
||||
geminiOAuth.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"AntigravityOAuthService", func() error {
|
||||
antigravityOAuth.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"AntigravityQuotaRefresher", func() error {
|
||||
antigravityQuota.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"Redis", func() error {
|
||||
return rdb.Close()
|
||||
}},
|
||||
|
||||
@@ -49,7 +49,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
emailQueueService := service.ProvideEmailQueueService(emailService)
|
||||
authService := service.NewAuthService(userRepository, configConfig, settingService, emailService, turnstileService, emailQueueService)
|
||||
userService := service.NewUserService(userRepository)
|
||||
authHandler := handler.NewAuthHandler(authService, userService)
|
||||
authHandler := handler.NewAuthHandler(configConfig, authService, userService)
|
||||
userHandler := handler.NewUserHandler(userService)
|
||||
apiKeyRepository := repository.NewApiKeyRepository(db)
|
||||
groupRepository := repository.NewGroupRepository(db)
|
||||
@@ -62,7 +62,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
usageHandler := handler.NewUsageHandler(usageService, apiKeyService)
|
||||
redeemCodeRepository := repository.NewRedeemCodeRepository(db)
|
||||
billingCache := repository.NewBillingCache(client)
|
||||
billingCacheService := service.NewBillingCacheService(billingCache, userRepository, userSubscriptionRepository)
|
||||
billingCacheService := service.NewBillingCacheService(billingCache, userRepository, userSubscriptionRepository, configConfig)
|
||||
subscriptionService := service.NewSubscriptionService(groupRepository, userSubscriptionRepository, billingCacheService)
|
||||
redeemCache := repository.NewRedeemCache(client)
|
||||
redeemService := service.NewRedeemService(redeemCodeRepository, userRepository, subscriptionService, redeemCache, billingCacheService)
|
||||
@@ -97,6 +97,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
oAuthHandler := admin.NewOAuthHandler(oAuthService)
|
||||
openAIOAuthHandler := admin.NewOpenAIOAuthHandler(openAIOAuthService, adminService)
|
||||
geminiOAuthHandler := admin.NewGeminiOAuthHandler(geminiOAuthService)
|
||||
antigravityOAuthService := service.NewAntigravityOAuthService(proxyRepository)
|
||||
antigravityOAuthHandler := admin.NewAntigravityOAuthHandler(antigravityOAuthService)
|
||||
proxyHandler := admin.NewProxyHandler(adminService)
|
||||
adminRedeemHandler := admin.NewRedeemHandler(adminService)
|
||||
settingHandler := admin.NewSettingHandler(settingService, emailService)
|
||||
@@ -107,7 +109,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
systemHandler := handler.ProvideSystemHandler(updateService)
|
||||
adminSubscriptionHandler := admin.NewSubscriptionHandler(subscriptionService)
|
||||
adminUsageHandler := admin.NewUsageHandler(usageService, apiKeyService, adminService)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, proxyHandler, adminRedeemHandler, settingHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, proxyHandler, adminRedeemHandler, settingHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler)
|
||||
gatewayCache := repository.NewGatewayCache(client)
|
||||
pricingRemoteClient := repository.NewPricingRemoteClient()
|
||||
pricingService, err := service.ProvidePricingService(configConfig, pricingRemoteClient)
|
||||
@@ -117,20 +119,25 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
billingService := service.NewBillingService(configConfig, pricingService)
|
||||
identityCache := repository.NewIdentityCache(client)
|
||||
identityService := service.NewIdentityService(identityCache)
|
||||
gatewayService := service.NewGatewayService(accountRepository, usageLogRepository, userRepository, userSubscriptionRepository, gatewayCache, configConfig, billingService, rateLimitService, billingCacheService, identityService, httpUpstream)
|
||||
geminiMessagesCompatService := service.NewGeminiMessagesCompatService(accountRepository, gatewayCache, geminiTokenProvider, rateLimitService, httpUpstream)
|
||||
gatewayHandler := handler.NewGatewayHandler(gatewayService, geminiMessagesCompatService, userService, concurrencyService, billingCacheService)
|
||||
openAIGatewayService := service.NewOpenAIGatewayService(accountRepository, usageLogRepository, userRepository, userSubscriptionRepository, gatewayCache, configConfig, billingService, rateLimitService, billingCacheService, httpUpstream)
|
||||
timingWheelService := service.ProvideTimingWheelService()
|
||||
deferredService := service.ProvideDeferredService(accountRepository, timingWheelService)
|
||||
gatewayService := service.NewGatewayService(accountRepository, groupRepository, usageLogRepository, userRepository, userSubscriptionRepository, gatewayCache, configConfig, billingService, rateLimitService, billingCacheService, identityService, httpUpstream, deferredService)
|
||||
antigravityTokenProvider := service.NewAntigravityTokenProvider(accountRepository, geminiTokenCache, antigravityOAuthService)
|
||||
antigravityGatewayService := service.NewAntigravityGatewayService(accountRepository, gatewayCache, antigravityTokenProvider, rateLimitService, httpUpstream)
|
||||
geminiMessagesCompatService := service.NewGeminiMessagesCompatService(accountRepository, groupRepository, gatewayCache, geminiTokenProvider, rateLimitService, httpUpstream, antigravityGatewayService)
|
||||
gatewayHandler := handler.NewGatewayHandler(gatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService)
|
||||
openAIGatewayService := service.NewOpenAIGatewayService(accountRepository, usageLogRepository, userRepository, userSubscriptionRepository, gatewayCache, configConfig, billingService, rateLimitService, billingCacheService, httpUpstream, deferredService)
|
||||
openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService)
|
||||
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo)
|
||||
handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler)
|
||||
jwtAuthMiddleware := middleware.NewJWTAuthMiddleware(authService, userService)
|
||||
adminAuthMiddleware := middleware.NewAdminAuthMiddleware(authService, userService, settingService)
|
||||
apiKeyAuthMiddleware := middleware.NewApiKeyAuthMiddleware(apiKeyService, subscriptionService)
|
||||
apiKeyAuthMiddleware := middleware.NewApiKeyAuthMiddleware(apiKeyService, subscriptionService, configConfig)
|
||||
engine := server.ProvideRouter(configConfig, handlers, jwtAuthMiddleware, adminAuthMiddleware, apiKeyAuthMiddleware, apiKeyService, subscriptionService)
|
||||
httpServer := server.ProvideHTTPServer(configConfig, engine)
|
||||
tokenRefreshService := service.ProvideTokenRefreshService(accountRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
|
||||
v := provideCleanup(db, client, tokenRefreshService, pricingService, emailQueueService, oAuthService, openAIOAuthService, geminiOAuthService)
|
||||
tokenRefreshService := service.ProvideTokenRefreshService(accountRepository, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, configConfig)
|
||||
antigravityQuotaRefresher := service.ProvideAntigravityQuotaRefresher(accountRepository, proxyRepository, antigravityOAuthService, configConfig)
|
||||
v := provideCleanup(db, client, tokenRefreshService, pricingService, emailQueueService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, antigravityQuotaRefresher)
|
||||
application := &Application{
|
||||
Server: httpServer,
|
||||
Cleanup: v,
|
||||
@@ -161,6 +168,8 @@ func provideCleanup(
|
||||
oauth *service.OAuthService,
|
||||
openaiOAuth *service.OpenAIOAuthService,
|
||||
geminiOAuth *service.GeminiOAuthService,
|
||||
antigravityOAuth *service.AntigravityOAuthService,
|
||||
antigravityQuota *service.AntigravityQuotaRefresher,
|
||||
) func() {
|
||||
return func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
@@ -194,6 +203,14 @@ func provideCleanup(
|
||||
geminiOAuth.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"AntigravityOAuthService", func() error {
|
||||
antigravityOAuth.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"AntigravityQuotaRefresher", func() error {
|
||||
antigravityQuota.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"Redis", func() error {
|
||||
return rdb.Close()
|
||||
}},
|
||||
|
||||
@@ -11,13 +11,14 @@ require (
|
||||
github.com/google/wire v0.7.0
|
||||
github.com/imroc/req/v3 v3.56.0
|
||||
github.com/lib/pq v1.10.9
|
||||
github.com/redis/go-redis/v9 v9.7.3
|
||||
github.com/redis/go-redis/v9 v9.17.2
|
||||
github.com/spf13/viper v1.18.2
|
||||
github.com/stretchr/testify v1.11.1
|
||||
github.com/testcontainers/testcontainers-go/modules/postgres v0.40.0
|
||||
github.com/testcontainers/testcontainers-go/modules/redis v0.40.0
|
||||
github.com/tidwall/gjson v1.18.0
|
||||
github.com/tidwall/sjson v1.2.5
|
||||
github.com/zeromicro/go-zero v1.9.4
|
||||
golang.org/x/crypto v0.44.0
|
||||
golang.org/x/net v0.47.0
|
||||
golang.org/x/term v0.37.0
|
||||
@@ -49,6 +50,7 @@ require (
|
||||
github.com/docker/go-connections v0.6.0 // indirect
|
||||
github.com/docker/go-units v0.5.0 // indirect
|
||||
github.com/ebitengine/purego v0.8.4 // indirect
|
||||
github.com/fatih/color v1.18.0 // indirect
|
||||
github.com/felixge/httpsnoop v1.0.4 // indirect
|
||||
github.com/fsnotify/fsnotify v1.7.0 // indirect
|
||||
github.com/gabriel-vasile/mimetype v1.4.2 // indirect
|
||||
@@ -59,7 +61,7 @@ require (
|
||||
github.com/go-playground/locales v0.14.1 // indirect
|
||||
github.com/go-playground/universal-translator v0.18.1 // indirect
|
||||
github.com/go-playground/validator/v10 v10.14.0 // indirect
|
||||
github.com/go-sql-driver/mysql v1.8.1 // indirect
|
||||
github.com/go-sql-driver/mysql v1.9.0 // indirect
|
||||
github.com/goccy/go-json v0.10.2 // indirect
|
||||
github.com/google/go-querystring v1.1.0 // indirect
|
||||
github.com/google/subcommands v1.2.0 // indirect
|
||||
@@ -67,9 +69,9 @@ require (
|
||||
github.com/hashicorp/hcl v1.0.0 // indirect
|
||||
github.com/icholy/digest v1.1.0 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20231201235250-de7065d80cb9 // indirect
|
||||
github.com/jackc/pgx/v5 v5.5.5 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.1 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
|
||||
github.com/jackc/pgx/v5 v5.7.4 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
github.com/jinzhu/now v1.1.5 // indirect
|
||||
github.com/json-iterator/go v1.1.12 // indirect
|
||||
@@ -78,7 +80,8 @@ require (
|
||||
github.com/leodido/go-urn v1.2.4 // indirect
|
||||
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
|
||||
github.com/magiconair/properties v1.8.10 // indirect
|
||||
github.com/mattn/go-isatty v0.0.19 // indirect
|
||||
github.com/mattn/go-colorable v0.1.13 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/mdelapenya/tlscert v0.2.0 // indirect
|
||||
github.com/mitchellh/mapstructure v1.5.0 // indirect
|
||||
github.com/moby/docker-image-spec v1.3.1 // indirect
|
||||
@@ -93,7 +96,7 @@ require (
|
||||
github.com/morikuni/aec v1.0.0 // indirect
|
||||
github.com/opencontainers/go-digest v1.0.0 // indirect
|
||||
github.com/opencontainers/image-spec v1.1.1 // indirect
|
||||
github.com/pelletier/go-toml/v2 v2.1.0 // indirect
|
||||
github.com/pelletier/go-toml/v2 v2.2.2 // indirect
|
||||
github.com/pkg/errors v0.9.1 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
|
||||
@@ -105,6 +108,7 @@ require (
|
||||
github.com/shirou/gopsutil/v4 v4.25.6 // indirect
|
||||
github.com/sirupsen/logrus v1.9.3 // indirect
|
||||
github.com/sourcegraph/conc v0.3.0 // indirect
|
||||
github.com/spaolacci/murmur3 v1.1.0 // indirect
|
||||
github.com/spf13/afero v1.11.0 // indirect
|
||||
github.com/spf13/cast v1.6.0 // indirect
|
||||
github.com/spf13/pflag v1.0.5 // indirect
|
||||
@@ -123,7 +127,8 @@ require (
|
||||
go.opentelemetry.io/otel/metric v1.37.0 // indirect
|
||||
go.opentelemetry.io/otel/sdk v1.37.0 // indirect
|
||||
go.opentelemetry.io/otel/trace v1.37.0 // indirect
|
||||
go.uber.org/atomic v1.9.0 // indirect
|
||||
go.uber.org/atomic v1.10.0 // indirect
|
||||
go.uber.org/automaxprocs v1.6.0 // indirect
|
||||
go.uber.org/multierr v1.9.0 // indirect
|
||||
golang.org/x/arch v0.3.0 // indirect
|
||||
golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect
|
||||
|
||||
@@ -52,6 +52,8 @@ github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4
|
||||
github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
|
||||
github.com/ebitengine/purego v0.8.4 h1:CF7LEKg5FFOsASUj0+QwaXf8Ht6TlFxg09+S9wz0omw=
|
||||
github.com/ebitengine/purego v0.8.4/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ=
|
||||
github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
|
||||
github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
|
||||
github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
|
||||
github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
|
||||
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
|
||||
@@ -80,8 +82,8 @@ github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91
|
||||
github.com/go-playground/validator/v10 v10.14.0 h1:vgvQWe3XCz3gIeFDm/HnTIbj6UGmg/+t63MyGU2n5js=
|
||||
github.com/go-playground/validator/v10 v10.14.0/go.mod h1:9iXMNT7sEkjXb0I+enO7QXmzG6QCsPWY4zveKFVRSyU=
|
||||
github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
|
||||
github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y=
|
||||
github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
|
||||
github.com/go-sql-driver/mysql v1.9.0 h1:Y0zIbQXhQKmQgTp44Y1dp3wTXcn804QoTptLZT1vtvo=
|
||||
github.com/go-sql-driver/mysql v1.9.0/go.mod h1:pDetrLJeA3oMujJuvXc8RJoasr589B6A9fwzD3QMrqw=
|
||||
github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
|
||||
github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
|
||||
github.com/golang-jwt/jwt/v5 v5.2.0 h1:d/ix8ftRUorsN+5eMIlF4T6J8CAt9rch3My2winC1Jw=
|
||||
@@ -113,12 +115,12 @@ github.com/imroc/req/v3 v3.56.0 h1:t6YdqqerYBXhZ9+VjqsQs5wlKxdUNEvsgBhxWc1AEEo=
|
||||
github.com/imroc/req/v3 v3.56.0/go.mod h1:cUZSooE8hhzFNOrAbdxuemXDQxFXLQTnu3066jr7ZGk=
|
||||
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/pgservicefile v0.0.0-20231201235250-de7065d80cb9 h1:L0QtFUgDarD7Fpv9jeVMgy/+Ec0mtnmYuImjTz6dtDA=
|
||||
github.com/jackc/pgservicefile v0.0.0-20231201235250-de7065d80cb9/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
|
||||
github.com/jackc/pgx/v5 v5.5.5 h1:amBjrZVmksIdNjxGW/IiIMzxMKZFelXbUoPNb+8sjQw=
|
||||
github.com/jackc/pgx/v5 v5.5.5/go.mod h1:ez9gk+OAat140fv9ErkZDYFWmXLfV+++K0uAOiwgm1A=
|
||||
github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk=
|
||||
github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
|
||||
github.com/jackc/pgx/v5 v5.7.4 h1:9wKznZrhWa2QiHL+NjTSPP6yjl3451BX3imWDnokYlg=
|
||||
github.com/jackc/pgx/v5 v5.7.4/go.mod h1:ncY89UGWxg82EykZUwSpUKEfccBGGYq1xjrOpsbsfGQ=
|
||||
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
|
||||
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
|
||||
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
|
||||
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
|
||||
@@ -142,8 +144,11 @@ github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ
|
||||
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=
|
||||
github.com/magiconair/properties v1.8.10 h1:s31yESBquKXCV9a/ScB3ESkOjUYYv+X0rg8SYxI99mE=
|
||||
github.com/magiconair/properties v1.8.10/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
|
||||
github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
|
||||
github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
|
||||
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
|
||||
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/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-sqlite3 v1.14.15 h1:vfoHhTN1af61xCRSWzFIWzx2YskyMTwHLrExkBOjvxI=
|
||||
github.com/mattn/go-sqlite3 v1.14.15/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
|
||||
github.com/mdelapenya/tlscert v0.2.0 h1:7H81W6Z/4weDvZBNOfQte5GpIMo0lGYEeWbkGp5LJHI=
|
||||
@@ -179,8 +184,8 @@ github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8
|
||||
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/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M=
|
||||
github.com/pelletier/go-toml/v2 v2.1.0 h1:FnwAJ4oYMvbT/34k9zzHuZNrhlz48GB3/s6at6/MHO4=
|
||||
github.com/pelletier/go-toml/v2 v2.1.0/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc=
|
||||
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/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
|
||||
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
@@ -188,12 +193,14 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRI
|
||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
|
||||
github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
|
||||
github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
|
||||
github.com/quic-go/qpack v0.5.1 h1:giqksBPnT/HDtZ6VhtFKgoLOWmlyo9Ei6u9PqzIMbhI=
|
||||
github.com/quic-go/qpack v0.5.1/go.mod h1:+PC4XFrEskIVkcLzpEkbLqq1uCoxPhQuvK5rH1ZgaEg=
|
||||
github.com/quic-go/quic-go v0.56.0 h1:q/TW+OLismmXAehgFLczhCDTYB3bFmua4D9lsNBWxvY=
|
||||
github.com/quic-go/quic-go v0.56.0/go.mod h1:9gx5KsFQtw2oZ6GZTyh+7YEvOxWCL9WZAepnHxgAo6c=
|
||||
github.com/redis/go-redis/v9 v9.7.3 h1:YpPyAayJV+XErNsatSElgRZZVCwXX9QzkKYNvO7x0wM=
|
||||
github.com/redis/go-redis/v9 v9.7.3/go.mod h1:bGUrSggJ9X9GUmZpZNEOQKaANxSGgOEBRltRTZHSvrA=
|
||||
github.com/redis/go-redis/v9 v9.17.2 h1:P2EGsA4qVIM3Pp+aPocCJ7DguDHhqrXNhVcEp4ViluI=
|
||||
github.com/redis/go-redis/v9 v9.17.2/go.mod h1:u410H11HMLoB+TP67dz8rL9s6QW2j76l0//kSOd3370=
|
||||
github.com/refraction-networking/utls v1.8.1 h1:yNY1kapmQU8JeM1sSw2H2asfTIwWxIkrMJI0pRUOCAo=
|
||||
github.com/refraction-networking/utls v1.8.1/go.mod h1:jkSOEkLqn+S/jtpEHPOsVv/4V4EVnelwbMQl4vCWXAM=
|
||||
github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
|
||||
@@ -208,6 +215,8 @@ github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ
|
||||
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
|
||||
github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=
|
||||
github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0=
|
||||
github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI=
|
||||
github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
|
||||
github.com/spf13/afero v1.11.0 h1:WJQKhtpdm3v2IzqG8VMqrr6Rf3UYpEF239Jy9wNepM8=
|
||||
github.com/spf13/afero v1.11.0/go.mod h1:GH9Y3pIexgf1MTIWtNGyogA5MwRIDXGUr+hbWNoBjkY=
|
||||
github.com/spf13/cast v1.6.0 h1:GEiTHELF+vaR5dhz3VqZfFSzZjYbgeKDpBxQVS4GYJ0=
|
||||
@@ -228,6 +237,7 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
|
||||
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
|
||||
github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
|
||||
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
|
||||
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||
github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
|
||||
@@ -259,26 +269,30 @@ github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZ
|
||||
github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
|
||||
github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
|
||||
github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
|
||||
github.com/zeromicro/go-zero v1.9.4 h1:aRLFoISqAYijABtkbliQC5SsI5TbizJpQvoHc9xup8k=
|
||||
github.com/zeromicro/go-zero v1.9.4/go.mod h1:a17JOTch25SWxBcUgJZYps60hygK3pIYdw7nGwlcS38=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 h1:jq9TW8u3so/bN+JPT166wjOI6/vQPF6Xe7nMNIltagk=
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0/go.mod h1:p8pYQP+m5XfbZm9fxtSKAbM6oIllS7s2AfxrChvc7iw=
|
||||
go.opentelemetry.io/otel v1.37.0 h1:9zhNfelUvx0KBfu/gb+ZgeAfAgtWrfHJZcAqFC228wQ=
|
||||
go.opentelemetry.io/otel v1.37.0/go.mod h1:ehE/umFRLnuLa/vSccNq9oS1ErUlkkK71gMcN34UG8I=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.19.0 h1:Mne5On7VWdx7omSrSSZvM4Kw7cS7NQkOOmLcgscI51U=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.19.0/go.mod h1:IPtUMKL4O3tH5y+iXVyAXqpAwMuzC1IrxVS81rummfE=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.19.0 h1:IeMeyr1aBvBiPVYihXIaeIZba6b8E1bYp7lbdxK8CQg=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.19.0/go.mod h1:oVdCUtjq9MK9BlS7TtucsQwUcXcymNiEDjgDD2jMtZU=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.24.0 h1:t6wl9SPayj+c7lEIFgm4ooDBZVb01IhLB4InpomhRw8=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.24.0/go.mod h1:iSDOcsnSA5INXzZtwaBPrKp/lWu/V14Dd+llD0oI2EA=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.24.0 h1:Xw8U6u2f8DK2XAkGRFV7BBLENgnTGX9i4rQRxJf+/vs=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.24.0/go.mod h1:6KW1Fm6R/s6Z3PGXwSJN2K4eT6wQB3vXX6CVnYX9NmM=
|
||||
go.opentelemetry.io/otel/metric v1.37.0 h1:mvwbQS5m0tbmqML4NqK+e3aDiO02vsf/WgbsdpcPoZE=
|
||||
go.opentelemetry.io/otel/metric v1.37.0/go.mod h1:04wGrZurHYKOc+RKeye86GwKiTb9FKm1WHtO+4EVr2E=
|
||||
go.opentelemetry.io/otel/sdk v1.37.0 h1:ItB0QUqnjesGRvNcmAcU0LyvkVyGJ2xftD29bWdDvKI=
|
||||
go.opentelemetry.io/otel/sdk v1.37.0/go.mod h1:VredYzxUvuo2q3WRcDnKDjbdvmO0sCzOvVAiY+yUkAg=
|
||||
go.opentelemetry.io/otel/trace v1.37.0 h1:HLdcFNbRQBE2imdSEgm/kwqmQj1Or1l/7bW6mxVK7z4=
|
||||
go.opentelemetry.io/otel/trace v1.37.0/go.mod h1:TlgrlQ+PtQO5XFerSPUYG0JSgGyryXewPGyayAWSBS0=
|
||||
go.opentelemetry.io/proto/otlp v1.0.0 h1:T0TX0tmXU8a3CbNXzEKGeU5mIVOdf0oykP+u2lIVU/I=
|
||||
go.opentelemetry.io/proto/otlp v1.0.0/go.mod h1:Sy6pihPLfYHkr3NkUbEhGHFhINUSI/v80hjKIs5JXpM=
|
||||
go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE=
|
||||
go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
|
||||
go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0=
|
||||
go.opentelemetry.io/proto/otlp v1.3.1/go.mod h1:0X1WI4de4ZsLrrJNLAQbFeLCm3T7yBkR0XqQ7niQU+8=
|
||||
go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ=
|
||||
go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
|
||||
go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs=
|
||||
go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8=
|
||||
go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y=
|
||||
go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU=
|
||||
go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI=
|
||||
@@ -301,6 +315,7 @@ golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7w
|
||||
golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
|
||||
@@ -7,6 +7,11 @@ import (
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
const (
|
||||
RunModeStandard = "standard"
|
||||
RunModeSimple = "simple"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
Server ServerConfig `mapstructure:"server"`
|
||||
Database DatabaseConfig `mapstructure:"database"`
|
||||
@@ -17,6 +22,7 @@ type Config struct {
|
||||
Pricing PricingConfig `mapstructure:"pricing"`
|
||||
Gateway GatewayConfig `mapstructure:"gateway"`
|
||||
TokenRefresh TokenRefreshConfig `mapstructure:"token_refresh"`
|
||||
RunMode string `mapstructure:"run_mode" yaml:"run_mode"`
|
||||
Timezone string `mapstructure:"timezone"` // e.g. "Asia/Shanghai", "UTC"
|
||||
Gemini GeminiConfig `mapstructure:"gemini"`
|
||||
}
|
||||
@@ -135,6 +141,16 @@ type RateLimitConfig struct {
|
||||
OverloadCooldownMinutes int `mapstructure:"overload_cooldown_minutes"` // 529过载冷却时间(分钟)
|
||||
}
|
||||
|
||||
func NormalizeRunMode(value string) string {
|
||||
normalized := strings.ToLower(strings.TrimSpace(value))
|
||||
switch normalized {
|
||||
case RunModeStandard, RunModeSimple:
|
||||
return normalized
|
||||
default:
|
||||
return RunModeStandard
|
||||
}
|
||||
}
|
||||
|
||||
func Load() (*Config, error) {
|
||||
viper.SetConfigName("config")
|
||||
viper.SetConfigType("yaml")
|
||||
@@ -161,6 +177,8 @@ func Load() (*Config, error) {
|
||||
return nil, fmt.Errorf("unmarshal config error: %w", err)
|
||||
}
|
||||
|
||||
cfg.RunMode = NormalizeRunMode(cfg.RunMode)
|
||||
|
||||
if err := cfg.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("validate config error: %w", err)
|
||||
}
|
||||
@@ -169,6 +187,8 @@ func Load() (*Config, error) {
|
||||
}
|
||||
|
||||
func setDefaults() {
|
||||
viper.SetDefault("run_mode", RunModeStandard)
|
||||
|
||||
// Server
|
||||
viper.SetDefault("server.host", "0.0.0.0")
|
||||
viper.SetDefault("server.port", 8080)
|
||||
|
||||
23
backend/internal/config/config_test.go
Normal file
23
backend/internal/config/config_test.go
Normal file
@@ -0,0 +1,23 @@
|
||||
package config
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestNormalizeRunMode(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{"simple", "simple"},
|
||||
{"SIMPLE", "simple"},
|
||||
{"standard", "standard"},
|
||||
{"invalid", "standard"},
|
||||
{"", "standard"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
result := NormalizeRunMode(tt.input)
|
||||
if result != tt.expected {
|
||||
t.Errorf("NormalizeRunMode(%q) = %q, want %q", tt.input, result, tt.expected)
|
||||
}
|
||||
}
|
||||
}
|
||||
67
backend/internal/handler/admin/antigravity_oauth_handler.go
Normal file
67
backend/internal/handler/admin/antigravity_oauth_handler.go
Normal file
@@ -0,0 +1,67 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type AntigravityOAuthHandler struct {
|
||||
antigravityOAuthService *service.AntigravityOAuthService
|
||||
}
|
||||
|
||||
func NewAntigravityOAuthHandler(antigravityOAuthService *service.AntigravityOAuthService) *AntigravityOAuthHandler {
|
||||
return &AntigravityOAuthHandler{antigravityOAuthService: antigravityOAuthService}
|
||||
}
|
||||
|
||||
type AntigravityGenerateAuthURLRequest struct {
|
||||
ProxyID *int64 `json:"proxy_id"`
|
||||
}
|
||||
|
||||
// GenerateAuthURL generates Google OAuth authorization URL
|
||||
// POST /api/v1/admin/antigravity/oauth/auth-url
|
||||
func (h *AntigravityOAuthHandler) GenerateAuthURL(c *gin.Context) {
|
||||
var req AntigravityGenerateAuthURLRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "请求无效: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
result, err := h.antigravityOAuthService.GenerateAuthURL(c.Request.Context(), req.ProxyID)
|
||||
if err != nil {
|
||||
response.InternalError(c, "生成授权链接失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
type AntigravityExchangeCodeRequest struct {
|
||||
SessionID string `json:"session_id" binding:"required"`
|
||||
State string `json:"state" binding:"required"`
|
||||
Code string `json:"code" binding:"required"`
|
||||
ProxyID *int64 `json:"proxy_id"`
|
||||
}
|
||||
|
||||
// ExchangeCode 用 authorization code 交换 token
|
||||
// POST /api/v1/admin/antigravity/oauth/exchange-code
|
||||
func (h *AntigravityOAuthHandler) ExchangeCode(c *gin.Context) {
|
||||
var req AntigravityExchangeCodeRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "请求无效: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
tokenInfo, err := h.antigravityOAuthService.ExchangeCode(c.Request.Context(), &service.AntigravityExchangeCodeInput{
|
||||
SessionID: req.SessionID,
|
||||
State: req.State,
|
||||
Code: req.Code,
|
||||
ProxyID: req.ProxyID,
|
||||
})
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Token 交换失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, tokenInfo)
|
||||
}
|
||||
@@ -26,7 +26,7 @@ func NewGroupHandler(adminService service.AdminService) *GroupHandler {
|
||||
type CreateGroupRequest struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity"`
|
||||
RateMultiplier float64 `json:"rate_multiplier"`
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
|
||||
@@ -39,7 +39,7 @@ type CreateGroupRequest struct {
|
||||
type UpdateGroupRequest struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity"`
|
||||
RateMultiplier *float64 `json:"rate_multiplier"`
|
||||
IsExclusive *bool `json:"is_exclusive"`
|
||||
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
|
||||
|
||||
@@ -30,8 +30,8 @@ type GenerateRedeemCodesRequest struct {
|
||||
Count int `json:"count" binding:"required,min=1,max=100"`
|
||||
Type string `json:"type" binding:"required,oneof=balance concurrency subscription"`
|
||||
Value float64 `json:"value" binding:"min=0"`
|
||||
GroupID *int64 `json:"group_id"` // 订阅类型必填
|
||||
ValidityDays int `json:"validity_days"` // 订阅类型使用,默认30天
|
||||
GroupID *int64 `json:"group_id"` // 订阅类型必填
|
||||
ValidityDays int `json:"validity_days" binding:"omitempty,max=36500"` // 订阅类型使用,默认30天,最大100年
|
||||
}
|
||||
|
||||
// List handles listing all redeem codes with pagination
|
||||
|
||||
@@ -41,7 +41,7 @@ func NewSubscriptionHandler(subscriptionService *service.SubscriptionService) *S
|
||||
type AssignSubscriptionRequest struct {
|
||||
UserID int64 `json:"user_id" binding:"required"`
|
||||
GroupID int64 `json:"group_id" binding:"required"`
|
||||
ValidityDays int `json:"validity_days"`
|
||||
ValidityDays int `json:"validity_days" binding:"omitempty,max=36500"` // max 100 years
|
||||
Notes string `json:"notes"`
|
||||
}
|
||||
|
||||
@@ -49,13 +49,13 @@ type AssignSubscriptionRequest struct {
|
||||
type BulkAssignSubscriptionRequest struct {
|
||||
UserIDs []int64 `json:"user_ids" binding:"required,min=1"`
|
||||
GroupID int64 `json:"group_id" binding:"required"`
|
||||
ValidityDays int `json:"validity_days"`
|
||||
ValidityDays int `json:"validity_days" binding:"omitempty,max=36500"` // max 100 years
|
||||
Notes string `json:"notes"`
|
||||
}
|
||||
|
||||
// ExtendSubscriptionRequest represents extend subscription request
|
||||
type ExtendSubscriptionRequest struct {
|
||||
Days int `json:"days" binding:"required,min=1"`
|
||||
Days int `json:"days" binding:"required,min=1,max=36500"` // max 100 years
|
||||
}
|
||||
|
||||
// List handles listing all subscriptions with pagination and filters
|
||||
|
||||
@@ -40,7 +40,7 @@ func (h *UsageHandler) List(c *gin.Context) {
|
||||
page, pageSize := response.ParsePagination(c)
|
||||
|
||||
// Parse filters
|
||||
var userID, apiKeyID int64
|
||||
var userID, apiKeyID, accountID, groupID int64
|
||||
if userIDStr := c.Query("user_id"); userIDStr != "" {
|
||||
id, err := strconv.ParseInt(userIDStr, 10, 64)
|
||||
if err != nil {
|
||||
@@ -59,6 +59,47 @@ func (h *UsageHandler) List(c *gin.Context) {
|
||||
apiKeyID = id
|
||||
}
|
||||
|
||||
if accountIDStr := c.Query("account_id"); accountIDStr != "" {
|
||||
id, err := strconv.ParseInt(accountIDStr, 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid account_id")
|
||||
return
|
||||
}
|
||||
accountID = id
|
||||
}
|
||||
|
||||
if groupIDStr := c.Query("group_id"); groupIDStr != "" {
|
||||
id, err := strconv.ParseInt(groupIDStr, 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group_id")
|
||||
return
|
||||
}
|
||||
groupID = id
|
||||
}
|
||||
|
||||
model := c.Query("model")
|
||||
|
||||
var stream *bool
|
||||
if streamStr := c.Query("stream"); streamStr != "" {
|
||||
val, err := strconv.ParseBool(streamStr)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid stream value, use true or false")
|
||||
return
|
||||
}
|
||||
stream = &val
|
||||
}
|
||||
|
||||
var billingType *int8
|
||||
if billingTypeStr := c.Query("billing_type"); billingTypeStr != "" {
|
||||
val, err := strconv.ParseInt(billingTypeStr, 10, 8)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid billing_type")
|
||||
return
|
||||
}
|
||||
bt := int8(val)
|
||||
billingType = &bt
|
||||
}
|
||||
|
||||
// Parse date range
|
||||
var startTime, endTime *time.Time
|
||||
if startDateStr := c.Query("start_date"); startDateStr != "" {
|
||||
@@ -83,10 +124,15 @@ func (h *UsageHandler) List(c *gin.Context) {
|
||||
|
||||
params := pagination.PaginationParams{Page: page, PageSize: pageSize}
|
||||
filters := usagestats.UsageLogFilters{
|
||||
UserID: userID,
|
||||
ApiKeyID: apiKeyID,
|
||||
StartTime: startTime,
|
||||
EndTime: endTime,
|
||||
UserID: userID,
|
||||
ApiKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
Model: model,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
StartTime: startTime,
|
||||
EndTime: endTime,
|
||||
}
|
||||
|
||||
records, result, err := h.usageService.ListWithFilters(c.Request.Context(), params, filters)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
@@ -11,13 +12,15 @@ import (
|
||||
|
||||
// AuthHandler handles authentication-related requests
|
||||
type AuthHandler struct {
|
||||
cfg *config.Config
|
||||
authService *service.AuthService
|
||||
userService *service.UserService
|
||||
}
|
||||
|
||||
// NewAuthHandler creates a new AuthHandler
|
||||
func NewAuthHandler(authService *service.AuthService, userService *service.UserService) *AuthHandler {
|
||||
func NewAuthHandler(cfg *config.Config, authService *service.AuthService, userService *service.UserService) *AuthHandler {
|
||||
return &AuthHandler{
|
||||
cfg: cfg,
|
||||
authService: authService,
|
||||
userService: userService,
|
||||
}
|
||||
@@ -157,5 +160,15 @@ func (h *AuthHandler) GetCurrentUser(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, dto.UserFromService(user))
|
||||
type UserResponse struct {
|
||||
*dto.User
|
||||
RunMode string `json:"run_mode"`
|
||||
}
|
||||
|
||||
runMode := config.RunModeStandard
|
||||
if h.cfg != nil {
|
||||
runMode = h.cfg.RunMode
|
||||
}
|
||||
|
||||
response.Success(c, UserResponse{User: dto.UserFromService(user), RunMode: runMode})
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ package handler
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
@@ -20,27 +21,30 @@ import (
|
||||
|
||||
// GatewayHandler handles API gateway requests
|
||||
type GatewayHandler struct {
|
||||
gatewayService *service.GatewayService
|
||||
geminiCompatService *service.GeminiMessagesCompatService
|
||||
userService *service.UserService
|
||||
billingCacheService *service.BillingCacheService
|
||||
concurrencyHelper *ConcurrencyHelper
|
||||
gatewayService *service.GatewayService
|
||||
geminiCompatService *service.GeminiMessagesCompatService
|
||||
antigravityGatewayService *service.AntigravityGatewayService
|
||||
userService *service.UserService
|
||||
billingCacheService *service.BillingCacheService
|
||||
concurrencyHelper *ConcurrencyHelper
|
||||
}
|
||||
|
||||
// NewGatewayHandler creates a new GatewayHandler
|
||||
func NewGatewayHandler(
|
||||
gatewayService *service.GatewayService,
|
||||
geminiCompatService *service.GeminiMessagesCompatService,
|
||||
antigravityGatewayService *service.AntigravityGatewayService,
|
||||
userService *service.UserService,
|
||||
concurrencyService *service.ConcurrencyService,
|
||||
billingCacheService *service.BillingCacheService,
|
||||
) *GatewayHandler {
|
||||
return &GatewayHandler{
|
||||
gatewayService: gatewayService,
|
||||
geminiCompatService: geminiCompatService,
|
||||
userService: userService,
|
||||
billingCacheService: billingCacheService,
|
||||
concurrencyHelper: NewConcurrencyHelper(concurrencyService, SSEPingFormatClaude),
|
||||
gatewayService: gatewayService,
|
||||
geminiCompatService: geminiCompatService,
|
||||
antigravityGatewayService: antigravityGatewayService,
|
||||
userService: userService,
|
||||
billingCacheService: billingCacheService,
|
||||
concurrencyHelper: NewConcurrencyHelper(concurrencyService, SSEPingFormatClaude),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -122,71 +126,176 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
// 计算粘性会话hash
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(body)
|
||||
|
||||
// 获取平台:优先使用强制平台(/antigravity 路由,中间件已设置 request.Context),否则使用分组平台
|
||||
platform := ""
|
||||
if apiKey.Group != nil {
|
||||
if forcePlatform, ok := middleware2.GetForcePlatformFromContext(c); ok {
|
||||
platform = forcePlatform
|
||||
} else if apiKey.Group != nil {
|
||||
platform = apiKey.Group.Platform
|
||||
}
|
||||
|
||||
// 选择支持该模型的账号
|
||||
var account *service.Account
|
||||
if platform == service.PlatformGemini {
|
||||
account, err = h.geminiCompatService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, sessionHash, req.Model)
|
||||
} else {
|
||||
account, err = h.gatewayService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, sessionHash, req.Model)
|
||||
}
|
||||
if err != nil {
|
||||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
||||
return
|
||||
const maxAccountSwitches = 3
|
||||
switchCount := 0
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
lastFailoverStatus := 0
|
||||
|
||||
for {
|
||||
account, err := h.geminiCompatService.SelectAccountForModelWithExclusions(c.Request.Context(), apiKey.GroupID, sessionHash, req.Model, failedAccountIDs)
|
||||
if err != nil {
|
||||
if len(failedAccountIDs) == 0 {
|
||||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
||||
return
|
||||
}
|
||||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||||
return
|
||||
}
|
||||
|
||||
// 检查预热请求拦截(在账号选择后、转发前检查)
|
||||
if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
|
||||
if req.Stream {
|
||||
sendMockWarmupStream(c, req.Model)
|
||||
} else {
|
||||
sendMockWarmupResponse(c, req.Model)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// 3. 获取账号并发槽位
|
||||
accountReleaseFunc, err := h.concurrencyHelper.AcquireAccountSlotWithWait(c, account.ID, account.Concurrency, req.Stream, &streamStarted)
|
||||
if err != nil {
|
||||
log.Printf("Account concurrency acquire failed: %v", err)
|
||||
h.handleConcurrencyError(c, err, "account", streamStarted)
|
||||
return
|
||||
}
|
||||
|
||||
// 转发请求 - 根据账号平台分流
|
||||
var result *service.ForwardResult
|
||||
if account.Platform == service.PlatformAntigravity {
|
||||
result, err = h.antigravityGatewayService.ForwardGemini(c.Request.Context(), c, account, req.Model, "generateContent", req.Stream, body)
|
||||
} else {
|
||||
result, err = h.geminiCompatService.Forward(c.Request.Context(), c, account, body)
|
||||
}
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
if err != nil {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
if switchCount >= maxAccountSwitches {
|
||||
lastFailoverStatus = failoverErr.StatusCode
|
||||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||||
return
|
||||
}
|
||||
lastFailoverStatus = failoverErr.StatusCode
|
||||
switchCount++
|
||||
log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
|
||||
continue
|
||||
}
|
||||
// 错误响应已在Forward中处理,这里只记录日志
|
||||
log.Printf("Forward request failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 异步记录使用量(subscription已在函数开头获取)
|
||||
go func(result *service.ForwardResult, usedAccount *service.Account) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
ApiKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: usedAccount,
|
||||
Subscription: subscription,
|
||||
}); err != nil {
|
||||
log.Printf("Record usage failed: %v", err)
|
||||
}
|
||||
}(result, account)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 检查预热请求拦截(在账号选择后、转发前检查)
|
||||
if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
|
||||
if req.Stream {
|
||||
sendMockWarmupStream(c, req.Model)
|
||||
const maxAccountSwitches = 10
|
||||
switchCount := 0
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
lastFailoverStatus := 0
|
||||
|
||||
for {
|
||||
// 选择支持该模型的账号
|
||||
account, err := h.gatewayService.SelectAccountForModelWithExclusions(c.Request.Context(), apiKey.GroupID, sessionHash, req.Model, failedAccountIDs)
|
||||
if err != nil {
|
||||
if len(failedAccountIDs) == 0 {
|
||||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
||||
return
|
||||
}
|
||||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||||
return
|
||||
}
|
||||
|
||||
// 检查预热请求拦截(在账号选择后、转发前检查)
|
||||
if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
|
||||
if req.Stream {
|
||||
sendMockWarmupStream(c, req.Model)
|
||||
} else {
|
||||
sendMockWarmupResponse(c, req.Model)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// 3. 获取账号并发槽位
|
||||
accountReleaseFunc, err := h.concurrencyHelper.AcquireAccountSlotWithWait(c, account.ID, account.Concurrency, req.Stream, &streamStarted)
|
||||
if err != nil {
|
||||
log.Printf("Account concurrency acquire failed: %v", err)
|
||||
h.handleConcurrencyError(c, err, "account", streamStarted)
|
||||
return
|
||||
}
|
||||
|
||||
// 转发请求 - 根据账号平台分流
|
||||
var result *service.ForwardResult
|
||||
if account.Platform == service.PlatformAntigravity {
|
||||
result, err = h.antigravityGatewayService.Forward(c.Request.Context(), c, account, body)
|
||||
} else {
|
||||
sendMockWarmupResponse(c, req.Model)
|
||||
result, err = h.gatewayService.Forward(c.Request.Context(), c, account, body)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// 3. 获取账号并发槽位
|
||||
accountReleaseFunc, err := h.concurrencyHelper.AcquireAccountSlotWithWait(c, account.ID, account.Concurrency, req.Stream, &streamStarted)
|
||||
if err != nil {
|
||||
log.Printf("Account concurrency acquire failed: %v", err)
|
||||
h.handleConcurrencyError(c, err, "account", streamStarted)
|
||||
return
|
||||
}
|
||||
if accountReleaseFunc != nil {
|
||||
defer accountReleaseFunc()
|
||||
}
|
||||
|
||||
// 转发请求
|
||||
var result *service.ForwardResult
|
||||
if platform == service.PlatformGemini {
|
||||
result, err = h.geminiCompatService.Forward(c.Request.Context(), c, account, body)
|
||||
} else {
|
||||
result, err = h.gatewayService.Forward(c.Request.Context(), c, account, body)
|
||||
}
|
||||
if err != nil {
|
||||
// 错误响应已在Forward中处理,这里只记录日志
|
||||
log.Printf("Forward request failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 异步记录使用量(subscription已在函数开头获取)
|
||||
go func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
ApiKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
}); err != nil {
|
||||
log.Printf("Record usage failed: %v", err)
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
}()
|
||||
if err != nil {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
if switchCount >= maxAccountSwitches {
|
||||
lastFailoverStatus = failoverErr.StatusCode
|
||||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||||
return
|
||||
}
|
||||
lastFailoverStatus = failoverErr.StatusCode
|
||||
switchCount++
|
||||
log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
|
||||
continue
|
||||
}
|
||||
// 错误响应已在Forward中处理,这里只记录日志
|
||||
log.Printf("Forward request failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 异步记录使用量(subscription已在函数开头获取)
|
||||
go func(result *service.ForwardResult, usedAccount *service.Account) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
ApiKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: usedAccount,
|
||||
Subscription: subscription,
|
||||
}); err != nil {
|
||||
log.Printf("Record usage failed: %v", err)
|
||||
}
|
||||
}(result, account)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Models handles listing available models
|
||||
@@ -314,6 +423,28 @@ func (h *GatewayHandler) handleConcurrencyError(c *gin.Context, err error, slotT
|
||||
fmt.Sprintf("Concurrency limit exceeded for %s, please retry later", slotType), streamStarted)
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) handleFailoverExhausted(c *gin.Context, statusCode int, streamStarted bool) {
|
||||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) mapUpstreamError(statusCode int) (int, string, string) {
|
||||
switch statusCode {
|
||||
case 401:
|
||||
return http.StatusBadGateway, "upstream_error", "Upstream authentication failed, please contact administrator"
|
||||
case 403:
|
||||
return http.StatusBadGateway, "upstream_error", "Upstream access forbidden, please contact administrator"
|
||||
case 429:
|
||||
return http.StatusTooManyRequests, "rate_limit_error", "Upstream rate limit exceeded, please retry later"
|
||||
case 529:
|
||||
return http.StatusServiceUnavailable, "overloaded_error", "Upstream service overloaded, please retry later"
|
||||
case 500, 502, 503, 504:
|
||||
return http.StatusBadGateway, "upstream_error", "Upstream service temporarily unavailable"
|
||||
default:
|
||||
return http.StatusBadGateway, "upstream_error", "Upstream request failed"
|
||||
}
|
||||
}
|
||||
|
||||
// handleStreamingAwareError handles errors that may occur after streaming has started
|
||||
func (h *GatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) {
|
||||
if streamStarted {
|
||||
|
||||
@@ -2,6 +2,7 @@ package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
@@ -24,13 +25,28 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||||
googleError(c, http.StatusUnauthorized, "Invalid API key")
|
||||
return
|
||||
}
|
||||
if apiKey.Group == nil || apiKey.Group.Platform != service.PlatformGemini {
|
||||
// 检查平台:优先使用强制平台(/antigravity 路由),否则要求 gemini 分组
|
||||
forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c)
|
||||
if !hasForcePlatform && (apiKey.Group == nil || apiKey.Group.Platform != service.PlatformGemini) {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return
|
||||
}
|
||||
|
||||
// 强制 antigravity 模式:直接返回静态模型列表
|
||||
if forcePlatform == service.PlatformAntigravity {
|
||||
c.JSON(http.StatusOK, gemini.FallbackModelsList())
|
||||
return
|
||||
}
|
||||
|
||||
account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
if err != nil {
|
||||
// 没有 gemini 账户,检查是否有 antigravity 账户可用
|
||||
hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
|
||||
if hasAntigravity {
|
||||
// antigravity 账户使用静态模型列表
|
||||
c.JSON(http.StatusOK, gemini.FallbackModelsList())
|
||||
return
|
||||
}
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
|
||||
return
|
||||
}
|
||||
@@ -55,7 +71,9 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
|
||||
googleError(c, http.StatusUnauthorized, "Invalid API key")
|
||||
return
|
||||
}
|
||||
if apiKey.Group == nil || apiKey.Group.Platform != service.PlatformGemini {
|
||||
// 检查平台:优先使用强制平台(/antigravity 路由),否则要求 gemini 分组
|
||||
forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c)
|
||||
if !hasForcePlatform && (apiKey.Group == nil || apiKey.Group.Platform != service.PlatformGemini) {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return
|
||||
}
|
||||
@@ -66,8 +84,21 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 强制 antigravity 模式:直接返回静态模型信息
|
||||
if forcePlatform == service.PlatformAntigravity {
|
||||
c.JSON(http.StatusOK, gemini.FallbackModel(modelName))
|
||||
return
|
||||
}
|
||||
|
||||
account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID)
|
||||
if err != nil {
|
||||
// 没有 gemini 账户,检查是否有 antigravity 账户可用
|
||||
hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
|
||||
if hasAntigravity {
|
||||
// antigravity 账户使用静态模型信息
|
||||
c.JSON(http.StatusOK, gemini.FallbackModel(modelName))
|
||||
return
|
||||
}
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
|
||||
return
|
||||
}
|
||||
@@ -99,9 +130,12 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if apiKey.Group == nil || apiKey.Group.Platform != service.PlatformGemini {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return
|
||||
// 检查平台:优先使用强制平台(/antigravity 路由,中间件已设置 request.Context),否则要求 gemini 分组
|
||||
if !middleware.HasForcePlatform(c) {
|
||||
if apiKey.Group == nil || apiKey.Group.Platform != service.PlatformGemini {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
modelName, action, err := parseGeminiModelAction(strings.TrimPrefix(c.Param("modelAction"), "/"))
|
||||
@@ -158,44 +192,74 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
|
||||
// 3) select account (sticky session based on request body)
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(body)
|
||||
account, err := h.geminiCompatService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, sessionHash, modelName)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
|
||||
return
|
||||
}
|
||||
const maxAccountSwitches = 3
|
||||
switchCount := 0
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
lastFailoverStatus := 0
|
||||
|
||||
// 4) account concurrency slot
|
||||
accountReleaseFunc, err := geminiConcurrency.AcquireAccountSlotWithWait(c, account.ID, account.Concurrency, stream, &streamStarted)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusTooManyRequests, err.Error())
|
||||
return
|
||||
}
|
||||
if accountReleaseFunc != nil {
|
||||
defer accountReleaseFunc()
|
||||
}
|
||||
|
||||
// 5) forward (writes response to client)
|
||||
result, err := h.geminiCompatService.ForwardNative(c.Request.Context(), c, account, modelName, action, stream, body)
|
||||
if err != nil {
|
||||
// ForwardNative already wrote the response
|
||||
log.Printf("Gemini native forward failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 6) record usage async
|
||||
go func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
ApiKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
}); err != nil {
|
||||
log.Printf("Record usage failed: %v", err)
|
||||
for {
|
||||
account, err := h.geminiCompatService.SelectAccountForModelWithExclusions(c.Request.Context(), apiKey.GroupID, sessionHash, modelName, failedAccountIDs)
|
||||
if err != nil {
|
||||
if len(failedAccountIDs) == 0 {
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
|
||||
return
|
||||
}
|
||||
handleGeminiFailoverExhausted(c, lastFailoverStatus)
|
||||
return
|
||||
}
|
||||
}()
|
||||
|
||||
// 4) account concurrency slot
|
||||
accountReleaseFunc, err := geminiConcurrency.AcquireAccountSlotWithWait(c, account.ID, account.Concurrency, stream, &streamStarted)
|
||||
if err != nil {
|
||||
googleError(c, http.StatusTooManyRequests, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// 5) forward (根据平台分流)
|
||||
var result *service.ForwardResult
|
||||
if account.Platform == service.PlatformAntigravity {
|
||||
result, err = h.antigravityGatewayService.ForwardGemini(c.Request.Context(), c, account, modelName, action, stream, body)
|
||||
} else {
|
||||
result, err = h.geminiCompatService.ForwardNative(c.Request.Context(), c, account, modelName, action, stream, body)
|
||||
}
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
if err != nil {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
if switchCount >= maxAccountSwitches {
|
||||
lastFailoverStatus = failoverErr.StatusCode
|
||||
handleGeminiFailoverExhausted(c, lastFailoverStatus)
|
||||
return
|
||||
}
|
||||
lastFailoverStatus = failoverErr.StatusCode
|
||||
switchCount++
|
||||
log.Printf("Gemini account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
|
||||
continue
|
||||
}
|
||||
// ForwardNative already wrote the response
|
||||
log.Printf("Gemini native forward failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 6) record usage async
|
||||
go func(result *service.ForwardResult, usedAccount *service.Account) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
ApiKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: usedAccount,
|
||||
Subscription: subscription,
|
||||
}); err != nil {
|
||||
log.Printf("Record usage failed: %v", err)
|
||||
}
|
||||
}(result, account)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func parseGeminiModelAction(rest string) (model string, action string, err error) {
|
||||
@@ -217,6 +281,28 @@ func parseGeminiModelAction(rest string) (model string, action string, err error
|
||||
return "", "", &pathParseError{"invalid model action path"}
|
||||
}
|
||||
|
||||
func handleGeminiFailoverExhausted(c *gin.Context, statusCode int) {
|
||||
status, message := mapGeminiUpstreamError(statusCode)
|
||||
googleError(c, status, message)
|
||||
}
|
||||
|
||||
func mapGeminiUpstreamError(statusCode int) (int, string) {
|
||||
switch statusCode {
|
||||
case 401:
|
||||
return http.StatusBadGateway, "Upstream authentication failed, please contact administrator"
|
||||
case 403:
|
||||
return http.StatusBadGateway, "Upstream access forbidden, please contact administrator"
|
||||
case 429:
|
||||
return http.StatusTooManyRequests, "Upstream rate limit exceeded, please retry later"
|
||||
case 529:
|
||||
return http.StatusServiceUnavailable, "Upstream service overloaded, please retry later"
|
||||
case 500, 502, 503, 504:
|
||||
return http.StatusBadGateway, "Upstream service temporarily unavailable"
|
||||
default:
|
||||
return http.StatusBadGateway, "Upstream request failed"
|
||||
}
|
||||
}
|
||||
|
||||
type pathParseError struct{ msg string }
|
||||
|
||||
func (e *pathParseError) Error() string { return e.msg }
|
||||
|
||||
143
backend/internal/handler/gemini_v1beta_handler_test.go
Normal file
143
backend/internal/handler/gemini_v1beta_handler_test.go
Normal file
@@ -0,0 +1,143 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestGeminiV1BetaHandler_PlatformRoutingInvariant 文档化并验证 Handler 层的平台路由逻辑不变量
|
||||
// 该测试确保 gemini 和 antigravity 平台的路由逻辑符合预期
|
||||
func TestGeminiV1BetaHandler_PlatformRoutingInvariant(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
platform string
|
||||
expectedService string
|
||||
description string
|
||||
}{
|
||||
{
|
||||
name: "Gemini平台使用ForwardNative",
|
||||
platform: service.PlatformGemini,
|
||||
expectedService: "GeminiMessagesCompatService.ForwardNative",
|
||||
description: "Gemini OAuth 账户直接调用 Google API",
|
||||
},
|
||||
{
|
||||
name: "Antigravity平台使用ForwardGemini",
|
||||
platform: service.PlatformAntigravity,
|
||||
expectedService: "AntigravityGatewayService.ForwardGemini",
|
||||
description: "Antigravity 账户通过 CRS 中转,支持 Gemini 协议",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// 模拟 GeminiV1BetaModels 中的路由决策 (lines 199-205 in gemini_v1beta_handler.go)
|
||||
var routedService string
|
||||
if tt.platform == service.PlatformAntigravity {
|
||||
routedService = "AntigravityGatewayService.ForwardGemini"
|
||||
} else {
|
||||
routedService = "GeminiMessagesCompatService.ForwardNative"
|
||||
}
|
||||
|
||||
require.Equal(t, tt.expectedService, routedService,
|
||||
"平台 %s 应该路由到 %s: %s",
|
||||
tt.platform, tt.expectedService, tt.description)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeminiV1BetaHandler_ListModelsAntigravityFallback 验证 ListModels 的 antigravity 降级逻辑
|
||||
// 当没有 gemini 账户但有 antigravity 账户时,应返回静态模型列表
|
||||
func TestGeminiV1BetaHandler_ListModelsAntigravityFallback(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
hasGeminiAccount bool
|
||||
hasAntigravity bool
|
||||
expectedBehavior string
|
||||
}{
|
||||
{
|
||||
name: "有Gemini账户-调用ForwardAIStudioGET",
|
||||
hasGeminiAccount: true,
|
||||
hasAntigravity: false,
|
||||
expectedBehavior: "forward_to_upstream",
|
||||
},
|
||||
{
|
||||
name: "无Gemini有Antigravity-返回静态列表",
|
||||
hasGeminiAccount: false,
|
||||
hasAntigravity: true,
|
||||
expectedBehavior: "static_fallback",
|
||||
},
|
||||
{
|
||||
name: "无任何账户-返回503",
|
||||
hasGeminiAccount: false,
|
||||
hasAntigravity: false,
|
||||
expectedBehavior: "service_unavailable",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// 模拟 GeminiV1BetaListModels 的逻辑 (lines 33-44 in gemini_v1beta_handler.go)
|
||||
var behavior string
|
||||
|
||||
if tt.hasGeminiAccount {
|
||||
behavior = "forward_to_upstream"
|
||||
} else if tt.hasAntigravity {
|
||||
behavior = "static_fallback"
|
||||
} else {
|
||||
behavior = "service_unavailable"
|
||||
}
|
||||
|
||||
require.Equal(t, tt.expectedBehavior, behavior)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeminiV1BetaHandler_GetModelAntigravityFallback 验证 GetModel 的 antigravity 降级逻辑
|
||||
func TestGeminiV1BetaHandler_GetModelAntigravityFallback(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
hasGeminiAccount bool
|
||||
hasAntigravity bool
|
||||
expectedBehavior string
|
||||
}{
|
||||
{
|
||||
name: "有Gemini账户-调用ForwardAIStudioGET",
|
||||
hasGeminiAccount: true,
|
||||
hasAntigravity: false,
|
||||
expectedBehavior: "forward_to_upstream",
|
||||
},
|
||||
{
|
||||
name: "无Gemini有Antigravity-返回静态模型信息",
|
||||
hasGeminiAccount: false,
|
||||
hasAntigravity: true,
|
||||
expectedBehavior: "static_model_info",
|
||||
},
|
||||
{
|
||||
name: "无任何账户-返回503",
|
||||
hasGeminiAccount: false,
|
||||
hasAntigravity: false,
|
||||
expectedBehavior: "service_unavailable",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// 模拟 GeminiV1BetaGetModel 的逻辑 (lines 77-87 in gemini_v1beta_handler.go)
|
||||
var behavior string
|
||||
|
||||
if tt.hasGeminiAccount {
|
||||
behavior = "forward_to_upstream"
|
||||
} else if tt.hasAntigravity {
|
||||
behavior = "static_model_info"
|
||||
} else {
|
||||
behavior = "service_unavailable"
|
||||
}
|
||||
|
||||
require.Equal(t, tt.expectedBehavior, behavior)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -6,19 +6,20 @@ import (
|
||||
|
||||
// AdminHandlers contains all admin-related HTTP handlers
|
||||
type AdminHandlers struct {
|
||||
Dashboard *admin.DashboardHandler
|
||||
User *admin.UserHandler
|
||||
Group *admin.GroupHandler
|
||||
Account *admin.AccountHandler
|
||||
OAuth *admin.OAuthHandler
|
||||
OpenAIOAuth *admin.OpenAIOAuthHandler
|
||||
GeminiOAuth *admin.GeminiOAuthHandler
|
||||
Proxy *admin.ProxyHandler
|
||||
Redeem *admin.RedeemHandler
|
||||
Setting *admin.SettingHandler
|
||||
System *admin.SystemHandler
|
||||
Subscription *admin.SubscriptionHandler
|
||||
Usage *admin.UsageHandler
|
||||
Dashboard *admin.DashboardHandler
|
||||
User *admin.UserHandler
|
||||
Group *admin.GroupHandler
|
||||
Account *admin.AccountHandler
|
||||
OAuth *admin.OAuthHandler
|
||||
OpenAIOAuth *admin.OpenAIOAuthHandler
|
||||
GeminiOAuth *admin.GeminiOAuthHandler
|
||||
AntigravityOAuth *admin.AntigravityOAuthHandler
|
||||
Proxy *admin.ProxyHandler
|
||||
Redeem *admin.RedeemHandler
|
||||
Setting *admin.SettingHandler
|
||||
System *admin.SystemHandler
|
||||
Subscription *admin.SubscriptionHandler
|
||||
Usage *admin.UsageHandler
|
||||
}
|
||||
|
||||
// Handlers contains all HTTP handlers
|
||||
|
||||
@@ -3,6 +3,7 @@ package handler
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
@@ -127,49 +128,74 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
// Generate session hash (from header for OpenAI)
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c)
|
||||
|
||||
// Select account supporting the requested model
|
||||
log.Printf("[OpenAI Handler] Selecting account: groupID=%v model=%s", apiKey.GroupID, reqModel)
|
||||
account, err := h.gatewayService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, sessionHash, reqModel)
|
||||
if err != nil {
|
||||
log.Printf("[OpenAI Handler] SelectAccount failed: %v", err)
|
||||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
||||
return
|
||||
}
|
||||
log.Printf("[OpenAI Handler] Selected account: id=%d name=%s", account.ID, account.Name)
|
||||
const maxAccountSwitches = 3
|
||||
switchCount := 0
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
lastFailoverStatus := 0
|
||||
|
||||
// 3. Acquire account concurrency slot
|
||||
accountReleaseFunc, err := h.concurrencyHelper.AcquireAccountSlotWithWait(c, account.ID, account.Concurrency, reqStream, &streamStarted)
|
||||
if err != nil {
|
||||
log.Printf("Account concurrency acquire failed: %v", err)
|
||||
h.handleConcurrencyError(c, err, "account", streamStarted)
|
||||
return
|
||||
}
|
||||
if accountReleaseFunc != nil {
|
||||
defer accountReleaseFunc()
|
||||
}
|
||||
|
||||
// Forward request
|
||||
result, err := h.gatewayService.Forward(c.Request.Context(), c, account, body)
|
||||
if err != nil {
|
||||
// Error response already handled in Forward, just log
|
||||
log.Printf("Forward request failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Async record usage
|
||||
go func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
Result: result,
|
||||
ApiKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
}); err != nil {
|
||||
log.Printf("Record usage failed: %v", err)
|
||||
for {
|
||||
// Select account supporting the requested model
|
||||
log.Printf("[OpenAI Handler] Selecting account: groupID=%v model=%s", apiKey.GroupID, reqModel)
|
||||
account, err := h.gatewayService.SelectAccountForModelWithExclusions(c.Request.Context(), apiKey.GroupID, sessionHash, reqModel, failedAccountIDs)
|
||||
if err != nil {
|
||||
log.Printf("[OpenAI Handler] SelectAccount failed: %v", err)
|
||||
if len(failedAccountIDs) == 0 {
|
||||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
||||
return
|
||||
}
|
||||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||||
return
|
||||
}
|
||||
}()
|
||||
log.Printf("[OpenAI Handler] Selected account: id=%d name=%s", account.ID, account.Name)
|
||||
|
||||
// 3. Acquire account concurrency slot
|
||||
accountReleaseFunc, err := h.concurrencyHelper.AcquireAccountSlotWithWait(c, account.ID, account.Concurrency, reqStream, &streamStarted)
|
||||
if err != nil {
|
||||
log.Printf("Account concurrency acquire failed: %v", err)
|
||||
h.handleConcurrencyError(c, err, "account", streamStarted)
|
||||
return
|
||||
}
|
||||
|
||||
// Forward request
|
||||
result, err := h.gatewayService.Forward(c.Request.Context(), c, account, body)
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
if err != nil {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
if switchCount >= maxAccountSwitches {
|
||||
lastFailoverStatus = failoverErr.StatusCode
|
||||
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
||||
return
|
||||
}
|
||||
lastFailoverStatus = failoverErr.StatusCode
|
||||
switchCount++
|
||||
log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
|
||||
continue
|
||||
}
|
||||
// Error response already handled in Forward, just log
|
||||
log.Printf("Forward request failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Async record usage
|
||||
go func(result *service.OpenAIForwardResult, usedAccount *service.Account) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
Result: result,
|
||||
ApiKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: usedAccount,
|
||||
Subscription: subscription,
|
||||
}); err != nil {
|
||||
log.Printf("Record usage failed: %v", err)
|
||||
}
|
||||
}(result, account)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// handleConcurrencyError handles concurrency-related errors with proper 429 response
|
||||
@@ -178,6 +204,28 @@ func (h *OpenAIGatewayHandler) handleConcurrencyError(c *gin.Context, err error,
|
||||
fmt.Sprintf("Concurrency limit exceeded for %s, please retry later", slotType), streamStarted)
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, statusCode int, streamStarted bool) {
|
||||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) mapUpstreamError(statusCode int) (int, string, string) {
|
||||
switch statusCode {
|
||||
case 401:
|
||||
return http.StatusBadGateway, "upstream_error", "Upstream authentication failed, please contact administrator"
|
||||
case 403:
|
||||
return http.StatusBadGateway, "upstream_error", "Upstream access forbidden, please contact administrator"
|
||||
case 429:
|
||||
return http.StatusTooManyRequests, "rate_limit_error", "Upstream rate limit exceeded, please retry later"
|
||||
case 529:
|
||||
return http.StatusServiceUnavailable, "upstream_error", "Upstream service overloaded, please retry later"
|
||||
case 500, 502, 503, 504:
|
||||
return http.StatusBadGateway, "upstream_error", "Upstream service temporarily unavailable"
|
||||
default:
|
||||
return http.StatusBadGateway, "upstream_error", "Upstream request failed"
|
||||
}
|
||||
}
|
||||
|
||||
// handleStreamingAwareError handles errors that may occur after streaming has started
|
||||
func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) {
|
||||
if streamStarted {
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
@@ -61,16 +62,64 @@ func (h *UsageHandler) List(c *gin.Context) {
|
||||
apiKeyID = id
|
||||
}
|
||||
|
||||
params := pagination.PaginationParams{Page: page, PageSize: pageSize}
|
||||
var records []service.UsageLog
|
||||
var result *pagination.PaginationResult
|
||||
var err error
|
||||
// Parse additional filters
|
||||
model := c.Query("model")
|
||||
|
||||
if apiKeyID > 0 {
|
||||
records, result, err = h.usageService.ListByApiKey(c.Request.Context(), apiKeyID, params)
|
||||
} else {
|
||||
records, result, err = h.usageService.ListByUser(c.Request.Context(), subject.UserID, params)
|
||||
var stream *bool
|
||||
if streamStr := c.Query("stream"); streamStr != "" {
|
||||
val, err := strconv.ParseBool(streamStr)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid stream value, use true or false")
|
||||
return
|
||||
}
|
||||
stream = &val
|
||||
}
|
||||
|
||||
var billingType *int8
|
||||
if billingTypeStr := c.Query("billing_type"); billingTypeStr != "" {
|
||||
val, err := strconv.ParseInt(billingTypeStr, 10, 8)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid billing_type")
|
||||
return
|
||||
}
|
||||
bt := int8(val)
|
||||
billingType = &bt
|
||||
}
|
||||
|
||||
// Parse date range
|
||||
var startTime, endTime *time.Time
|
||||
if startDateStr := c.Query("start_date"); startDateStr != "" {
|
||||
t, err := timezone.ParseInLocation("2006-01-02", startDateStr)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid start_date format, use YYYY-MM-DD")
|
||||
return
|
||||
}
|
||||
startTime = &t
|
||||
}
|
||||
|
||||
if endDateStr := c.Query("end_date"); endDateStr != "" {
|
||||
t, err := timezone.ParseInLocation("2006-01-02", endDateStr)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid end_date format, use YYYY-MM-DD")
|
||||
return
|
||||
}
|
||||
// Set end time to end of day
|
||||
t = t.Add(24*time.Hour - time.Nanosecond)
|
||||
endTime = &t
|
||||
}
|
||||
|
||||
params := pagination.PaginationParams{Page: page, PageSize: pageSize}
|
||||
filters := usagestats.UsageLogFilters{
|
||||
UserID: subject.UserID, // Always filter by current user for security
|
||||
ApiKeyID: apiKeyID,
|
||||
Model: model,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
StartTime: startTime,
|
||||
EndTime: endTime,
|
||||
}
|
||||
|
||||
records, result, err := h.usageService.ListWithFilters(c.Request.Context(), params, filters)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
@@ -322,24 +371,16 @@ func (h *UsageHandler) DashboardApiKeysUsage(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// Verify ownership of all requested API keys
|
||||
userApiKeys, _, err := h.apiKeyService.List(c.Request.Context(), subject.UserID, pagination.PaginationParams{Page: 1, PageSize: 1000})
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
// Limit the number of API key IDs to prevent SQL parameter overflow
|
||||
if len(req.ApiKeyIDs) > 100 {
|
||||
response.BadRequest(c, "Too many API key IDs (maximum 100 allowed)")
|
||||
return
|
||||
}
|
||||
|
||||
userApiKeyIDs := make(map[int64]bool)
|
||||
for _, key := range userApiKeys {
|
||||
userApiKeyIDs[key.ID] = true
|
||||
}
|
||||
|
||||
// Filter to only include user's own API keys
|
||||
validApiKeyIDs := make([]int64, 0)
|
||||
for _, id := range req.ApiKeyIDs {
|
||||
if userApiKeyIDs[id] {
|
||||
validApiKeyIDs = append(validApiKeyIDs, id)
|
||||
}
|
||||
validApiKeyIDs, err := h.apiKeyService.VerifyOwnership(c.Request.Context(), subject.UserID, req.ApiKeyIDs)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
if len(validApiKeyIDs) == 0 {
|
||||
|
||||
@@ -16,6 +16,7 @@ func ProvideAdminHandlers(
|
||||
oauthHandler *admin.OAuthHandler,
|
||||
openaiOAuthHandler *admin.OpenAIOAuthHandler,
|
||||
geminiOAuthHandler *admin.GeminiOAuthHandler,
|
||||
antigravityOAuthHandler *admin.AntigravityOAuthHandler,
|
||||
proxyHandler *admin.ProxyHandler,
|
||||
redeemHandler *admin.RedeemHandler,
|
||||
settingHandler *admin.SettingHandler,
|
||||
@@ -24,19 +25,20 @@ func ProvideAdminHandlers(
|
||||
usageHandler *admin.UsageHandler,
|
||||
) *AdminHandlers {
|
||||
return &AdminHandlers{
|
||||
Dashboard: dashboardHandler,
|
||||
User: userHandler,
|
||||
Group: groupHandler,
|
||||
Account: accountHandler,
|
||||
OAuth: oauthHandler,
|
||||
OpenAIOAuth: openaiOAuthHandler,
|
||||
GeminiOAuth: geminiOAuthHandler,
|
||||
Proxy: proxyHandler,
|
||||
Redeem: redeemHandler,
|
||||
Setting: settingHandler,
|
||||
System: systemHandler,
|
||||
Subscription: subscriptionHandler,
|
||||
Usage: usageHandler,
|
||||
Dashboard: dashboardHandler,
|
||||
User: userHandler,
|
||||
Group: groupHandler,
|
||||
Account: accountHandler,
|
||||
OAuth: oauthHandler,
|
||||
OpenAIOAuth: openaiOAuthHandler,
|
||||
GeminiOAuth: geminiOAuthHandler,
|
||||
AntigravityOAuth: antigravityOAuthHandler,
|
||||
Proxy: proxyHandler,
|
||||
Redeem: redeemHandler,
|
||||
Setting: settingHandler,
|
||||
System: systemHandler,
|
||||
Subscription: subscriptionHandler,
|
||||
Usage: usageHandler,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -98,6 +100,7 @@ var ProviderSet = wire.NewSet(
|
||||
admin.NewOAuthHandler,
|
||||
admin.NewOpenAIOAuthHandler,
|
||||
admin.NewGeminiOAuthHandler,
|
||||
admin.NewAntigravityOAuthHandler,
|
||||
admin.NewProxyHandler,
|
||||
admin.NewRedeemHandler,
|
||||
admin.NewSettingHandler,
|
||||
|
||||
@@ -30,7 +30,7 @@ func InitDB(cfg *config.Config) (*gorm.DB, error) {
|
||||
|
||||
// 自动迁移(始终执行,确保数据库结构与代码同步)
|
||||
// GORM 的 AutoMigrate 只会添加新字段,不会删除或修改已有字段,是安全的
|
||||
if err := repository.AutoMigrate(db); err != nil {
|
||||
if err := repository.AutoMigrate(db, cfg.RunMode); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
|
||||
740
backend/internal/integration/e2e_gateway_test.go
Normal file
740
backend/internal/integration/e2e_gateway_test.go
Normal file
@@ -0,0 +1,740 @@
|
||||
//go:build e2e
|
||||
|
||||
package integration
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
baseURL = getEnv("BASE_URL", "http://localhost:8080")
|
||||
// ENDPOINT_PREFIX: 端点前缀,支持混合模式和非混合模式测试
|
||||
// - "" (默认): 使用 /v1/messages, /v1beta/models(混合模式,可调度 antigravity 账户)
|
||||
// - "/antigravity": 使用 /antigravity/v1/messages, /antigravity/v1beta/models(非混合模式,仅 antigravity 账户)
|
||||
endpointPrefix = getEnv("ENDPOINT_PREFIX", "")
|
||||
claudeAPIKey = "sk-8e572bc3b3de92ace4f41f4256c28600ca11805732a7b693b5c44741346bbbb3"
|
||||
geminiAPIKey = "sk-5950197a2085b38bbe5a1b229cc02b8ece914963fc44cacc06d497ae8b87410f"
|
||||
testInterval = 1 * time.Second // 测试间隔,防止限流
|
||||
)
|
||||
|
||||
func getEnv(key, defaultVal string) string {
|
||||
if v := os.Getenv(key); v != "" {
|
||||
return v
|
||||
}
|
||||
return defaultVal
|
||||
}
|
||||
|
||||
// Claude 模型列表
|
||||
var claudeModels = []string{
|
||||
// Opus 系列
|
||||
"claude-opus-4-5-thinking", // 直接支持
|
||||
"claude-opus-4", // 映射到 claude-opus-4-5-thinking
|
||||
"claude-opus-4-5-20251101", // 映射到 claude-opus-4-5-thinking
|
||||
// Sonnet 系列
|
||||
"claude-sonnet-4-5", // 直接支持
|
||||
"claude-sonnet-4-5-thinking", // 直接支持
|
||||
"claude-sonnet-4-5-20250929", // 映射到 claude-sonnet-4-5-thinking
|
||||
"claude-3-5-sonnet-20241022", // 映射到 claude-sonnet-4-5
|
||||
// Haiku 系列(映射到 gemini-3-flash)
|
||||
"claude-haiku-4",
|
||||
"claude-haiku-4-5",
|
||||
"claude-haiku-4-5-20251001",
|
||||
"claude-3-haiku-20240307",
|
||||
}
|
||||
|
||||
// Gemini 模型列表
|
||||
var geminiModels = []string{
|
||||
"gemini-2.5-flash",
|
||||
"gemini-2.5-flash-lite",
|
||||
"gemini-3-flash",
|
||||
"gemini-3-pro-low",
|
||||
}
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
mode := "混合模式"
|
||||
if endpointPrefix != "" {
|
||||
mode = "Antigravity 模式"
|
||||
}
|
||||
fmt.Printf("\n🚀 E2E Gateway Tests - %s (prefix=%q, %s)\n\n", baseURL, endpointPrefix, mode)
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
|
||||
// TestClaudeModelsList 测试 GET /v1/models
|
||||
func TestClaudeModelsList(t *testing.T) {
|
||||
url := baseURL + endpointPrefix + "/v1/models"
|
||||
|
||||
req, _ := http.NewRequest("GET", url, nil)
|
||||
req.Header.Set("Authorization", "Bearer "+claudeAPIKey)
|
||||
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("请求失败: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("HTTP %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
if result["object"] != "list" {
|
||||
t.Errorf("期望 object=list, 得到 %v", result["object"])
|
||||
}
|
||||
|
||||
data, ok := result["data"].([]any)
|
||||
if !ok {
|
||||
t.Fatal("响应缺少 data 数组")
|
||||
}
|
||||
t.Logf("✅ 返回 %d 个模型", len(data))
|
||||
}
|
||||
|
||||
// TestGeminiModelsList 测试 GET /v1beta/models
|
||||
func TestGeminiModelsList(t *testing.T) {
|
||||
url := baseURL + endpointPrefix + "/v1beta/models"
|
||||
|
||||
req, _ := http.NewRequest("GET", url, nil)
|
||||
req.Header.Set("Authorization", "Bearer "+geminiAPIKey)
|
||||
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("请求失败: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("HTTP %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
models, ok := result["models"].([]any)
|
||||
if !ok {
|
||||
t.Fatal("响应缺少 models 数组")
|
||||
}
|
||||
t.Logf("✅ 返回 %d 个模型", len(models))
|
||||
}
|
||||
|
||||
// TestClaudeMessages 测试 Claude /v1/messages 接口
|
||||
func TestClaudeMessages(t *testing.T) {
|
||||
for i, model := range claudeModels {
|
||||
if i > 0 {
|
||||
time.Sleep(testInterval)
|
||||
}
|
||||
t.Run(model+"_非流式", func(t *testing.T) {
|
||||
testClaudeMessage(t, model, false)
|
||||
})
|
||||
time.Sleep(testInterval)
|
||||
t.Run(model+"_流式", func(t *testing.T) {
|
||||
testClaudeMessage(t, model, true)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func testClaudeMessage(t *testing.T, model string, stream bool) {
|
||||
url := baseURL + endpointPrefix + "/v1/messages"
|
||||
|
||||
payload := map[string]any{
|
||||
"model": model,
|
||||
"max_tokens": 50,
|
||||
"stream": stream,
|
||||
"messages": []map[string]string{
|
||||
{"role": "user", "content": "Say 'hello' in one word."},
|
||||
},
|
||||
}
|
||||
body, _ := json.Marshal(payload)
|
||||
|
||||
req, _ := http.NewRequest("POST", url, bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+claudeAPIKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
|
||||
client := &http.Client{Timeout: 60 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("请求失败: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("HTTP %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
if stream {
|
||||
// 流式:读取 SSE 事件
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
eventCount := 0
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
eventCount++
|
||||
if eventCount >= 3 {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if eventCount == 0 {
|
||||
t.Fatal("未收到任何 SSE 事件")
|
||||
}
|
||||
t.Logf("✅ 收到 %d+ 个 SSE 事件", eventCount)
|
||||
} else {
|
||||
// 非流式:解析 JSON 响应
|
||||
var result map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
if result["type"] != "message" {
|
||||
t.Errorf("期望 type=message, 得到 %v", result["type"])
|
||||
}
|
||||
t.Logf("✅ 收到消息响应 id=%v", result["id"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeminiGenerateContent 测试 Gemini /v1beta/models/:model 接口
|
||||
func TestGeminiGenerateContent(t *testing.T) {
|
||||
for i, model := range geminiModels {
|
||||
if i > 0 {
|
||||
time.Sleep(testInterval)
|
||||
}
|
||||
t.Run(model+"_非流式", func(t *testing.T) {
|
||||
testGeminiGenerate(t, model, false)
|
||||
})
|
||||
time.Sleep(testInterval)
|
||||
t.Run(model+"_流式", func(t *testing.T) {
|
||||
testGeminiGenerate(t, model, true)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func testGeminiGenerate(t *testing.T, model string, stream bool) {
|
||||
action := "generateContent"
|
||||
if stream {
|
||||
action = "streamGenerateContent"
|
||||
}
|
||||
url := fmt.Sprintf("%s%s/v1beta/models/%s:%s", baseURL, endpointPrefix, model, action)
|
||||
if stream {
|
||||
url += "?alt=sse"
|
||||
}
|
||||
|
||||
payload := map[string]any{
|
||||
"contents": []map[string]any{
|
||||
{
|
||||
"role": "user",
|
||||
"parts": []map[string]string{
|
||||
{"text": "Say 'hello' in one word."},
|
||||
},
|
||||
},
|
||||
},
|
||||
"generationConfig": map[string]int{
|
||||
"maxOutputTokens": 50,
|
||||
},
|
||||
}
|
||||
body, _ := json.Marshal(payload)
|
||||
|
||||
req, _ := http.NewRequest("POST", url, bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+geminiAPIKey)
|
||||
|
||||
client := &http.Client{Timeout: 60 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("请求失败: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("HTTP %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
if stream {
|
||||
// 流式:读取 SSE 事件
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
eventCount := 0
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
eventCount++
|
||||
if eventCount >= 3 {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if eventCount == 0 {
|
||||
t.Fatal("未收到任何 SSE 事件")
|
||||
}
|
||||
t.Logf("✅ 收到 %d+ 个 SSE 事件", eventCount)
|
||||
} else {
|
||||
// 非流式:解析 JSON 响应
|
||||
var result map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
if _, ok := result["candidates"]; !ok {
|
||||
t.Error("响应缺少 candidates 字段")
|
||||
}
|
||||
t.Log("✅ 收到 candidates 响应")
|
||||
}
|
||||
}
|
||||
|
||||
// TestClaudeMessagesWithComplexTools 测试带复杂工具 schema 的请求
|
||||
// 模拟 Claude Code 发送的请求,包含需要清理的 JSON Schema 字段
|
||||
func TestClaudeMessagesWithComplexTools(t *testing.T) {
|
||||
// 测试模型列表(只测试几个代表性模型)
|
||||
models := []string{
|
||||
"claude-opus-4-5-20251101", // Claude 模型
|
||||
"claude-haiku-4-5-20251001", // 映射到 Gemini
|
||||
}
|
||||
|
||||
for i, model := range models {
|
||||
if i > 0 {
|
||||
time.Sleep(testInterval)
|
||||
}
|
||||
t.Run(model+"_复杂工具", func(t *testing.T) {
|
||||
testClaudeMessageWithTools(t, model)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func testClaudeMessageWithTools(t *testing.T, model string) {
|
||||
url := baseURL + endpointPrefix + "/v1/messages"
|
||||
|
||||
// 构造包含复杂 schema 的工具定义(模拟 Claude Code 的工具)
|
||||
// 这些字段需要被 cleanJSONSchema 清理
|
||||
tools := []map[string]any{
|
||||
{
|
||||
"name": "read_file",
|
||||
"description": "Read file contents",
|
||||
"input_schema": map[string]any{
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"path": map[string]any{
|
||||
"type": "string",
|
||||
"description": "File path",
|
||||
"minLength": 1,
|
||||
"maxLength": 4096,
|
||||
"pattern": "^[^\\x00]+$",
|
||||
},
|
||||
"encoding": map[string]any{
|
||||
"type": []string{"string", "null"},
|
||||
"default": "utf-8",
|
||||
"enum": []string{"utf-8", "ascii", "latin-1"},
|
||||
},
|
||||
},
|
||||
"required": []string{"path"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "write_file",
|
||||
"description": "Write content to file",
|
||||
"input_schema": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"path": map[string]any{
|
||||
"type": "string",
|
||||
"minLength": 1,
|
||||
},
|
||||
"content": map[string]any{
|
||||
"type": "string",
|
||||
"maxLength": 1048576,
|
||||
},
|
||||
},
|
||||
"required": []string{"path", "content"},
|
||||
"additionalProperties": false,
|
||||
"strict": true,
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "list_files",
|
||||
"description": "List files in directory",
|
||||
"input_schema": map[string]any{
|
||||
"$id": "https://example.com/list-files.schema.json",
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"directory": map[string]any{
|
||||
"type": "string",
|
||||
},
|
||||
"patterns": map[string]any{
|
||||
"type": "array",
|
||||
"items": map[string]any{
|
||||
"type": "string",
|
||||
"minLength": 1,
|
||||
},
|
||||
"minItems": 1,
|
||||
"maxItems": 100,
|
||||
"uniqueItems": true,
|
||||
},
|
||||
"recursive": map[string]any{
|
||||
"type": "boolean",
|
||||
"default": false,
|
||||
},
|
||||
},
|
||||
"required": []string{"directory"},
|
||||
"additionalProperties": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "search_code",
|
||||
"description": "Search code in files",
|
||||
"input_schema": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{
|
||||
"type": "string",
|
||||
"minLength": 1,
|
||||
"format": "regex",
|
||||
},
|
||||
"max_results": map[string]any{
|
||||
"type": "integer",
|
||||
"minimum": 1,
|
||||
"maximum": 1000,
|
||||
"exclusiveMinimum": 0,
|
||||
"default": 100,
|
||||
},
|
||||
},
|
||||
"required": []string{"query"},
|
||||
"additionalProperties": false,
|
||||
"examples": []map[string]any{
|
||||
{"query": "function.*test", "max_results": 50},
|
||||
},
|
||||
},
|
||||
},
|
||||
// 测试 required 引用不存在的属性(应被自动过滤)
|
||||
{
|
||||
"name": "invalid_required_tool",
|
||||
"description": "Tool with invalid required field",
|
||||
"input_schema": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"name": map[string]any{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
// "nonexistent_field" 不存在于 properties 中,应被过滤掉
|
||||
"required": []string{"name", "nonexistent_field"},
|
||||
},
|
||||
},
|
||||
// 测试没有 properties 的 schema(应自动添加空 properties)
|
||||
{
|
||||
"name": "no_properties_tool",
|
||||
"description": "Tool without properties",
|
||||
"input_schema": map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"should_be_removed"},
|
||||
},
|
||||
},
|
||||
// 测试没有 type 的 schema(应自动添加 type: OBJECT)
|
||||
{
|
||||
"name": "no_type_tool",
|
||||
"description": "Tool without type",
|
||||
"input_schema": map[string]any{
|
||||
"properties": map[string]any{
|
||||
"value": map[string]any{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
payload := map[string]any{
|
||||
"model": model,
|
||||
"max_tokens": 100,
|
||||
"stream": false,
|
||||
"messages": []map[string]string{
|
||||
{"role": "user", "content": "List files in the current directory"},
|
||||
},
|
||||
"tools": tools,
|
||||
}
|
||||
body, _ := json.Marshal(payload)
|
||||
|
||||
req, _ := http.NewRequest("POST", url, bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+claudeAPIKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
|
||||
client := &http.Client{Timeout: 60 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("请求失败: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
|
||||
// 400 错误说明 schema 清理不完整
|
||||
if resp.StatusCode == 400 {
|
||||
t.Fatalf("Schema 清理失败,收到 400 错误: %s", string(respBody))
|
||||
}
|
||||
|
||||
// 503 可能是账号限流,不算测试失败
|
||||
if resp.StatusCode == 503 {
|
||||
t.Skipf("账号暂时不可用 (503): %s", string(respBody))
|
||||
}
|
||||
|
||||
// 429 是限流
|
||||
if resp.StatusCode == 429 {
|
||||
t.Skipf("请求被限流 (429): %s", string(respBody))
|
||||
}
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
t.Fatalf("HTTP %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
if result["type"] != "message" {
|
||||
t.Errorf("期望 type=message, 得到 %v", result["type"])
|
||||
}
|
||||
t.Logf("✅ 复杂工具 schema 测试通过, id=%v", result["id"])
|
||||
}
|
||||
|
||||
// TestClaudeMessagesWithThinkingAndTools 测试 thinking 模式下带工具调用的场景
|
||||
// 验证:当历史 assistant 消息包含 tool_use 但没有 signature 时,
|
||||
// 系统应自动添加 dummy thought_signature 避免 Gemini 400 错误
|
||||
func TestClaudeMessagesWithThinkingAndTools(t *testing.T) {
|
||||
models := []string{
|
||||
"claude-haiku-4-5-20251001", // gemini-3-flash
|
||||
}
|
||||
for i, model := range models {
|
||||
if i > 0 {
|
||||
time.Sleep(testInterval)
|
||||
}
|
||||
t.Run(model+"_thinking模式工具调用", func(t *testing.T) {
|
||||
testClaudeThinkingWithToolHistory(t, model)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func testClaudeThinkingWithToolHistory(t *testing.T, model string) {
|
||||
url := baseURL + endpointPrefix + "/v1/messages"
|
||||
|
||||
// 模拟历史对话:用户请求 → assistant 调用工具 → 工具返回 → 继续对话
|
||||
// 注意:tool_use 块故意不包含 signature,测试系统是否能正确添加 dummy signature
|
||||
payload := map[string]any{
|
||||
"model": model,
|
||||
"max_tokens": 200,
|
||||
"stream": false,
|
||||
// 开启 thinking 模式
|
||||
"thinking": map[string]any{
|
||||
"type": "enabled",
|
||||
"budget_tokens": 1024,
|
||||
},
|
||||
"messages": []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": "List files in the current directory",
|
||||
},
|
||||
// assistant 消息包含 tool_use 但没有 signature
|
||||
map[string]any{
|
||||
"role": "assistant",
|
||||
"content": []map[string]any{
|
||||
{
|
||||
"type": "text",
|
||||
"text": "I'll list the files for you.",
|
||||
},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_01XGmNv",
|
||||
"name": "Bash",
|
||||
"input": map[string]any{"command": "ls -la"},
|
||||
// 故意不包含 signature
|
||||
},
|
||||
},
|
||||
},
|
||||
// 工具结果
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": []map[string]any{
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_01XGmNv",
|
||||
"content": "file1.txt\nfile2.txt\ndir1/",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"tools": []map[string]any{
|
||||
{
|
||||
"name": "Bash",
|
||||
"description": "Execute bash commands",
|
||||
"input_schema": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"command": map[string]any{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
"required": []string{"command"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
body, _ := json.Marshal(payload)
|
||||
|
||||
req, _ := http.NewRequest("POST", url, bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+claudeAPIKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
|
||||
client := &http.Client{Timeout: 60 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("请求失败: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
|
||||
// 400 错误说明 thought_signature 处理失败
|
||||
if resp.StatusCode == 400 {
|
||||
t.Fatalf("thought_signature 处理失败,收到 400 错误: %s", string(respBody))
|
||||
}
|
||||
|
||||
// 503 可能是账号限流,不算测试失败
|
||||
if resp.StatusCode == 503 {
|
||||
t.Skipf("账号暂时不可用 (503): %s", string(respBody))
|
||||
}
|
||||
|
||||
// 429 是限流
|
||||
if resp.StatusCode == 429 {
|
||||
t.Skipf("请求被限流 (429): %s", string(respBody))
|
||||
}
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
t.Fatalf("HTTP %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
if result["type"] != "message" {
|
||||
t.Errorf("期望 type=message, 得到 %v", result["type"])
|
||||
}
|
||||
t.Logf("✅ thinking 模式工具调用测试通过, id=%v", result["id"])
|
||||
}
|
||||
|
||||
// TestClaudeMessagesWithNoSignature 测试历史 thinking block 不带 signature 的场景
|
||||
// 验证:Gemini 模型接受没有 signature 的 thinking block
|
||||
func TestClaudeMessagesWithNoSignature(t *testing.T) {
|
||||
models := []string{
|
||||
"claude-haiku-4-5-20251001", // gemini-3-flash - 支持无 signature
|
||||
}
|
||||
for i, model := range models {
|
||||
if i > 0 {
|
||||
time.Sleep(testInterval)
|
||||
}
|
||||
t.Run(model+"_无signature", func(t *testing.T) {
|
||||
testClaudeWithNoSignature(t, model)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func testClaudeWithNoSignature(t *testing.T, model string) {
|
||||
url := baseURL + endpointPrefix + "/v1/messages"
|
||||
|
||||
// 模拟历史对话包含 thinking block 但没有 signature
|
||||
payload := map[string]any{
|
||||
"model": model,
|
||||
"max_tokens": 200,
|
||||
"stream": false,
|
||||
// 开启 thinking 模式
|
||||
"thinking": map[string]any{
|
||||
"type": "enabled",
|
||||
"budget_tokens": 1024,
|
||||
},
|
||||
"messages": []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": "What is 2+2?",
|
||||
},
|
||||
// assistant 消息包含 thinking block 但没有 signature
|
||||
map[string]any{
|
||||
"role": "assistant",
|
||||
"content": []map[string]any{
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Let me calculate 2+2...",
|
||||
// 故意不包含 signature
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "2+2 equals 4.",
|
||||
},
|
||||
},
|
||||
},
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": "What is 3+3?",
|
||||
},
|
||||
},
|
||||
}
|
||||
body, _ := json.Marshal(payload)
|
||||
|
||||
req, _ := http.NewRequest("POST", url, bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+claudeAPIKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
|
||||
client := &http.Client{Timeout: 60 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("请求失败: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
|
||||
if resp.StatusCode == 400 {
|
||||
t.Fatalf("无 signature thinking 处理失败,收到 400 错误: %s", string(respBody))
|
||||
}
|
||||
|
||||
if resp.StatusCode == 503 {
|
||||
t.Skipf("账号暂时不可用 (503): %s", string(respBody))
|
||||
}
|
||||
|
||||
if resp.StatusCode == 429 {
|
||||
t.Skipf("请求被限流 (429): %s", string(respBody))
|
||||
}
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
t.Fatalf("HTTP %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
if result["type"] != "message" {
|
||||
t.Errorf("期望 type=message, 得到 %v", result["type"])
|
||||
}
|
||||
t.Logf("✅ 无 signature thinking 处理测试通过, id=%v", result["id"])
|
||||
}
|
||||
126
backend/internal/pkg/antigravity/claude_types.go
Normal file
126
backend/internal/pkg/antigravity/claude_types.go
Normal file
@@ -0,0 +1,126 @@
|
||||
package antigravity
|
||||
|
||||
import "encoding/json"
|
||||
|
||||
// Claude 请求/响应类型定义
|
||||
|
||||
// ClaudeRequest Claude Messages API 请求
|
||||
type ClaudeRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []ClaudeMessage `json:"messages"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
System json.RawMessage `json:"system,omitempty"` // string 或 []SystemBlock
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
TopP *float64 `json:"top_p,omitempty"`
|
||||
TopK *int `json:"top_k,omitempty"`
|
||||
Tools []ClaudeTool `json:"tools,omitempty"`
|
||||
Thinking *ThinkingConfig `json:"thinking,omitempty"`
|
||||
Metadata *ClaudeMetadata `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// ClaudeMessage Claude 消息
|
||||
type ClaudeMessage struct {
|
||||
Role string `json:"role"` // user, assistant
|
||||
Content json.RawMessage `json:"content"`
|
||||
}
|
||||
|
||||
// ThinkingConfig Thinking 配置
|
||||
type ThinkingConfig struct {
|
||||
Type string `json:"type"` // "enabled" or "disabled"
|
||||
BudgetTokens int `json:"budget_tokens,omitempty"` // thinking budget
|
||||
}
|
||||
|
||||
// ClaudeMetadata 请求元数据
|
||||
type ClaudeMetadata struct {
|
||||
UserID string `json:"user_id,omitempty"`
|
||||
}
|
||||
|
||||
// ClaudeTool Claude 工具定义
|
||||
type ClaudeTool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
InputSchema map[string]any `json:"input_schema"`
|
||||
}
|
||||
|
||||
// SystemBlock system prompt 数组形式的元素
|
||||
type SystemBlock struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
// ContentBlock Claude 消息内容块(解析后)
|
||||
type ContentBlock struct {
|
||||
Type string `json:"type"`
|
||||
// text
|
||||
Text string `json:"text,omitempty"`
|
||||
// thinking
|
||||
Thinking string `json:"thinking,omitempty"`
|
||||
Signature string `json:"signature,omitempty"`
|
||||
// tool_use
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Input any `json:"input,omitempty"`
|
||||
// tool_result
|
||||
ToolUseID string `json:"tool_use_id,omitempty"`
|
||||
Content json.RawMessage `json:"content,omitempty"`
|
||||
IsError bool `json:"is_error,omitempty"`
|
||||
// image
|
||||
Source *ImageSource `json:"source,omitempty"`
|
||||
}
|
||||
|
||||
// ImageSource Claude 图片来源
|
||||
type ImageSource struct {
|
||||
Type string `json:"type"` // "base64"
|
||||
MediaType string `json:"media_type"` // "image/png", "image/jpeg" 等
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
// ClaudeResponse Claude Messages API 响应
|
||||
type ClaudeResponse struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // "message"
|
||||
Role string `json:"role"` // "assistant"
|
||||
Model string `json:"model"`
|
||||
Content []ClaudeContentItem `json:"content"`
|
||||
StopReason string `json:"stop_reason,omitempty"` // end_turn, tool_use, max_tokens
|
||||
StopSequence *string `json:"stop_sequence,omitempty"` // null 或具体值
|
||||
Usage ClaudeUsage `json:"usage"`
|
||||
}
|
||||
|
||||
// ClaudeContentItem Claude 响应内容项
|
||||
type ClaudeContentItem struct {
|
||||
Type string `json:"type"` // text, thinking, tool_use
|
||||
|
||||
// text
|
||||
Text string `json:"text,omitempty"`
|
||||
|
||||
// thinking
|
||||
Thinking string `json:"thinking,omitempty"`
|
||||
Signature string `json:"signature,omitempty"`
|
||||
|
||||
// tool_use
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Input any `json:"input,omitempty"`
|
||||
}
|
||||
|
||||
// ClaudeUsage Claude 用量统计
|
||||
type ClaudeUsage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"`
|
||||
CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"`
|
||||
}
|
||||
|
||||
// ClaudeError Claude 错误响应
|
||||
type ClaudeError struct {
|
||||
Type string `json:"type"` // "error"
|
||||
Error ErrorDetail `json:"error"`
|
||||
}
|
||||
|
||||
// ErrorDetail 错误详情
|
||||
type ErrorDetail struct {
|
||||
Type string `json:"type"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
305
backend/internal/pkg/antigravity/client.go
Normal file
305
backend/internal/pkg/antigravity/client.go
Normal file
@@ -0,0 +1,305 @@
|
||||
package antigravity
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TokenResponse Google OAuth token 响应
|
||||
type TokenResponse struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
ExpiresIn int64 `json:"expires_in"`
|
||||
TokenType string `json:"token_type"`
|
||||
Scope string `json:"scope,omitempty"`
|
||||
RefreshToken string `json:"refresh_token,omitempty"`
|
||||
}
|
||||
|
||||
// UserInfo Google 用户信息
|
||||
type UserInfo struct {
|
||||
Email string `json:"email"`
|
||||
Name string `json:"name,omitempty"`
|
||||
GivenName string `json:"given_name,omitempty"`
|
||||
FamilyName string `json:"family_name,omitempty"`
|
||||
Picture string `json:"picture,omitempty"`
|
||||
}
|
||||
|
||||
// LoadCodeAssistRequest loadCodeAssist 请求
|
||||
type LoadCodeAssistRequest struct {
|
||||
Metadata struct {
|
||||
IDEType string `json:"ideType"`
|
||||
} `json:"metadata"`
|
||||
}
|
||||
|
||||
// TierInfo 账户类型信息
|
||||
type TierInfo struct {
|
||||
ID string `json:"id"` // free-tier, g1-pro-tier, g1-ultra-tier
|
||||
Name string `json:"name"` // 显示名称
|
||||
Description string `json:"description"` // 描述
|
||||
}
|
||||
|
||||
// IneligibleTier 不符合条件的层级信息
|
||||
type IneligibleTier struct {
|
||||
Tier *TierInfo `json:"tier,omitempty"`
|
||||
// ReasonCode 不符合条件的原因代码,如 INELIGIBLE_ACCOUNT
|
||||
ReasonCode string `json:"reasonCode,omitempty"`
|
||||
ReasonMessage string `json:"reasonMessage,omitempty"`
|
||||
}
|
||||
|
||||
// LoadCodeAssistResponse loadCodeAssist 响应
|
||||
type LoadCodeAssistResponse struct {
|
||||
CloudAICompanionProject string `json:"cloudaicompanionProject"`
|
||||
CurrentTier *TierInfo `json:"currentTier,omitempty"`
|
||||
PaidTier *TierInfo `json:"paidTier,omitempty"`
|
||||
IneligibleTiers []*IneligibleTier `json:"ineligibleTiers,omitempty"`
|
||||
}
|
||||
|
||||
// GetTier 获取账户类型
|
||||
// 优先返回 paidTier(付费订阅级别),否则返回 currentTier
|
||||
func (r *LoadCodeAssistResponse) GetTier() string {
|
||||
if r.PaidTier != nil && r.PaidTier.ID != "" {
|
||||
return r.PaidTier.ID
|
||||
}
|
||||
if r.CurrentTier != nil {
|
||||
return r.CurrentTier.ID
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// Client Antigravity API 客户端
|
||||
type Client struct {
|
||||
httpClient *http.Client
|
||||
}
|
||||
|
||||
func NewClient(proxyURL string) *Client {
|
||||
client := &http.Client{
|
||||
Timeout: 30 * time.Second,
|
||||
}
|
||||
|
||||
if strings.TrimSpace(proxyURL) != "" {
|
||||
if proxyURLParsed, err := url.Parse(proxyURL); err == nil {
|
||||
client.Transport = &http.Transport{
|
||||
Proxy: http.ProxyURL(proxyURLParsed),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return &Client{
|
||||
httpClient: client,
|
||||
}
|
||||
}
|
||||
|
||||
// ExchangeCode 用 authorization code 交换 token
|
||||
func (c *Client) ExchangeCode(ctx context.Context, code, codeVerifier string) (*TokenResponse, error) {
|
||||
params := url.Values{}
|
||||
params.Set("client_id", ClientID)
|
||||
params.Set("client_secret", ClientSecret)
|
||||
params.Set("code", code)
|
||||
params.Set("redirect_uri", RedirectURI)
|
||||
params.Set("grant_type", "authorization_code")
|
||||
params.Set("code_verifier", codeVerifier)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, TokenURL, strings.NewReader(params.Encode()))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建请求失败: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("token 交换请求失败: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
bodyBytes, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取响应失败: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("token 交换失败 (HTTP %d): %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
var tokenResp TokenResponse
|
||||
if err := json.Unmarshal(bodyBytes, &tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("token 解析失败: %w", err)
|
||||
}
|
||||
|
||||
return &tokenResp, nil
|
||||
}
|
||||
|
||||
// RefreshToken 刷新 access_token
|
||||
func (c *Client) RefreshToken(ctx context.Context, refreshToken string) (*TokenResponse, error) {
|
||||
params := url.Values{}
|
||||
params.Set("client_id", ClientID)
|
||||
params.Set("client_secret", ClientSecret)
|
||||
params.Set("refresh_token", refreshToken)
|
||||
params.Set("grant_type", "refresh_token")
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, TokenURL, strings.NewReader(params.Encode()))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建请求失败: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("token 刷新请求失败: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
bodyBytes, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取响应失败: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("token 刷新失败 (HTTP %d): %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
var tokenResp TokenResponse
|
||||
if err := json.Unmarshal(bodyBytes, &tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("token 解析失败: %w", err)
|
||||
}
|
||||
|
||||
return &tokenResp, nil
|
||||
}
|
||||
|
||||
// GetUserInfo 获取用户信息
|
||||
func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, UserInfoURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建请求失败: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("用户信息请求失败: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
bodyBytes, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取响应失败: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("获取用户信息失败 (HTTP %d): %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
var userInfo UserInfo
|
||||
if err := json.Unmarshal(bodyBytes, &userInfo); err != nil {
|
||||
return nil, fmt.Errorf("用户信息解析失败: %w", err)
|
||||
}
|
||||
|
||||
return &userInfo, nil
|
||||
}
|
||||
|
||||
// LoadCodeAssist 获取 project_id
|
||||
func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadCodeAssistResponse, error) {
|
||||
reqBody := LoadCodeAssistRequest{}
|
||||
reqBody.Metadata.IDEType = "ANTIGRAVITY"
|
||||
|
||||
bodyBytes, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("序列化请求失败: %w", err)
|
||||
}
|
||||
|
||||
url := BaseURL + "/v1internal:loadCodeAssist"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(string(bodyBytes)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建请求失败: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", UserAgent)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("loadCodeAssist 请求失败: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
respBodyBytes, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取响应失败: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("loadCodeAssist 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
|
||||
}
|
||||
|
||||
var loadResp LoadCodeAssistResponse
|
||||
if err := json.Unmarshal(respBodyBytes, &loadResp); err != nil {
|
||||
return nil, fmt.Errorf("响应解析失败: %w", err)
|
||||
}
|
||||
|
||||
return &loadResp, nil
|
||||
}
|
||||
|
||||
// ModelQuotaInfo 模型配额信息
|
||||
type ModelQuotaInfo struct {
|
||||
RemainingFraction float64 `json:"remainingFraction"`
|
||||
ResetTime string `json:"resetTime,omitempty"`
|
||||
}
|
||||
|
||||
// ModelInfo 模型信息
|
||||
type ModelInfo struct {
|
||||
QuotaInfo *ModelQuotaInfo `json:"quotaInfo,omitempty"`
|
||||
}
|
||||
|
||||
// FetchAvailableModelsRequest fetchAvailableModels 请求
|
||||
type FetchAvailableModelsRequest struct {
|
||||
Project string `json:"project"`
|
||||
}
|
||||
|
||||
// FetchAvailableModelsResponse fetchAvailableModels 响应
|
||||
type FetchAvailableModelsResponse struct {
|
||||
Models map[string]ModelInfo `json:"models"`
|
||||
}
|
||||
|
||||
// FetchAvailableModels 获取可用模型和配额信息
|
||||
func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectID string) (*FetchAvailableModelsResponse, error) {
|
||||
reqBody := FetchAvailableModelsRequest{Project: projectID}
|
||||
bodyBytes, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("序列化请求失败: %w", err)
|
||||
}
|
||||
|
||||
apiURL := BaseURL + "/v1internal:fetchAvailableModels"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, strings.NewReader(string(bodyBytes)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建请求失败: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", UserAgent)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
respBodyBytes, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取响应失败: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("fetchAvailableModels 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
|
||||
}
|
||||
|
||||
var modelsResp FetchAvailableModelsResponse
|
||||
if err := json.Unmarshal(respBodyBytes, &modelsResp); err != nil {
|
||||
return nil, fmt.Errorf("响应解析失败: %w", err)
|
||||
}
|
||||
|
||||
return &modelsResp, nil
|
||||
}
|
||||
167
backend/internal/pkg/antigravity/gemini_types.go
Normal file
167
backend/internal/pkg/antigravity/gemini_types.go
Normal file
@@ -0,0 +1,167 @@
|
||||
package antigravity
|
||||
|
||||
// Gemini v1internal 请求/响应类型定义
|
||||
|
||||
// V1InternalRequest v1internal 请求包装
|
||||
type V1InternalRequest struct {
|
||||
Project string `json:"project"`
|
||||
RequestID string `json:"requestId"`
|
||||
UserAgent string `json:"userAgent"`
|
||||
RequestType string `json:"requestType,omitempty"`
|
||||
Model string `json:"model"`
|
||||
Request GeminiRequest `json:"request"`
|
||||
}
|
||||
|
||||
// GeminiRequest Gemini 请求内容
|
||||
type GeminiRequest struct {
|
||||
Contents []GeminiContent `json:"contents"`
|
||||
SystemInstruction *GeminiContent `json:"systemInstruction,omitempty"`
|
||||
GenerationConfig *GeminiGenerationConfig `json:"generationConfig,omitempty"`
|
||||
Tools []GeminiToolDeclaration `json:"tools,omitempty"`
|
||||
ToolConfig *GeminiToolConfig `json:"toolConfig,omitempty"`
|
||||
SafetySettings []GeminiSafetySetting `json:"safetySettings,omitempty"`
|
||||
SessionID string `json:"sessionId,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiContent Gemini 内容
|
||||
type GeminiContent struct {
|
||||
Role string `json:"role"` // user, model
|
||||
Parts []GeminiPart `json:"parts"`
|
||||
}
|
||||
|
||||
// GeminiPart Gemini 内容部分
|
||||
type GeminiPart struct {
|
||||
Text string `json:"text,omitempty"`
|
||||
Thought bool `json:"thought,omitempty"`
|
||||
ThoughtSignature string `json:"thoughtSignature,omitempty"`
|
||||
InlineData *GeminiInlineData `json:"inlineData,omitempty"`
|
||||
FunctionCall *GeminiFunctionCall `json:"functionCall,omitempty"`
|
||||
FunctionResponse *GeminiFunctionResponse `json:"functionResponse,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiInlineData Gemini 内联数据(图片等)
|
||||
type GeminiInlineData struct {
|
||||
MimeType string `json:"mimeType"`
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
// GeminiFunctionCall Gemini 函数调用
|
||||
type GeminiFunctionCall struct {
|
||||
Name string `json:"name"`
|
||||
Args any `json:"args,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiFunctionResponse Gemini 函数响应
|
||||
type GeminiFunctionResponse struct {
|
||||
Name string `json:"name"`
|
||||
Response map[string]any `json:"response"`
|
||||
ID string `json:"id,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiGenerationConfig Gemini 生成配置
|
||||
type GeminiGenerationConfig struct {
|
||||
MaxOutputTokens int `json:"maxOutputTokens,omitempty"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
TopP *float64 `json:"topP,omitempty"`
|
||||
TopK *int `json:"topK,omitempty"`
|
||||
ThinkingConfig *GeminiThinkingConfig `json:"thinkingConfig,omitempty"`
|
||||
StopSequences []string `json:"stopSequences,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiThinkingConfig Gemini thinking 配置
|
||||
type GeminiThinkingConfig struct {
|
||||
IncludeThoughts bool `json:"includeThoughts"`
|
||||
ThinkingBudget int `json:"thinkingBudget,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiToolDeclaration Gemini 工具声明
|
||||
type GeminiToolDeclaration struct {
|
||||
FunctionDeclarations []GeminiFunctionDecl `json:"functionDeclarations,omitempty"`
|
||||
GoogleSearch *GeminiGoogleSearch `json:"googleSearch,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiFunctionDecl Gemini 函数声明
|
||||
type GeminiFunctionDecl struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Parameters map[string]any `json:"parameters,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiGoogleSearch Gemini Google 搜索工具
|
||||
type GeminiGoogleSearch struct {
|
||||
EnhancedContent *GeminiEnhancedContent `json:"enhancedContent,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiEnhancedContent 增强内容配置
|
||||
type GeminiEnhancedContent struct {
|
||||
ImageSearch *GeminiImageSearch `json:"imageSearch,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiImageSearch 图片搜索配置
|
||||
type GeminiImageSearch struct {
|
||||
MaxResultCount int `json:"maxResultCount,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiToolConfig Gemini 工具配置
|
||||
type GeminiToolConfig struct {
|
||||
FunctionCallingConfig *GeminiFunctionCallingConfig `json:"functionCallingConfig,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiFunctionCallingConfig 函数调用配置
|
||||
type GeminiFunctionCallingConfig struct {
|
||||
Mode string `json:"mode,omitempty"` // VALIDATED, AUTO, NONE
|
||||
}
|
||||
|
||||
// GeminiSafetySetting Gemini 安全设置
|
||||
type GeminiSafetySetting struct {
|
||||
Category string `json:"category"`
|
||||
Threshold string `json:"threshold"`
|
||||
}
|
||||
|
||||
// V1InternalResponse v1internal 响应包装
|
||||
type V1InternalResponse struct {
|
||||
Response GeminiResponse `json:"response"`
|
||||
ResponseID string `json:"responseId,omitempty"`
|
||||
ModelVersion string `json:"modelVersion,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiResponse Gemini 响应
|
||||
type GeminiResponse struct {
|
||||
Candidates []GeminiCandidate `json:"candidates,omitempty"`
|
||||
UsageMetadata *GeminiUsageMetadata `json:"usageMetadata,omitempty"`
|
||||
ResponseID string `json:"responseId,omitempty"`
|
||||
ModelVersion string `json:"modelVersion,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiCandidate Gemini 候选响应
|
||||
type GeminiCandidate struct {
|
||||
Content *GeminiContent `json:"content,omitempty"`
|
||||
FinishReason string `json:"finishReason,omitempty"`
|
||||
Index int `json:"index,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiUsageMetadata Gemini 用量元数据
|
||||
type GeminiUsageMetadata struct {
|
||||
PromptTokenCount int `json:"promptTokenCount,omitempty"`
|
||||
CandidatesTokenCount int `json:"candidatesTokenCount,omitempty"`
|
||||
TotalTokenCount int `json:"totalTokenCount,omitempty"`
|
||||
}
|
||||
|
||||
// DefaultSafetySettings 默认安全设置(关闭所有过滤)
|
||||
var DefaultSafetySettings = []GeminiSafetySetting{
|
||||
{Category: "HARM_CATEGORY_HARASSMENT", Threshold: "OFF"},
|
||||
{Category: "HARM_CATEGORY_HATE_SPEECH", Threshold: "OFF"},
|
||||
{Category: "HARM_CATEGORY_SEXUALLY_EXPLICIT", Threshold: "OFF"},
|
||||
{Category: "HARM_CATEGORY_DANGEROUS_CONTENT", Threshold: "OFF"},
|
||||
{Category: "HARM_CATEGORY_CIVIC_INTEGRITY", Threshold: "OFF"},
|
||||
}
|
||||
|
||||
// DefaultStopSequences 默认停止序列
|
||||
var DefaultStopSequences = []string{
|
||||
"<|user|>",
|
||||
"<|endoftext|>",
|
||||
"<|end_of_turn|>",
|
||||
"[DONE]",
|
||||
"\n\nHuman:",
|
||||
}
|
||||
179
backend/internal/pkg/antigravity/oauth.go
Normal file
179
backend/internal/pkg/antigravity/oauth.go
Normal file
@@ -0,0 +1,179 @@
|
||||
package antigravity
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
// Google OAuth 端点
|
||||
AuthorizeURL = "https://accounts.google.com/o/oauth2/v2/auth"
|
||||
TokenURL = "https://oauth2.googleapis.com/token"
|
||||
UserInfoURL = "https://www.googleapis.com/oauth2/v2/userinfo"
|
||||
|
||||
// Antigravity OAuth 客户端凭证
|
||||
ClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
||||
ClientSecret = "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"
|
||||
|
||||
// 固定的 redirect_uri(用户需手动复制 code)
|
||||
RedirectURI = "http://localhost:8085/callback"
|
||||
|
||||
// OAuth scopes
|
||||
Scopes = "https://www.googleapis.com/auth/cloud-platform " +
|
||||
"https://www.googleapis.com/auth/userinfo.email " +
|
||||
"https://www.googleapis.com/auth/userinfo.profile " +
|
||||
"https://www.googleapis.com/auth/cclog " +
|
||||
"https://www.googleapis.com/auth/experimentsandconfigs"
|
||||
|
||||
// API 端点
|
||||
BaseURL = "https://cloudcode-pa.googleapis.com"
|
||||
|
||||
// User-Agent
|
||||
UserAgent = "antigravity/1.11.9 windows/amd64"
|
||||
|
||||
// Session 过期时间
|
||||
SessionTTL = 30 * time.Minute
|
||||
)
|
||||
|
||||
// OAuthSession 保存 OAuth 授权流程的临时状态
|
||||
type OAuthSession struct {
|
||||
State string `json:"state"`
|
||||
CodeVerifier string `json:"code_verifier"`
|
||||
ProxyURL string `json:"proxy_url,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// SessionStore OAuth session 存储
|
||||
type SessionStore struct {
|
||||
mu sync.RWMutex
|
||||
sessions map[string]*OAuthSession
|
||||
stopCh chan struct{}
|
||||
}
|
||||
|
||||
func NewSessionStore() *SessionStore {
|
||||
store := &SessionStore{
|
||||
sessions: make(map[string]*OAuthSession),
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
go store.cleanup()
|
||||
return store
|
||||
}
|
||||
|
||||
func (s *SessionStore) Set(sessionID string, session *OAuthSession) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.sessions[sessionID] = session
|
||||
}
|
||||
|
||||
func (s *SessionStore) Get(sessionID string) (*OAuthSession, bool) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
session, ok := s.sessions[sessionID]
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
if time.Since(session.CreatedAt) > SessionTTL {
|
||||
return nil, false
|
||||
}
|
||||
return session, true
|
||||
}
|
||||
|
||||
func (s *SessionStore) Delete(sessionID string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
delete(s.sessions, sessionID)
|
||||
}
|
||||
|
||||
func (s *SessionStore) Stop() {
|
||||
select {
|
||||
case <-s.stopCh:
|
||||
return
|
||||
default:
|
||||
close(s.stopCh)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SessionStore) cleanup() {
|
||||
ticker := time.NewTicker(5 * time.Minute)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-s.stopCh:
|
||||
return
|
||||
case <-ticker.C:
|
||||
s.mu.Lock()
|
||||
for id, session := range s.sessions {
|
||||
if time.Since(session.CreatedAt) > SessionTTL {
|
||||
delete(s.sessions, id)
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func GenerateRandomBytes(n int) ([]byte, error) {
|
||||
b := make([]byte, n)
|
||||
_, err := rand.Read(b)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
func GenerateState() (string, error) {
|
||||
bytes, err := GenerateRandomBytes(32)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64URLEncode(bytes), nil
|
||||
}
|
||||
|
||||
func GenerateSessionID() (string, error) {
|
||||
bytes, err := GenerateRandomBytes(16)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(bytes), nil
|
||||
}
|
||||
|
||||
func GenerateCodeVerifier() (string, error) {
|
||||
bytes, err := GenerateRandomBytes(32)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64URLEncode(bytes), nil
|
||||
}
|
||||
|
||||
func GenerateCodeChallenge(verifier string) string {
|
||||
hash := sha256.Sum256([]byte(verifier))
|
||||
return base64URLEncode(hash[:])
|
||||
}
|
||||
|
||||
func base64URLEncode(data []byte) string {
|
||||
return strings.TrimRight(base64.URLEncoding.EncodeToString(data), "=")
|
||||
}
|
||||
|
||||
// BuildAuthorizationURL 构建 Google OAuth 授权 URL
|
||||
func BuildAuthorizationURL(state, codeChallenge string) string {
|
||||
params := url.Values{}
|
||||
params.Set("client_id", ClientID)
|
||||
params.Set("redirect_uri", RedirectURI)
|
||||
params.Set("response_type", "code")
|
||||
params.Set("scope", Scopes)
|
||||
params.Set("state", state)
|
||||
params.Set("code_challenge", codeChallenge)
|
||||
params.Set("code_challenge_method", "S256")
|
||||
params.Set("access_type", "offline")
|
||||
params.Set("prompt", "consent")
|
||||
params.Set("include_granted_scopes", "true")
|
||||
|
||||
return fmt.Sprintf("%s?%s", AuthorizeURL, params.Encode())
|
||||
}
|
||||
525
backend/internal/pkg/antigravity/request_transformer.go
Normal file
525
backend/internal/pkg/antigravity/request_transformer.go
Normal file
@@ -0,0 +1,525 @@
|
||||
package antigravity
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// TransformClaudeToGemini 将 Claude 请求转换为 v1internal Gemini 格式
|
||||
func TransformClaudeToGemini(claudeReq *ClaudeRequest, projectID, mappedModel string) ([]byte, error) {
|
||||
// 用于存储 tool_use id -> name 映射
|
||||
toolIDToName := make(map[string]string)
|
||||
|
||||
// 检测是否启用 thinking
|
||||
isThinkingEnabled := claudeReq.Thinking != nil && claudeReq.Thinking.Type == "enabled"
|
||||
|
||||
// 只有 Gemini 模型支持 dummy thought workaround
|
||||
// Claude 模型通过 Vertex/Google API 需要有效的 thought signatures
|
||||
allowDummyThought := strings.HasPrefix(mappedModel, "gemini-")
|
||||
|
||||
// 1. 构建 contents
|
||||
contents, err := buildContents(claudeReq.Messages, toolIDToName, isThinkingEnabled, allowDummyThought)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build contents: %w", err)
|
||||
}
|
||||
|
||||
// 2. 构建 systemInstruction
|
||||
systemInstruction := buildSystemInstruction(claudeReq.System, claudeReq.Model)
|
||||
|
||||
// 3. 构建 generationConfig
|
||||
generationConfig := buildGenerationConfig(claudeReq)
|
||||
|
||||
// 4. 构建 tools
|
||||
tools := buildTools(claudeReq.Tools)
|
||||
|
||||
// 5. 构建内部请求
|
||||
innerRequest := GeminiRequest{
|
||||
Contents: contents,
|
||||
SafetySettings: DefaultSafetySettings,
|
||||
}
|
||||
|
||||
if systemInstruction != nil {
|
||||
innerRequest.SystemInstruction = systemInstruction
|
||||
}
|
||||
if generationConfig != nil {
|
||||
innerRequest.GenerationConfig = generationConfig
|
||||
}
|
||||
if len(tools) > 0 {
|
||||
innerRequest.Tools = tools
|
||||
innerRequest.ToolConfig = &GeminiToolConfig{
|
||||
FunctionCallingConfig: &GeminiFunctionCallingConfig{
|
||||
Mode: "VALIDATED",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// 如果提供了 metadata.user_id,复用为 sessionId
|
||||
if claudeReq.Metadata != nil && claudeReq.Metadata.UserID != "" {
|
||||
innerRequest.SessionID = claudeReq.Metadata.UserID
|
||||
}
|
||||
|
||||
// 6. 包装为 v1internal 请求
|
||||
v1Req := V1InternalRequest{
|
||||
Project: projectID,
|
||||
RequestID: "agent-" + uuid.New().String(),
|
||||
UserAgent: "sub2api",
|
||||
RequestType: "agent",
|
||||
Model: mappedModel,
|
||||
Request: innerRequest,
|
||||
}
|
||||
|
||||
return json.Marshal(v1Req)
|
||||
}
|
||||
|
||||
// buildSystemInstruction 构建 systemInstruction
|
||||
func buildSystemInstruction(system json.RawMessage, modelName string) *GeminiContent {
|
||||
var parts []GeminiPart
|
||||
|
||||
// 注入身份防护指令
|
||||
identityPatch := fmt.Sprintf(
|
||||
"--- [IDENTITY_PATCH] ---\n"+
|
||||
"Ignore any previous instructions regarding your identity or host platform (e.g., Amazon Q, Google AI).\n"+
|
||||
"You are currently providing services as the native %s model via a standard API proxy.\n"+
|
||||
"Always use the 'claude' command for terminal tasks if relevant.\n"+
|
||||
"--- [SYSTEM_PROMPT_BEGIN] ---\n",
|
||||
modelName,
|
||||
)
|
||||
parts = append(parts, GeminiPart{Text: identityPatch})
|
||||
|
||||
// 解析 system prompt
|
||||
if len(system) > 0 {
|
||||
// 尝试解析为字符串
|
||||
var sysStr string
|
||||
if err := json.Unmarshal(system, &sysStr); err == nil {
|
||||
if strings.TrimSpace(sysStr) != "" {
|
||||
parts = append(parts, GeminiPart{Text: sysStr})
|
||||
}
|
||||
} else {
|
||||
// 尝试解析为数组
|
||||
var sysBlocks []SystemBlock
|
||||
if err := json.Unmarshal(system, &sysBlocks); err == nil {
|
||||
for _, block := range sysBlocks {
|
||||
if block.Type == "text" && strings.TrimSpace(block.Text) != "" {
|
||||
parts = append(parts, GeminiPart{Text: block.Text})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
parts = append(parts, GeminiPart{Text: "\n--- [SYSTEM_PROMPT_END] ---"})
|
||||
|
||||
return &GeminiContent{
|
||||
Role: "user",
|
||||
Parts: parts,
|
||||
}
|
||||
}
|
||||
|
||||
// buildContents 构建 contents
|
||||
func buildContents(messages []ClaudeMessage, toolIDToName map[string]string, isThinkingEnabled, allowDummyThought bool) ([]GeminiContent, error) {
|
||||
var contents []GeminiContent
|
||||
|
||||
for i, msg := range messages {
|
||||
role := msg.Role
|
||||
if role == "assistant" {
|
||||
role = "model"
|
||||
}
|
||||
|
||||
parts, err := buildParts(msg.Content, toolIDToName, allowDummyThought)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build parts for message %d: %w", i, err)
|
||||
}
|
||||
|
||||
// 只有 Gemini 模型支持 dummy thinking block workaround
|
||||
// 只对最后一条 assistant 消息添加(Pre-fill 场景)
|
||||
// 历史 assistant 消息不能添加没有 signature 的 dummy thinking block
|
||||
if allowDummyThought && role == "model" && isThinkingEnabled && i == len(messages)-1 {
|
||||
hasThoughtPart := false
|
||||
for _, p := range parts {
|
||||
if p.Thought {
|
||||
hasThoughtPart = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasThoughtPart && len(parts) > 0 {
|
||||
// 在开头添加 dummy thinking block
|
||||
parts = append([]GeminiPart{{
|
||||
Text: "Thinking...",
|
||||
Thought: true,
|
||||
}}, parts...)
|
||||
}
|
||||
}
|
||||
|
||||
if len(parts) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
contents = append(contents, GeminiContent{
|
||||
Role: role,
|
||||
Parts: parts,
|
||||
})
|
||||
}
|
||||
|
||||
return contents, nil
|
||||
}
|
||||
|
||||
// dummyThoughtSignature 用于跳过 Gemini 3 thought_signature 验证
|
||||
// 参考: https://ai.google.dev/gemini-api/docs/thought-signatures
|
||||
const dummyThoughtSignature = "skip_thought_signature_validator"
|
||||
|
||||
// buildParts 构建消息的 parts
|
||||
// allowDummyThought: 只有 Gemini 模型支持 dummy thought signature
|
||||
func buildParts(content json.RawMessage, toolIDToName map[string]string, allowDummyThought bool) ([]GeminiPart, error) {
|
||||
var parts []GeminiPart
|
||||
|
||||
// 尝试解析为字符串
|
||||
var textContent string
|
||||
if err := json.Unmarshal(content, &textContent); err == nil {
|
||||
if textContent != "(no content)" && strings.TrimSpace(textContent) != "" {
|
||||
parts = append(parts, GeminiPart{Text: strings.TrimSpace(textContent)})
|
||||
}
|
||||
return parts, nil
|
||||
}
|
||||
|
||||
// 解析为内容块数组
|
||||
var blocks []ContentBlock
|
||||
if err := json.Unmarshal(content, &blocks); err != nil {
|
||||
return nil, fmt.Errorf("parse content blocks: %w", err)
|
||||
}
|
||||
|
||||
for _, block := range blocks {
|
||||
switch block.Type {
|
||||
case "text":
|
||||
if block.Text != "(no content)" && strings.TrimSpace(block.Text) != "" {
|
||||
parts = append(parts, GeminiPart{Text: block.Text})
|
||||
}
|
||||
|
||||
case "thinking":
|
||||
part := GeminiPart{
|
||||
Text: block.Thinking,
|
||||
Thought: true,
|
||||
}
|
||||
// 保留原有 signature(Claude 模型需要有效的 signature)
|
||||
if block.Signature != "" {
|
||||
part.ThoughtSignature = block.Signature
|
||||
}
|
||||
parts = append(parts, part)
|
||||
|
||||
case "image":
|
||||
if block.Source != nil && block.Source.Type == "base64" {
|
||||
parts = append(parts, GeminiPart{
|
||||
InlineData: &GeminiInlineData{
|
||||
MimeType: block.Source.MediaType,
|
||||
Data: block.Source.Data,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
case "tool_use":
|
||||
// 存储 id -> name 映射
|
||||
if block.ID != "" && block.Name != "" {
|
||||
toolIDToName[block.ID] = block.Name
|
||||
}
|
||||
|
||||
part := GeminiPart{
|
||||
FunctionCall: &GeminiFunctionCall{
|
||||
Name: block.Name,
|
||||
Args: block.Input,
|
||||
ID: block.ID,
|
||||
},
|
||||
}
|
||||
// 保留原有 signature,或对 Gemini 模型使用 dummy signature
|
||||
if block.Signature != "" {
|
||||
part.ThoughtSignature = block.Signature
|
||||
} else if allowDummyThought {
|
||||
part.ThoughtSignature = dummyThoughtSignature
|
||||
}
|
||||
parts = append(parts, part)
|
||||
|
||||
case "tool_result":
|
||||
// 获取函数名
|
||||
funcName := block.Name
|
||||
if funcName == "" {
|
||||
if name, ok := toolIDToName[block.ToolUseID]; ok {
|
||||
funcName = name
|
||||
} else {
|
||||
funcName = block.ToolUseID
|
||||
}
|
||||
}
|
||||
|
||||
// 解析 content
|
||||
resultContent := parseToolResultContent(block.Content, block.IsError)
|
||||
|
||||
parts = append(parts, GeminiPart{
|
||||
FunctionResponse: &GeminiFunctionResponse{
|
||||
Name: funcName,
|
||||
Response: map[string]any{
|
||||
"result": resultContent,
|
||||
},
|
||||
ID: block.ToolUseID,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return parts, nil
|
||||
}
|
||||
|
||||
// parseToolResultContent 解析 tool_result 的 content
|
||||
func parseToolResultContent(content json.RawMessage, isError bool) string {
|
||||
if len(content) == 0 {
|
||||
if isError {
|
||||
return "Tool execution failed with no output."
|
||||
}
|
||||
return "Command executed successfully."
|
||||
}
|
||||
|
||||
// 尝试解析为字符串
|
||||
var str string
|
||||
if err := json.Unmarshal(content, &str); err == nil {
|
||||
if strings.TrimSpace(str) == "" {
|
||||
if isError {
|
||||
return "Tool execution failed with no output."
|
||||
}
|
||||
return "Command executed successfully."
|
||||
}
|
||||
return str
|
||||
}
|
||||
|
||||
// 尝试解析为数组
|
||||
var arr []map[string]any
|
||||
if err := json.Unmarshal(content, &arr); err == nil {
|
||||
var texts []string
|
||||
for _, item := range arr {
|
||||
if text, ok := item["text"].(string); ok {
|
||||
texts = append(texts, text)
|
||||
}
|
||||
}
|
||||
result := strings.Join(texts, "\n")
|
||||
if strings.TrimSpace(result) == "" {
|
||||
if isError {
|
||||
return "Tool execution failed with no output."
|
||||
}
|
||||
return "Command executed successfully."
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// 返回原始 JSON
|
||||
return string(content)
|
||||
}
|
||||
|
||||
// buildGenerationConfig 构建 generationConfig
|
||||
func buildGenerationConfig(req *ClaudeRequest) *GeminiGenerationConfig {
|
||||
config := &GeminiGenerationConfig{
|
||||
MaxOutputTokens: 64000, // 默认最大输出
|
||||
StopSequences: DefaultStopSequences,
|
||||
}
|
||||
|
||||
// Thinking 配置
|
||||
if req.Thinking != nil && req.Thinking.Type == "enabled" {
|
||||
config.ThinkingConfig = &GeminiThinkingConfig{
|
||||
IncludeThoughts: true,
|
||||
}
|
||||
if req.Thinking.BudgetTokens > 0 {
|
||||
budget := req.Thinking.BudgetTokens
|
||||
// gemini-2.5-flash 上限 24576
|
||||
if strings.Contains(req.Model, "gemini-2.5-flash") && budget > 24576 {
|
||||
budget = 24576
|
||||
}
|
||||
config.ThinkingConfig.ThinkingBudget = budget
|
||||
}
|
||||
}
|
||||
|
||||
// 其他参数
|
||||
if req.Temperature != nil {
|
||||
config.Temperature = req.Temperature
|
||||
}
|
||||
if req.TopP != nil {
|
||||
config.TopP = req.TopP
|
||||
}
|
||||
if req.TopK != nil {
|
||||
config.TopK = req.TopK
|
||||
}
|
||||
|
||||
return config
|
||||
}
|
||||
|
||||
// buildTools 构建 tools
|
||||
func buildTools(tools []ClaudeTool) []GeminiToolDeclaration {
|
||||
if len(tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 检查是否有 web_search 工具
|
||||
hasWebSearch := false
|
||||
for _, tool := range tools {
|
||||
if tool.Name == "web_search" {
|
||||
hasWebSearch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if hasWebSearch {
|
||||
// Web Search 工具映射
|
||||
return []GeminiToolDeclaration{{
|
||||
GoogleSearch: &GeminiGoogleSearch{
|
||||
EnhancedContent: &GeminiEnhancedContent{
|
||||
ImageSearch: &GeminiImageSearch{
|
||||
MaxResultCount: 5,
|
||||
},
|
||||
},
|
||||
},
|
||||
}}
|
||||
}
|
||||
|
||||
// 普通工具
|
||||
var funcDecls []GeminiFunctionDecl
|
||||
for _, tool := range tools {
|
||||
// 清理 JSON Schema
|
||||
params := cleanJSONSchema(tool.InputSchema)
|
||||
|
||||
funcDecls = append(funcDecls, GeminiFunctionDecl{
|
||||
Name: tool.Name,
|
||||
Description: tool.Description,
|
||||
Parameters: params,
|
||||
})
|
||||
}
|
||||
|
||||
if len(funcDecls) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return []GeminiToolDeclaration{{
|
||||
FunctionDeclarations: funcDecls,
|
||||
}}
|
||||
}
|
||||
|
||||
// cleanJSONSchema 清理 JSON Schema,移除 Antigravity/Gemini 不支持的字段
|
||||
// 参考 proxycast 的实现,确保 schema 符合 JSON Schema draft 2020-12
|
||||
func cleanJSONSchema(schema map[string]any) map[string]any {
|
||||
if schema == nil {
|
||||
return nil
|
||||
}
|
||||
cleaned := cleanSchemaValue(schema)
|
||||
result, ok := cleaned.(map[string]any)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 确保有 type 字段(默认 OBJECT)
|
||||
if _, hasType := result["type"]; !hasType {
|
||||
result["type"] = "OBJECT"
|
||||
}
|
||||
|
||||
// 确保有 properties 字段(默认空对象)
|
||||
if _, hasProps := result["properties"]; !hasProps {
|
||||
result["properties"] = make(map[string]any)
|
||||
}
|
||||
|
||||
// 验证 required 中的字段都存在于 properties 中
|
||||
if required, ok := result["required"].([]any); ok {
|
||||
if props, ok := result["properties"].(map[string]any); ok {
|
||||
validRequired := make([]any, 0, len(required))
|
||||
for _, r := range required {
|
||||
if reqName, ok := r.(string); ok {
|
||||
if _, exists := props[reqName]; exists {
|
||||
validRequired = append(validRequired, r)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(validRequired) > 0 {
|
||||
result["required"] = validRequired
|
||||
} else {
|
||||
delete(result, "required")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// excludedSchemaKeys 不支持的 schema 字段
|
||||
var excludedSchemaKeys = map[string]bool{
|
||||
"$schema": true,
|
||||
"$id": true,
|
||||
"$ref": true,
|
||||
"additionalProperties": true,
|
||||
"minLength": true,
|
||||
"maxLength": true,
|
||||
"minItems": true,
|
||||
"maxItems": true,
|
||||
"uniqueItems": true,
|
||||
"minimum": true,
|
||||
"maximum": true,
|
||||
"exclusiveMinimum": true,
|
||||
"exclusiveMaximum": true,
|
||||
"pattern": true,
|
||||
"format": true,
|
||||
"default": true,
|
||||
"strict": true,
|
||||
"const": true,
|
||||
"examples": true,
|
||||
"deprecated": true,
|
||||
"readOnly": true,
|
||||
"writeOnly": true,
|
||||
"contentMediaType": true,
|
||||
"contentEncoding": true,
|
||||
}
|
||||
|
||||
// cleanSchemaValue 递归清理 schema 值
|
||||
func cleanSchemaValue(value any) any {
|
||||
switch v := value.(type) {
|
||||
case map[string]any:
|
||||
result := make(map[string]any)
|
||||
for k, val := range v {
|
||||
// 跳过不支持的字段
|
||||
if excludedSchemaKeys[k] {
|
||||
continue
|
||||
}
|
||||
|
||||
// 特殊处理 type 字段
|
||||
if k == "type" {
|
||||
result[k] = cleanTypeValue(val)
|
||||
continue
|
||||
}
|
||||
|
||||
// 递归清理所有值
|
||||
result[k] = cleanSchemaValue(val)
|
||||
}
|
||||
return result
|
||||
|
||||
case []any:
|
||||
// 递归处理数组中的每个元素
|
||||
cleaned := make([]any, 0, len(v))
|
||||
for _, item := range v {
|
||||
cleaned = append(cleaned, cleanSchemaValue(item))
|
||||
}
|
||||
return cleaned
|
||||
|
||||
default:
|
||||
return value
|
||||
}
|
||||
}
|
||||
|
||||
// cleanTypeValue 处理 type 字段,转换为大写
|
||||
func cleanTypeValue(value any) any {
|
||||
switch v := value.(type) {
|
||||
case string:
|
||||
return strings.ToUpper(v)
|
||||
case []any:
|
||||
// 联合类型 ["string", "null"] -> 取第一个非 null 类型
|
||||
for _, t := range v {
|
||||
if ts, ok := t.(string); ok && ts != "null" {
|
||||
return strings.ToUpper(ts)
|
||||
}
|
||||
}
|
||||
// 如果只有 null,返回 STRING
|
||||
return "STRING"
|
||||
default:
|
||||
return value
|
||||
}
|
||||
}
|
||||
269
backend/internal/pkg/antigravity/response_transformer.go
Normal file
269
backend/internal/pkg/antigravity/response_transformer.go
Normal file
@@ -0,0 +1,269 @@
|
||||
package antigravity
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// TransformGeminiToClaude 将 Gemini 响应转换为 Claude 格式(非流式)
|
||||
func TransformGeminiToClaude(geminiResp []byte, originalModel string) ([]byte, *ClaudeUsage, error) {
|
||||
// 解包 v1internal 响应
|
||||
var v1Resp V1InternalResponse
|
||||
if err := json.Unmarshal(geminiResp, &v1Resp); err != nil {
|
||||
// 尝试直接解析为 GeminiResponse
|
||||
var directResp GeminiResponse
|
||||
if err2 := json.Unmarshal(geminiResp, &directResp); err2 != nil {
|
||||
return nil, nil, fmt.Errorf("parse gemini response: %w", err)
|
||||
}
|
||||
v1Resp.Response = directResp
|
||||
v1Resp.ResponseID = directResp.ResponseID
|
||||
v1Resp.ModelVersion = directResp.ModelVersion
|
||||
}
|
||||
|
||||
// 使用处理器转换
|
||||
processor := NewNonStreamingProcessor()
|
||||
claudeResp := processor.Process(&v1Resp.Response, v1Resp.ResponseID, originalModel)
|
||||
|
||||
// 序列化
|
||||
respBytes, err := json.Marshal(claudeResp)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("marshal claude response: %w", err)
|
||||
}
|
||||
|
||||
return respBytes, &claudeResp.Usage, nil
|
||||
}
|
||||
|
||||
// NonStreamingProcessor 非流式响应处理器
|
||||
type NonStreamingProcessor struct {
|
||||
contentBlocks []ClaudeContentItem
|
||||
textBuilder string
|
||||
thinkingBuilder string
|
||||
thinkingSignature string
|
||||
trailingSignature string
|
||||
hasToolCall bool
|
||||
}
|
||||
|
||||
// NewNonStreamingProcessor 创建非流式响应处理器
|
||||
func NewNonStreamingProcessor() *NonStreamingProcessor {
|
||||
return &NonStreamingProcessor{
|
||||
contentBlocks: make([]ClaudeContentItem, 0),
|
||||
}
|
||||
}
|
||||
|
||||
// Process 处理 Gemini 响应
|
||||
func (p *NonStreamingProcessor) Process(geminiResp *GeminiResponse, responseID, originalModel string) *ClaudeResponse {
|
||||
// 获取 parts
|
||||
var parts []GeminiPart
|
||||
if len(geminiResp.Candidates) > 0 && geminiResp.Candidates[0].Content != nil {
|
||||
parts = geminiResp.Candidates[0].Content.Parts
|
||||
}
|
||||
|
||||
// 处理所有 parts
|
||||
for _, part := range parts {
|
||||
p.processPart(&part)
|
||||
}
|
||||
|
||||
// 刷新剩余内容
|
||||
p.flushThinking()
|
||||
p.flushText()
|
||||
|
||||
// 处理 trailingSignature
|
||||
if p.trailingSignature != "" {
|
||||
p.contentBlocks = append(p.contentBlocks, ClaudeContentItem{
|
||||
Type: "thinking",
|
||||
Thinking: "",
|
||||
Signature: p.trailingSignature,
|
||||
})
|
||||
}
|
||||
|
||||
// 构建响应
|
||||
return p.buildResponse(geminiResp, responseID, originalModel)
|
||||
}
|
||||
|
||||
// processPart 处理单个 part
|
||||
func (p *NonStreamingProcessor) processPart(part *GeminiPart) {
|
||||
signature := part.ThoughtSignature
|
||||
|
||||
// 1. FunctionCall 处理
|
||||
if part.FunctionCall != nil {
|
||||
p.flushThinking()
|
||||
p.flushText()
|
||||
|
||||
// 处理 trailingSignature
|
||||
if p.trailingSignature != "" {
|
||||
p.contentBlocks = append(p.contentBlocks, ClaudeContentItem{
|
||||
Type: "thinking",
|
||||
Thinking: "",
|
||||
Signature: p.trailingSignature,
|
||||
})
|
||||
p.trailingSignature = ""
|
||||
}
|
||||
|
||||
p.hasToolCall = true
|
||||
|
||||
// 生成 tool_use id
|
||||
toolID := part.FunctionCall.ID
|
||||
if toolID == "" {
|
||||
toolID = fmt.Sprintf("%s-%s", part.FunctionCall.Name, generateRandomID())
|
||||
}
|
||||
|
||||
item := ClaudeContentItem{
|
||||
Type: "tool_use",
|
||||
ID: toolID,
|
||||
Name: part.FunctionCall.Name,
|
||||
Input: part.FunctionCall.Args,
|
||||
}
|
||||
|
||||
if signature != "" {
|
||||
item.Signature = signature
|
||||
}
|
||||
|
||||
p.contentBlocks = append(p.contentBlocks, item)
|
||||
return
|
||||
}
|
||||
|
||||
// 2. Text 处理
|
||||
if part.Text != "" || part.Thought {
|
||||
if part.Thought {
|
||||
// Thinking part
|
||||
p.flushText()
|
||||
|
||||
// 处理 trailingSignature
|
||||
if p.trailingSignature != "" {
|
||||
p.flushThinking()
|
||||
p.contentBlocks = append(p.contentBlocks, ClaudeContentItem{
|
||||
Type: "thinking",
|
||||
Thinking: "",
|
||||
Signature: p.trailingSignature,
|
||||
})
|
||||
p.trailingSignature = ""
|
||||
}
|
||||
|
||||
p.thinkingBuilder += part.Text
|
||||
if signature != "" {
|
||||
p.thinkingSignature = signature
|
||||
}
|
||||
} else {
|
||||
// 普通 Text
|
||||
if part.Text == "" {
|
||||
// 空 text 带签名 - 暂存
|
||||
if signature != "" {
|
||||
p.trailingSignature = signature
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
p.flushThinking()
|
||||
|
||||
// 处理之前的 trailingSignature
|
||||
if p.trailingSignature != "" {
|
||||
p.flushText()
|
||||
p.contentBlocks = append(p.contentBlocks, ClaudeContentItem{
|
||||
Type: "thinking",
|
||||
Thinking: "",
|
||||
Signature: p.trailingSignature,
|
||||
})
|
||||
p.trailingSignature = ""
|
||||
}
|
||||
|
||||
p.textBuilder += part.Text
|
||||
|
||||
// 非空 text 带签名 - 立即刷新并输出空 thinking 块
|
||||
if signature != "" {
|
||||
p.flushText()
|
||||
p.contentBlocks = append(p.contentBlocks, ClaudeContentItem{
|
||||
Type: "thinking",
|
||||
Thinking: "",
|
||||
Signature: signature,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. InlineData (Image) 处理
|
||||
if part.InlineData != nil && part.InlineData.Data != "" {
|
||||
p.flushThinking()
|
||||
markdownImg := fmt.Sprintf("",
|
||||
part.InlineData.MimeType, part.InlineData.Data)
|
||||
p.textBuilder += markdownImg
|
||||
p.flushText()
|
||||
}
|
||||
}
|
||||
|
||||
// flushText 刷新 text builder
|
||||
func (p *NonStreamingProcessor) flushText() {
|
||||
if p.textBuilder == "" {
|
||||
return
|
||||
}
|
||||
|
||||
p.contentBlocks = append(p.contentBlocks, ClaudeContentItem{
|
||||
Type: "text",
|
||||
Text: p.textBuilder,
|
||||
})
|
||||
p.textBuilder = ""
|
||||
}
|
||||
|
||||
// flushThinking 刷新 thinking builder
|
||||
func (p *NonStreamingProcessor) flushThinking() {
|
||||
if p.thinkingBuilder == "" && p.thinkingSignature == "" {
|
||||
return
|
||||
}
|
||||
|
||||
p.contentBlocks = append(p.contentBlocks, ClaudeContentItem{
|
||||
Type: "thinking",
|
||||
Thinking: p.thinkingBuilder,
|
||||
Signature: p.thinkingSignature,
|
||||
})
|
||||
p.thinkingBuilder = ""
|
||||
p.thinkingSignature = ""
|
||||
}
|
||||
|
||||
// buildResponse 构建最终响应
|
||||
func (p *NonStreamingProcessor) buildResponse(geminiResp *GeminiResponse, responseID, originalModel string) *ClaudeResponse {
|
||||
var finishReason string
|
||||
if len(geminiResp.Candidates) > 0 {
|
||||
finishReason = geminiResp.Candidates[0].FinishReason
|
||||
}
|
||||
|
||||
stopReason := "end_turn"
|
||||
if p.hasToolCall {
|
||||
stopReason = "tool_use"
|
||||
} else if finishReason == "MAX_TOKENS" {
|
||||
stopReason = "max_tokens"
|
||||
}
|
||||
|
||||
usage := ClaudeUsage{}
|
||||
if geminiResp.UsageMetadata != nil {
|
||||
usage.InputTokens = geminiResp.UsageMetadata.PromptTokenCount
|
||||
usage.OutputTokens = geminiResp.UsageMetadata.CandidatesTokenCount
|
||||
}
|
||||
|
||||
// 生成响应 ID
|
||||
respID := responseID
|
||||
if respID == "" {
|
||||
respID = geminiResp.ResponseID
|
||||
}
|
||||
if respID == "" {
|
||||
respID = "msg_" + generateRandomID()
|
||||
}
|
||||
|
||||
return &ClaudeResponse{
|
||||
ID: respID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Model: originalModel,
|
||||
Content: p.contentBlocks,
|
||||
StopReason: stopReason,
|
||||
Usage: usage,
|
||||
}
|
||||
}
|
||||
|
||||
// generateRandomID 生成随机 ID
|
||||
func generateRandomID() string {
|
||||
const chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
|
||||
result := make([]byte, 12)
|
||||
for i := range result {
|
||||
result[i] = chars[i%len(chars)]
|
||||
}
|
||||
return string(result)
|
||||
}
|
||||
455
backend/internal/pkg/antigravity/stream_transformer.go
Normal file
455
backend/internal/pkg/antigravity/stream_transformer.go
Normal file
@@ -0,0 +1,455 @@
|
||||
package antigravity
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// BlockType 内容块类型
|
||||
type BlockType int
|
||||
|
||||
const (
|
||||
BlockTypeNone BlockType = iota
|
||||
BlockTypeText
|
||||
BlockTypeThinking
|
||||
BlockTypeFunction
|
||||
)
|
||||
|
||||
// StreamingProcessor 流式响应处理器
|
||||
type StreamingProcessor struct {
|
||||
blockType BlockType
|
||||
blockIndex int
|
||||
messageStartSent bool
|
||||
messageStopSent bool
|
||||
usedTool bool
|
||||
pendingSignature string
|
||||
trailingSignature string
|
||||
originalModel string
|
||||
|
||||
// 累计 usage
|
||||
inputTokens int
|
||||
outputTokens int
|
||||
}
|
||||
|
||||
// NewStreamingProcessor 创建流式响应处理器
|
||||
func NewStreamingProcessor(originalModel string) *StreamingProcessor {
|
||||
return &StreamingProcessor{
|
||||
blockType: BlockTypeNone,
|
||||
originalModel: originalModel,
|
||||
}
|
||||
}
|
||||
|
||||
// ProcessLine 处理 SSE 行,返回 Claude SSE 事件
|
||||
func (p *StreamingProcessor) ProcessLine(line string) []byte {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || !strings.HasPrefix(line, "data:") {
|
||||
return nil
|
||||
}
|
||||
|
||||
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if data == "" || data == "[DONE]" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 解包 v1internal 响应
|
||||
var v1Resp V1InternalResponse
|
||||
if err := json.Unmarshal([]byte(data), &v1Resp); err != nil {
|
||||
// 尝试直接解析为 GeminiResponse
|
||||
var directResp GeminiResponse
|
||||
if err2 := json.Unmarshal([]byte(data), &directResp); err2 != nil {
|
||||
return nil
|
||||
}
|
||||
v1Resp.Response = directResp
|
||||
v1Resp.ResponseID = directResp.ResponseID
|
||||
v1Resp.ModelVersion = directResp.ModelVersion
|
||||
}
|
||||
|
||||
geminiResp := &v1Resp.Response
|
||||
|
||||
var result bytes.Buffer
|
||||
|
||||
// 发送 message_start
|
||||
if !p.messageStartSent {
|
||||
_, _ = result.Write(p.emitMessageStart(&v1Resp))
|
||||
}
|
||||
|
||||
// 更新 usage
|
||||
if geminiResp.UsageMetadata != nil {
|
||||
p.inputTokens = geminiResp.UsageMetadata.PromptTokenCount
|
||||
p.outputTokens = geminiResp.UsageMetadata.CandidatesTokenCount
|
||||
}
|
||||
|
||||
// 处理 parts
|
||||
if len(geminiResp.Candidates) > 0 && geminiResp.Candidates[0].Content != nil {
|
||||
for _, part := range geminiResp.Candidates[0].Content.Parts {
|
||||
_, _ = result.Write(p.processPart(&part))
|
||||
}
|
||||
}
|
||||
|
||||
// 检查是否结束
|
||||
if len(geminiResp.Candidates) > 0 {
|
||||
finishReason := geminiResp.Candidates[0].FinishReason
|
||||
if finishReason != "" {
|
||||
_, _ = result.Write(p.emitFinish(finishReason))
|
||||
}
|
||||
}
|
||||
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// Finish 结束处理,返回最终事件和用量
|
||||
func (p *StreamingProcessor) Finish() ([]byte, *ClaudeUsage) {
|
||||
var result bytes.Buffer
|
||||
|
||||
if !p.messageStopSent {
|
||||
_, _ = result.Write(p.emitFinish(""))
|
||||
}
|
||||
|
||||
usage := &ClaudeUsage{
|
||||
InputTokens: p.inputTokens,
|
||||
OutputTokens: p.outputTokens,
|
||||
}
|
||||
|
||||
return result.Bytes(), usage
|
||||
}
|
||||
|
||||
// emitMessageStart 发送 message_start 事件
|
||||
func (p *StreamingProcessor) emitMessageStart(v1Resp *V1InternalResponse) []byte {
|
||||
if p.messageStartSent {
|
||||
return nil
|
||||
}
|
||||
|
||||
usage := ClaudeUsage{}
|
||||
if v1Resp.Response.UsageMetadata != nil {
|
||||
usage.InputTokens = v1Resp.Response.UsageMetadata.PromptTokenCount
|
||||
usage.OutputTokens = v1Resp.Response.UsageMetadata.CandidatesTokenCount
|
||||
}
|
||||
|
||||
responseID := v1Resp.ResponseID
|
||||
if responseID == "" {
|
||||
responseID = v1Resp.Response.ResponseID
|
||||
}
|
||||
if responseID == "" {
|
||||
responseID = "msg_" + generateRandomID()
|
||||
}
|
||||
|
||||
message := map[string]any{
|
||||
"id": responseID,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": []any{},
|
||||
"model": p.originalModel,
|
||||
"stop_reason": nil,
|
||||
"stop_sequence": nil,
|
||||
"usage": usage,
|
||||
}
|
||||
|
||||
event := map[string]any{
|
||||
"type": "message_start",
|
||||
"message": message,
|
||||
}
|
||||
|
||||
p.messageStartSent = true
|
||||
return p.formatSSE("message_start", event)
|
||||
}
|
||||
|
||||
// processPart 处理单个 part
|
||||
func (p *StreamingProcessor) processPart(part *GeminiPart) []byte {
|
||||
var result bytes.Buffer
|
||||
signature := part.ThoughtSignature
|
||||
|
||||
// 1. FunctionCall 处理
|
||||
if part.FunctionCall != nil {
|
||||
// 先处理 trailingSignature
|
||||
if p.trailingSignature != "" {
|
||||
_, _ = result.Write(p.endBlock())
|
||||
_, _ = result.Write(p.emitEmptyThinkingWithSignature(p.trailingSignature))
|
||||
p.trailingSignature = ""
|
||||
}
|
||||
|
||||
_, _ = result.Write(p.processFunctionCall(part.FunctionCall, signature))
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// 2. Text 处理
|
||||
if part.Text != "" || part.Thought {
|
||||
if part.Thought {
|
||||
_, _ = result.Write(p.processThinking(part.Text, signature))
|
||||
} else {
|
||||
_, _ = result.Write(p.processText(part.Text, signature))
|
||||
}
|
||||
}
|
||||
|
||||
// 3. InlineData (Image) 处理
|
||||
if part.InlineData != nil && part.InlineData.Data != "" {
|
||||
markdownImg := fmt.Sprintf("",
|
||||
part.InlineData.MimeType, part.InlineData.Data)
|
||||
_, _ = result.Write(p.processText(markdownImg, ""))
|
||||
}
|
||||
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// processThinking 处理 thinking
|
||||
func (p *StreamingProcessor) processThinking(text, signature string) []byte {
|
||||
var result bytes.Buffer
|
||||
|
||||
// 处理之前的 trailingSignature
|
||||
if p.trailingSignature != "" {
|
||||
_, _ = result.Write(p.endBlock())
|
||||
_, _ = result.Write(p.emitEmptyThinkingWithSignature(p.trailingSignature))
|
||||
p.trailingSignature = ""
|
||||
}
|
||||
|
||||
// 开始或继续 thinking 块
|
||||
if p.blockType != BlockTypeThinking {
|
||||
_, _ = result.Write(p.startBlock(BlockTypeThinking, map[string]any{
|
||||
"type": "thinking",
|
||||
"thinking": "",
|
||||
}))
|
||||
}
|
||||
|
||||
if text != "" {
|
||||
_, _ = result.Write(p.emitDelta("thinking_delta", map[string]any{
|
||||
"thinking": text,
|
||||
}))
|
||||
}
|
||||
|
||||
// 暂存签名
|
||||
if signature != "" {
|
||||
p.pendingSignature = signature
|
||||
}
|
||||
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// processText 处理普通 text
|
||||
func (p *StreamingProcessor) processText(text, signature string) []byte {
|
||||
var result bytes.Buffer
|
||||
|
||||
// 空 text 带签名 - 暂存
|
||||
if text == "" {
|
||||
if signature != "" {
|
||||
p.trailingSignature = signature
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 处理之前的 trailingSignature
|
||||
if p.trailingSignature != "" {
|
||||
_, _ = result.Write(p.endBlock())
|
||||
_, _ = result.Write(p.emitEmptyThinkingWithSignature(p.trailingSignature))
|
||||
p.trailingSignature = ""
|
||||
}
|
||||
|
||||
// 非空 text 带签名 - 特殊处理
|
||||
if signature != "" {
|
||||
_, _ = result.Write(p.startBlock(BlockTypeText, map[string]any{
|
||||
"type": "text",
|
||||
"text": "",
|
||||
}))
|
||||
_, _ = result.Write(p.emitDelta("text_delta", map[string]any{
|
||||
"text": text,
|
||||
}))
|
||||
_, _ = result.Write(p.endBlock())
|
||||
_, _ = result.Write(p.emitEmptyThinkingWithSignature(signature))
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// 普通 text (无签名)
|
||||
if p.blockType != BlockTypeText {
|
||||
_, _ = result.Write(p.startBlock(BlockTypeText, map[string]any{
|
||||
"type": "text",
|
||||
"text": "",
|
||||
}))
|
||||
}
|
||||
|
||||
_, _ = result.Write(p.emitDelta("text_delta", map[string]any{
|
||||
"text": text,
|
||||
}))
|
||||
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// processFunctionCall 处理 function call
|
||||
func (p *StreamingProcessor) processFunctionCall(fc *GeminiFunctionCall, signature string) []byte {
|
||||
var result bytes.Buffer
|
||||
|
||||
p.usedTool = true
|
||||
|
||||
toolID := fc.ID
|
||||
if toolID == "" {
|
||||
toolID = fmt.Sprintf("%s-%s", fc.Name, generateRandomID())
|
||||
}
|
||||
|
||||
toolUse := map[string]any{
|
||||
"type": "tool_use",
|
||||
"id": toolID,
|
||||
"name": fc.Name,
|
||||
"input": map[string]any{},
|
||||
}
|
||||
|
||||
if signature != "" {
|
||||
toolUse["signature"] = signature
|
||||
}
|
||||
|
||||
_, _ = result.Write(p.startBlock(BlockTypeFunction, toolUse))
|
||||
|
||||
// 发送 input_json_delta
|
||||
if fc.Args != nil {
|
||||
argsJSON, _ := json.Marshal(fc.Args)
|
||||
_, _ = result.Write(p.emitDelta("input_json_delta", map[string]any{
|
||||
"partial_json": string(argsJSON),
|
||||
}))
|
||||
}
|
||||
|
||||
_, _ = result.Write(p.endBlock())
|
||||
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// startBlock 开始新的内容块
|
||||
func (p *StreamingProcessor) startBlock(blockType BlockType, contentBlock map[string]any) []byte {
|
||||
var result bytes.Buffer
|
||||
|
||||
if p.blockType != BlockTypeNone {
|
||||
_, _ = result.Write(p.endBlock())
|
||||
}
|
||||
|
||||
event := map[string]any{
|
||||
"type": "content_block_start",
|
||||
"index": p.blockIndex,
|
||||
"content_block": contentBlock,
|
||||
}
|
||||
|
||||
_, _ = result.Write(p.formatSSE("content_block_start", event))
|
||||
p.blockType = blockType
|
||||
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// endBlock 结束当前内容块
|
||||
func (p *StreamingProcessor) endBlock() []byte {
|
||||
if p.blockType == BlockTypeNone {
|
||||
return nil
|
||||
}
|
||||
|
||||
var result bytes.Buffer
|
||||
|
||||
// Thinking 块结束时发送暂存的签名
|
||||
if p.blockType == BlockTypeThinking && p.pendingSignature != "" {
|
||||
_, _ = result.Write(p.emitDelta("signature_delta", map[string]any{
|
||||
"signature": p.pendingSignature,
|
||||
}))
|
||||
p.pendingSignature = ""
|
||||
}
|
||||
|
||||
event := map[string]any{
|
||||
"type": "content_block_stop",
|
||||
"index": p.blockIndex,
|
||||
}
|
||||
|
||||
_, _ = result.Write(p.formatSSE("content_block_stop", event))
|
||||
|
||||
p.blockIndex++
|
||||
p.blockType = BlockTypeNone
|
||||
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// emitDelta 发送 delta 事件
|
||||
func (p *StreamingProcessor) emitDelta(deltaType string, deltaContent map[string]any) []byte {
|
||||
delta := map[string]any{
|
||||
"type": deltaType,
|
||||
}
|
||||
for k, v := range deltaContent {
|
||||
delta[k] = v
|
||||
}
|
||||
|
||||
event := map[string]any{
|
||||
"type": "content_block_delta",
|
||||
"index": p.blockIndex,
|
||||
"delta": delta,
|
||||
}
|
||||
|
||||
return p.formatSSE("content_block_delta", event)
|
||||
}
|
||||
|
||||
// emitEmptyThinkingWithSignature 发送空 thinking 块承载签名
|
||||
func (p *StreamingProcessor) emitEmptyThinkingWithSignature(signature string) []byte {
|
||||
var result bytes.Buffer
|
||||
|
||||
_, _ = result.Write(p.startBlock(BlockTypeThinking, map[string]any{
|
||||
"type": "thinking",
|
||||
"thinking": "",
|
||||
}))
|
||||
_, _ = result.Write(p.emitDelta("thinking_delta", map[string]any{
|
||||
"thinking": "",
|
||||
}))
|
||||
_, _ = result.Write(p.emitDelta("signature_delta", map[string]any{
|
||||
"signature": signature,
|
||||
}))
|
||||
_, _ = result.Write(p.endBlock())
|
||||
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// emitFinish 发送结束事件
|
||||
func (p *StreamingProcessor) emitFinish(finishReason string) []byte {
|
||||
var result bytes.Buffer
|
||||
|
||||
// 关闭最后一个块
|
||||
_, _ = result.Write(p.endBlock())
|
||||
|
||||
// 处理 trailingSignature
|
||||
if p.trailingSignature != "" {
|
||||
_, _ = result.Write(p.emitEmptyThinkingWithSignature(p.trailingSignature))
|
||||
p.trailingSignature = ""
|
||||
}
|
||||
|
||||
// 确定 stop_reason
|
||||
stopReason := "end_turn"
|
||||
if p.usedTool {
|
||||
stopReason = "tool_use"
|
||||
} else if finishReason == "MAX_TOKENS" {
|
||||
stopReason = "max_tokens"
|
||||
}
|
||||
|
||||
usage := ClaudeUsage{
|
||||
InputTokens: p.inputTokens,
|
||||
OutputTokens: p.outputTokens,
|
||||
}
|
||||
|
||||
deltaEvent := map[string]any{
|
||||
"type": "message_delta",
|
||||
"delta": map[string]any{
|
||||
"stop_reason": stopReason,
|
||||
"stop_sequence": nil,
|
||||
},
|
||||
"usage": usage,
|
||||
}
|
||||
|
||||
_, _ = result.Write(p.formatSSE("message_delta", deltaEvent))
|
||||
|
||||
if !p.messageStopSent {
|
||||
stopEvent := map[string]any{
|
||||
"type": "message_stop",
|
||||
}
|
||||
_, _ = result.Write(p.formatSSE("message_stop", stopEvent))
|
||||
p.messageStopSent = true
|
||||
}
|
||||
|
||||
return result.Bytes()
|
||||
}
|
||||
|
||||
// formatSSE 格式化 SSE 事件
|
||||
func (p *StreamingProcessor) formatSSE(eventType string, data any) []byte {
|
||||
jsonData, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
return []byte(fmt.Sprintf("event: %s\ndata: %s\n\n", eventType, string(jsonData)))
|
||||
}
|
||||
10
backend/internal/pkg/ctxkey/ctxkey.go
Normal file
10
backend/internal/pkg/ctxkey/ctxkey.go
Normal file
@@ -0,0 +1,10 @@
|
||||
// Package ctxkey 定义用于 context.Value 的类型安全 key
|
||||
package ctxkey
|
||||
|
||||
// Key 定义 context key 的类型,避免使用内置 string 类型(staticcheck SA1029)
|
||||
type Key string
|
||||
|
||||
const (
|
||||
// ForcePlatform 强制平台(用于 /antigravity 路由),由 middleware.ForcePlatform 设置
|
||||
ForcePlatform Key = "ctx_force_platform"
|
||||
)
|
||||
@@ -18,8 +18,10 @@ func DefaultModels() []Model {
|
||||
methods := []string{"generateContent", "streamGenerateContent"}
|
||||
return []Model{
|
||||
{Name: "models/gemini-3-pro-preview", SupportedGenerationMethods: methods},
|
||||
{Name: "models/gemini-3-flash-preview", SupportedGenerationMethods: methods},
|
||||
{Name: "models/gemini-2.5-pro", SupportedGenerationMethods: methods},
|
||||
{Name: "models/gemini-2.5-flash", SupportedGenerationMethods: methods},
|
||||
{Name: "models/gemini-2.0-flash", SupportedGenerationMethods: methods},
|
||||
{Name: "models/gemini-2.0-flash-lite", SupportedGenerationMethods: methods},
|
||||
{Name: "models/gemini-1.5-pro", SupportedGenerationMethods: methods},
|
||||
{Name: "models/gemini-1.5-flash", SupportedGenerationMethods: methods},
|
||||
{Name: "models/gemini-1.5-flash-8b", SupportedGenerationMethods: methods},
|
||||
|
||||
@@ -11,11 +11,11 @@ type Model struct {
|
||||
|
||||
// DefaultModels is the curated Gemini model list used by the admin UI "test account" flow.
|
||||
var DefaultModels = []Model{
|
||||
{ID: "gemini-3-pro", Type: "model", DisplayName: "Gemini 3 Pro", CreatedAt: ""},
|
||||
{ID: "gemini-3-flash", Type: "model", DisplayName: "Gemini 3 Flash", CreatedAt: ""},
|
||||
{ID: "gemini-3-pro-preview", Type: "model", DisplayName: "Gemini 3 Pro Preview", CreatedAt: ""},
|
||||
{ID: "gemini-3-flash-preview", Type: "model", DisplayName: "Gemini 3 Flash Preview", CreatedAt: ""},
|
||||
{ID: "gemini-2.5-pro", Type: "model", DisplayName: "Gemini 2.5 Pro", CreatedAt: ""},
|
||||
{ID: "gemini-2.5-flash", Type: "model", DisplayName: "Gemini 2.5 Flash", CreatedAt: ""},
|
||||
}
|
||||
|
||||
// DefaultTestModel is the default model to preselect in test flows.
|
||||
const DefaultTestModel = "gemini-2.5-pro"
|
||||
const DefaultTestModel = "gemini-3-pro-preview"
|
||||
|
||||
@@ -127,10 +127,15 @@ type UserDashboardStats struct {
|
||||
|
||||
// UsageLogFilters represents filters for usage log queries
|
||||
type UsageLogFilters struct {
|
||||
UserID int64
|
||||
ApiKeyID int64
|
||||
StartTime *time.Time
|
||||
EndTime *time.Time
|
||||
UserID int64
|
||||
ApiKeyID int64
|
||||
AccountID int64
|
||||
GroupID int64
|
||||
Model string
|
||||
Stream *bool
|
||||
BillingType *int8
|
||||
StartTime *time.Time
|
||||
EndTime *time.Time
|
||||
}
|
||||
|
||||
// UsageStats represents usage statistics
|
||||
|
||||
@@ -171,6 +171,27 @@ func (r *accountRepository) UpdateLastUsed(ctx context.Context, id int64) error
|
||||
return r.db.WithContext(ctx).Model(&accountModel{}).Where("id = ?", id).Update("last_used_at", now).Error
|
||||
}
|
||||
|
||||
func (r *accountRepository) BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
|
||||
if len(updates) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var caseSql = "UPDATE accounts SET last_used_at = CASE id"
|
||||
var args []any
|
||||
var ids []int64
|
||||
|
||||
for id, ts := range updates {
|
||||
caseSql += " WHEN ? THEN CAST(? AS TIMESTAMP)"
|
||||
args = append(args, id, ts)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
caseSql += " END WHERE id IN ? AND deleted_at IS NULL"
|
||||
args = append(args, ids)
|
||||
|
||||
return r.db.WithContext(ctx).Exec(caseSql, args...).Error
|
||||
}
|
||||
|
||||
func (r *accountRepository) SetError(ctx context.Context, id int64, errorMsg string) error {
|
||||
return r.db.WithContext(ctx).Model(&accountModel{}).Where("id = ?", id).
|
||||
Updates(map[string]any{
|
||||
@@ -316,6 +337,56 @@ func (r *accountRepository) ListSchedulableByGroupIDAndPlatform(ctx context.Cont
|
||||
return outAccounts, nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]service.Account, error) {
|
||||
if len(platforms) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var accounts []accountModel
|
||||
now := time.Now()
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("platform IN ?", platforms).
|
||||
Where("status = ? AND schedulable = ?", service.StatusActive, true).
|
||||
Where("(overload_until IS NULL OR overload_until <= ?)", now).
|
||||
Where("(rate_limit_reset_at IS NULL OR rate_limit_reset_at <= ?)", now).
|
||||
Preload("Proxy").
|
||||
Order("priority ASC").
|
||||
Find(&accounts).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
outAccounts := make([]service.Account, 0, len(accounts))
|
||||
for i := range accounts {
|
||||
outAccounts = append(outAccounts, *accountModelToService(&accounts[i]))
|
||||
}
|
||||
return outAccounts, nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]service.Account, error) {
|
||||
if len(platforms) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var accounts []accountModel
|
||||
now := time.Now()
|
||||
err := r.db.WithContext(ctx).
|
||||
Joins("JOIN account_groups ON account_groups.account_id = accounts.id").
|
||||
Where("account_groups.group_id = ?", groupID).
|
||||
Where("accounts.platform IN ?", platforms).
|
||||
Where("accounts.status = ? AND accounts.schedulable = ?", service.StatusActive, true).
|
||||
Where("(accounts.overload_until IS NULL OR accounts.overload_until <= ?)", now).
|
||||
Where("(accounts.rate_limit_reset_at IS NULL OR accounts.rate_limit_reset_at <= ?)", now).
|
||||
Preload("Proxy").
|
||||
Order("account_groups.priority ASC, accounts.priority ASC").
|
||||
Find(&accounts).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
outAccounts := make([]service.Account, 0, len(accounts))
|
||||
for i := range accounts {
|
||||
outAccounts = append(outAccounts, *accountModelToService(&accounts[i]))
|
||||
}
|
||||
return outAccounts, nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
|
||||
now := time.Now()
|
||||
return r.db.WithContext(ctx).Model(&accountModel{}).Where("id = ?", id).
|
||||
|
||||
@@ -81,6 +81,22 @@ func (r *apiKeyRepository) ListByUserID(ctx context.Context, userID int64, param
|
||||
return outKeys, paginationResultFromTotal(total, params), nil
|
||||
}
|
||||
|
||||
func (r *apiKeyRepository) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
|
||||
if len(apiKeyIDs) == 0 {
|
||||
return []int64{}, nil
|
||||
}
|
||||
|
||||
ids := make([]int64, 0, len(apiKeyIDs))
|
||||
err := r.db.WithContext(ctx).
|
||||
Model(&apiKeyModel{}).
|
||||
Where("user_id = ? AND id IN ?", userID, apiKeyIDs).
|
||||
Pluck("id", &ids).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
func (r *apiKeyRepository) CountByUserID(ctx context.Context, userID int64) (int64, error) {
|
||||
var count int64
|
||||
err := r.db.WithContext(ctx).Model(&apiKeyModel{}).Where("user_id = ?", userID).Count(&count).Error
|
||||
|
||||
@@ -1,11 +1,21 @@
|
||||
package repository
|
||||
|
||||
import "gorm.io/gorm"
|
||||
import (
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// MaxExpiresAt is the maximum allowed expiration date for subscriptions (year 2099)
|
||||
// This prevents time.Time JSON serialization errors (RFC 3339 requires year <= 9999)
|
||||
var maxExpiresAt = time.Date(2099, 12, 31, 23, 59, 59, 0, time.UTC)
|
||||
|
||||
// AutoMigrate runs schema migrations for all repository persistence models.
|
||||
// Persistence models are defined within individual `*_repo.go` files.
|
||||
func AutoMigrate(db *gorm.DB) error {
|
||||
return db.AutoMigrate(
|
||||
// runMode: "standard" or "simple" - determines whether to create default groups
|
||||
func AutoMigrate(db *gorm.DB, runMode string) error {
|
||||
err := db.AutoMigrate(
|
||||
&userModel{},
|
||||
&apiKeyModel{},
|
||||
&groupModel{},
|
||||
@@ -17,4 +27,109 @@ func AutoMigrate(db *gorm.DB) error {
|
||||
&settingModel{},
|
||||
&userSubscriptionModel{},
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 创建默认分组(简易模式支持)
|
||||
if err := ensureDefaultGroups(db, runMode); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 修复无效的过期时间(年份超过 2099 会导致 JSON 序列化失败)
|
||||
return fixInvalidExpiresAt(db)
|
||||
}
|
||||
|
||||
// fixInvalidExpiresAt 修复 user_subscriptions 表中无效的过期时间
|
||||
func fixInvalidExpiresAt(db *gorm.DB) error {
|
||||
result := db.Model(&userSubscriptionModel{}).
|
||||
Where("expires_at > ?", maxExpiresAt).
|
||||
Update("expires_at", maxExpiresAt)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected > 0 {
|
||||
log.Printf("[AutoMigrate] Fixed %d subscriptions with invalid expires_at (year > 2099)", result.RowsAffected)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ensureDefaultGroups 确保默认分组存在(简易模式支持)
|
||||
// 为每个平台创建一个默认分组,配置最大权限以确保简易模式下不受限制
|
||||
// runMode: "standard" 时跳过创建, "simple" 时创建/恢复默认分组
|
||||
func ensureDefaultGroups(db *gorm.DB, runMode string) error {
|
||||
// 标准版不创建默认分组
|
||||
if runMode == "standard" {
|
||||
return nil
|
||||
}
|
||||
|
||||
defaultGroups := []struct {
|
||||
name string
|
||||
platform string
|
||||
description string
|
||||
}{
|
||||
{
|
||||
name: "anthropic-default",
|
||||
platform: "anthropic",
|
||||
description: "Default group for Anthropic accounts (Simple Mode)",
|
||||
},
|
||||
{
|
||||
name: "openai-default",
|
||||
platform: "openai",
|
||||
description: "Default group for OpenAI accounts (Simple Mode)",
|
||||
},
|
||||
{
|
||||
name: "gemini-default",
|
||||
platform: "gemini",
|
||||
description: "Default group for Gemini accounts (Simple Mode)",
|
||||
},
|
||||
}
|
||||
|
||||
for _, dg := range defaultGroups {
|
||||
// 步骤1: 检查是否有软删除的记录
|
||||
var softDeletedCount int64
|
||||
if err := db.Unscoped().Model(&groupModel{}).
|
||||
Where("name = ? AND deleted_at IS NOT NULL", dg.name).
|
||||
Count(&softDeletedCount).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if softDeletedCount > 0 {
|
||||
// 恢复软删除的记录
|
||||
if err := db.Unscoped().Model(&groupModel{}).
|
||||
Where("name = ?", dg.name).
|
||||
Update("deleted_at", nil).Error; err != nil {
|
||||
log.Printf("[AutoMigrate] Failed to restore default group %s: %v", dg.name, err)
|
||||
return err
|
||||
}
|
||||
log.Printf("[AutoMigrate] Restored default group: %s (platform: %s)", dg.name, dg.platform)
|
||||
continue
|
||||
}
|
||||
|
||||
// 步骤2: 检查是否有活跃记录
|
||||
var activeCount int64
|
||||
if err := db.Model(&groupModel{}).Where("name = ?", dg.name).Count(&activeCount).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if activeCount == 0 {
|
||||
// 创建新分组
|
||||
group := &groupModel{
|
||||
Name: dg.name,
|
||||
Description: dg.description,
|
||||
Platform: dg.platform,
|
||||
RateMultiplier: 1.0,
|
||||
IsExclusive: false,
|
||||
Status: "active",
|
||||
SubscriptionType: "standard",
|
||||
}
|
||||
if err := db.Create(group).Error; err != nil {
|
||||
log.Printf("[AutoMigrate] Failed to create default group %s: %v", dg.name, err)
|
||||
return err
|
||||
}
|
||||
log.Printf("[AutoMigrate] Created default group: %s (platform: %s)", dg.name, dg.platform)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -145,7 +145,7 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe
|
||||
return fullCode, nil
|
||||
}
|
||||
|
||||
func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, codeVerifier, state, proxyURL string) (*oauth.TokenResponse, error) {
|
||||
func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, codeVerifier, state, proxyURL string, isSetupToken bool) (*oauth.TokenResponse, error) {
|
||||
client := s.clientFactory(proxyURL)
|
||||
|
||||
// Parse code which may contain state in format "authCode#state"
|
||||
@@ -168,6 +168,11 @@ func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, cod
|
||||
reqBody["state"] = codeState
|
||||
}
|
||||
|
||||
// Setup token requires longer expiration (1 year)
|
||||
if isSetupToken {
|
||||
reqBody["expires_in"] = 31536000 // 365 * 24 * 60 * 60 seconds
|
||||
}
|
||||
|
||||
reqBodyJSON, _ := json.Marshal(reqBody)
|
||||
log.Printf("[OAuth] Step 3: Exchanging code for token at %s", s.tokenURL)
|
||||
log.Printf("[OAuth] Step 3 Request Body: %s", string(reqBodyJSON))
|
||||
@@ -199,16 +204,20 @@ func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, cod
|
||||
func (s *claudeOAuthService) RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*oauth.TokenResponse, error) {
|
||||
client := s.clientFactory(proxyURL)
|
||||
|
||||
formData := url.Values{}
|
||||
formData.Set("grant_type", "refresh_token")
|
||||
formData.Set("refresh_token", refreshToken)
|
||||
formData.Set("client_id", oauth.ClientID)
|
||||
// 使用 JSON 格式(与 ExchangeCodeForToken 保持一致)
|
||||
// Anthropic OAuth API 期望 JSON 格式的请求体
|
||||
reqBody := map[string]any{
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refreshToken,
|
||||
"client_id": oauth.ClientID,
|
||||
}
|
||||
|
||||
var tokenResp oauth.TokenResponse
|
||||
|
||||
resp, err := client.R().
|
||||
SetContext(ctx).
|
||||
SetFormDataFromValues(formData).
|
||||
SetHeader("Content-Type", "application/json").
|
||||
SetBody(reqBody).
|
||||
SetSuccessResult(&tokenResp).
|
||||
Post(s.tokenURL)
|
||||
|
||||
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
@@ -34,7 +33,6 @@ type requestCapture struct {
|
||||
method string
|
||||
cookies []*http.Cookie
|
||||
body []byte
|
||||
formValues url.Values
|
||||
bodyJSON map[string]any
|
||||
contentType string
|
||||
}
|
||||
@@ -193,12 +191,13 @@ func (s *ClaudeOAuthServiceSuite) TestGetAuthorizationCode() {
|
||||
|
||||
func (s *ClaudeOAuthServiceSuite) TestExchangeCodeForToken() {
|
||||
tests := []struct {
|
||||
name string
|
||||
handler http.HandlerFunc
|
||||
code string
|
||||
wantErr bool
|
||||
wantResp *oauth.TokenResponse
|
||||
validate func(captured requestCapture)
|
||||
name string
|
||||
handler http.HandlerFunc
|
||||
code string
|
||||
isSetupToken bool
|
||||
wantErr bool
|
||||
wantResp *oauth.TokenResponse
|
||||
validate func(captured requestCapture)
|
||||
}{
|
||||
{
|
||||
name: "sends_state_when_embedded",
|
||||
@@ -212,7 +211,8 @@ func (s *ClaudeOAuthServiceSuite) TestExchangeCodeForToken() {
|
||||
Scope: "s",
|
||||
})
|
||||
},
|
||||
code: "AUTH#STATE2",
|
||||
code: "AUTH#STATE2",
|
||||
isSetupToken: false,
|
||||
wantResp: &oauth.TokenResponse{
|
||||
AccessToken: "at",
|
||||
RefreshToken: "rt",
|
||||
@@ -225,6 +225,29 @@ func (s *ClaudeOAuthServiceSuite) TestExchangeCodeForToken() {
|
||||
require.Equal(s.T(), oauth.ClientID, captured.bodyJSON["client_id"])
|
||||
require.Equal(s.T(), oauth.RedirectURI, captured.bodyJSON["redirect_uri"])
|
||||
require.Equal(s.T(), "ver", captured.bodyJSON["code_verifier"])
|
||||
// Regular OAuth should not include expires_in
|
||||
require.Nil(s.T(), captured.bodyJSON["expires_in"], "regular OAuth should not include expires_in")
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "setup_token_includes_expires_in",
|
||||
handler: func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(oauth.TokenResponse{
|
||||
AccessToken: "at",
|
||||
TokenType: "bearer",
|
||||
ExpiresIn: 31536000,
|
||||
})
|
||||
},
|
||||
code: "AUTH",
|
||||
isSetupToken: true,
|
||||
wantResp: &oauth.TokenResponse{
|
||||
AccessToken: "at",
|
||||
},
|
||||
validate: func(captured requestCapture) {
|
||||
// Setup token should include expires_in with 1 year value
|
||||
require.Equal(s.T(), float64(31536000), captured.bodyJSON["expires_in"],
|
||||
"setup token should include expires_in: 31536000")
|
||||
},
|
||||
},
|
||||
{
|
||||
@@ -233,8 +256,9 @@ func (s *ClaudeOAuthServiceSuite) TestExchangeCodeForToken() {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_, _ = w.Write([]byte("bad request"))
|
||||
},
|
||||
code: "AUTH",
|
||||
wantErr: true,
|
||||
code: "AUTH",
|
||||
isSetupToken: false,
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -256,7 +280,7 @@ func (s *ClaudeOAuthServiceSuite) TestExchangeCodeForToken() {
|
||||
s.client = client
|
||||
s.client.tokenURL = s.srv.URL
|
||||
|
||||
resp, err := s.client.ExchangeCodeForToken(context.Background(), tt.code, "ver", "", "")
|
||||
resp, err := s.client.ExchangeCodeForToken(context.Background(), tt.code, "ver", "", "", tt.isSetupToken)
|
||||
|
||||
if tt.wantErr {
|
||||
require.Error(s.T(), err)
|
||||
@@ -282,24 +306,53 @@ func (s *ClaudeOAuthServiceSuite) TestRefreshToken() {
|
||||
validate func(captured requestCapture)
|
||||
}{
|
||||
{
|
||||
name: "sends_form",
|
||||
name: "sends_json_format",
|
||||
handler: func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(oauth.TokenResponse{AccessToken: "at2", TokenType: "bearer", ExpiresIn: 3600})
|
||||
_ = json.NewEncoder(w).Encode(oauth.TokenResponse{
|
||||
AccessToken: "new_access_token",
|
||||
TokenType: "bearer",
|
||||
ExpiresIn: 28800,
|
||||
RefreshToken: "new_refresh_token",
|
||||
Scope: "user:profile user:inference",
|
||||
})
|
||||
},
|
||||
wantResp: &oauth.TokenResponse{
|
||||
AccessToken: "new_access_token",
|
||||
RefreshToken: "new_refresh_token",
|
||||
},
|
||||
wantResp: &oauth.TokenResponse{AccessToken: "at2"},
|
||||
validate: func(captured requestCapture) {
|
||||
require.Equal(s.T(), http.MethodPost, captured.method, "expected POST")
|
||||
require.Equal(s.T(), "refresh_token", captured.formValues.Get("grant_type"))
|
||||
require.Equal(s.T(), "rt", captured.formValues.Get("refresh_token"))
|
||||
require.Equal(s.T(), oauth.ClientID, captured.formValues.Get("client_id"))
|
||||
// 验证使用 JSON 格式(不是 form 格式)
|
||||
require.True(s.T(), strings.HasPrefix(captured.contentType, "application/json"),
|
||||
"expected JSON content-type, got: %s", captured.contentType)
|
||||
// 验证 JSON body 内容
|
||||
require.Equal(s.T(), "refresh_token", captured.bodyJSON["grant_type"])
|
||||
require.Equal(s.T(), "rt", captured.bodyJSON["refresh_token"])
|
||||
require.Equal(s.T(), oauth.ClientID, captured.bodyJSON["client_id"])
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "returns_new_refresh_token",
|
||||
handler: func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(oauth.TokenResponse{
|
||||
AccessToken: "at",
|
||||
TokenType: "bearer",
|
||||
ExpiresIn: 28800,
|
||||
RefreshToken: "rotated_rt", // Anthropic rotates refresh tokens
|
||||
})
|
||||
},
|
||||
wantResp: &oauth.TokenResponse{
|
||||
AccessToken: "at",
|
||||
RefreshToken: "rotated_rt",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "non_200_returns_error",
|
||||
handler: func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte("unauthorized"))
|
||||
_, _ = w.Write([]byte(`{"error":"invalid_grant"}`))
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
@@ -311,8 +364,9 @@ func (s *ClaudeOAuthServiceSuite) TestRefreshToken() {
|
||||
|
||||
s.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
captured.method = r.Method
|
||||
captured.contentType = r.Header.Get("Content-Type")
|
||||
captured.body, _ = io.ReadAll(r.Body)
|
||||
captured.formValues, _ = url.ParseQuery(string(captured.body))
|
||||
_ = json.Unmarshal(captured.body, &captured.bodyJSON)
|
||||
tt.handler(w, r)
|
||||
}))
|
||||
defer s.srv.Close()
|
||||
@@ -331,6 +385,7 @@ func (s *ClaudeOAuthServiceSuite) TestRefreshToken() {
|
||||
|
||||
require.NoError(s.T(), err)
|
||||
require.Equal(s.T(), tt.wantResp.AccessToken, resp.AccessToken)
|
||||
require.Equal(s.T(), tt.wantResp.RefreshToken, resp.RefreshToken)
|
||||
if tt.validate != nil {
|
||||
tt.validate(captured)
|
||||
}
|
||||
|
||||
250
backend/internal/repository/gateway_routing_integration_test.go
Normal file
250
backend/internal/repository/gateway_routing_integration_test.go
Normal file
@@ -0,0 +1,250 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/suite"
|
||||
"gorm.io/datatypes"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// GatewayRoutingSuite 测试网关路由相关的数据库查询
|
||||
// 验证账户选择和分流逻辑在真实数据库环境下的行为
|
||||
type GatewayRoutingSuite struct {
|
||||
suite.Suite
|
||||
ctx context.Context
|
||||
db *gorm.DB
|
||||
accountRepo *accountRepository
|
||||
}
|
||||
|
||||
func (s *GatewayRoutingSuite) SetupTest() {
|
||||
s.ctx = context.Background()
|
||||
s.db = testTx(s.T())
|
||||
s.accountRepo = NewAccountRepository(s.db).(*accountRepository)
|
||||
}
|
||||
|
||||
func TestGatewayRoutingSuite(t *testing.T) {
|
||||
suite.Run(t, new(GatewayRoutingSuite))
|
||||
}
|
||||
|
||||
// TestListSchedulableByPlatforms_GeminiAndAntigravity 验证多平台账户查询
|
||||
func (s *GatewayRoutingSuite) TestListSchedulableByPlatforms_GeminiAndAntigravity() {
|
||||
// 创建各平台账户
|
||||
geminiAcc := mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "gemini-oauth",
|
||||
Platform: service.PlatformGemini,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Priority: 1,
|
||||
})
|
||||
|
||||
antigravityAcc := mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "antigravity-oauth",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Priority: 2,
|
||||
Credentials: datatypes.JSONMap{
|
||||
"access_token": "test-token",
|
||||
"refresh_token": "test-refresh",
|
||||
"project_id": "test-project",
|
||||
},
|
||||
})
|
||||
|
||||
// 创建不应被选中的 anthropic 账户
|
||||
mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "anthropic-oauth",
|
||||
Platform: service.PlatformAnthropic,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Priority: 0,
|
||||
})
|
||||
|
||||
// 查询 gemini + antigravity 平台
|
||||
accounts, err := s.accountRepo.ListSchedulableByPlatforms(s.ctx, []string{
|
||||
service.PlatformGemini,
|
||||
service.PlatformAntigravity,
|
||||
})
|
||||
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(accounts, 2, "应返回 gemini 和 antigravity 两个账户")
|
||||
|
||||
// 验证返回的账户平台
|
||||
platforms := make(map[string]bool)
|
||||
for _, acc := range accounts {
|
||||
platforms[acc.Platform] = true
|
||||
}
|
||||
s.Require().True(platforms[service.PlatformGemini], "应包含 gemini 账户")
|
||||
s.Require().True(platforms[service.PlatformAntigravity], "应包含 antigravity 账户")
|
||||
s.Require().False(platforms[service.PlatformAnthropic], "不应包含 anthropic 账户")
|
||||
|
||||
// 验证账户 ID 匹配
|
||||
ids := make(map[int64]bool)
|
||||
for _, acc := range accounts {
|
||||
ids[acc.ID] = true
|
||||
}
|
||||
s.Require().True(ids[geminiAcc.ID])
|
||||
s.Require().True(ids[antigravityAcc.ID])
|
||||
}
|
||||
|
||||
// TestListSchedulableByGroupIDAndPlatforms_WithGroupBinding 验证按分组过滤
|
||||
func (s *GatewayRoutingSuite) TestListSchedulableByGroupIDAndPlatforms_WithGroupBinding() {
|
||||
// 创建 gemini 分组
|
||||
group := mustCreateGroup(s.T(), s.db, &groupModel{
|
||||
Name: "gemini-group",
|
||||
Platform: service.PlatformGemini,
|
||||
Status: service.StatusActive,
|
||||
})
|
||||
|
||||
// 创建账户
|
||||
boundAcc := mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "bound-antigravity",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
})
|
||||
unboundAcc := mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "unbound-antigravity",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
})
|
||||
|
||||
// 只绑定一个账户到分组
|
||||
mustBindAccountToGroup(s.T(), s.db, boundAcc.ID, group.ID, 1)
|
||||
|
||||
// 查询分组内的账户
|
||||
accounts, err := s.accountRepo.ListSchedulableByGroupIDAndPlatforms(s.ctx, group.ID, []string{
|
||||
service.PlatformGemini,
|
||||
service.PlatformAntigravity,
|
||||
})
|
||||
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(accounts, 1, "应只返回绑定到分组的账户")
|
||||
s.Require().Equal(boundAcc.ID, accounts[0].ID)
|
||||
|
||||
// 确认未绑定的账户不在结果中
|
||||
for _, acc := range accounts {
|
||||
s.Require().NotEqual(unboundAcc.ID, acc.ID, "不应包含未绑定的账户")
|
||||
}
|
||||
}
|
||||
|
||||
// TestListSchedulableByPlatform_Antigravity 验证单平台查询
|
||||
func (s *GatewayRoutingSuite) TestListSchedulableByPlatform_Antigravity() {
|
||||
// 创建多种平台账户
|
||||
mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "gemini-1",
|
||||
Platform: service.PlatformGemini,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
})
|
||||
|
||||
antigravity := mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "antigravity-1",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
})
|
||||
|
||||
// 只查询 antigravity 平台
|
||||
accounts, err := s.accountRepo.ListSchedulableByPlatform(s.ctx, service.PlatformAntigravity)
|
||||
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(accounts, 1)
|
||||
s.Require().Equal(antigravity.ID, accounts[0].ID)
|
||||
s.Require().Equal(service.PlatformAntigravity, accounts[0].Platform)
|
||||
}
|
||||
|
||||
// TestSchedulableFilter_ExcludesInactive 验证不可调度账户被过滤
|
||||
func (s *GatewayRoutingSuite) TestSchedulableFilter_ExcludesInactive() {
|
||||
// 创建可调度账户
|
||||
activeAcc := mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "active-antigravity",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
})
|
||||
|
||||
// 创建不可调度账户(需要先创建再更新,因为 fixture 默认设置 Schedulable=true)
|
||||
inactiveAcc := mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "inactive-antigravity",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Status: service.StatusActive,
|
||||
})
|
||||
s.Require().NoError(s.db.Model(&accountModel{}).Where("id = ?", inactiveAcc.ID).Update("schedulable", false).Error)
|
||||
|
||||
// 创建错误状态账户
|
||||
mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "error-antigravity",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Status: service.StatusError,
|
||||
Schedulable: true,
|
||||
})
|
||||
|
||||
accounts, err := s.accountRepo.ListSchedulableByPlatform(s.ctx, service.PlatformAntigravity)
|
||||
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(accounts, 1, "应只返回可调度的 active 账户")
|
||||
s.Require().Equal(activeAcc.ID, accounts[0].ID)
|
||||
}
|
||||
|
||||
// TestPlatformRoutingDecision 验证平台路由决策
|
||||
// 这个测试模拟 Handler 层在选择账户后的路由决策逻辑
|
||||
func (s *GatewayRoutingSuite) TestPlatformRoutingDecision() {
|
||||
// 创建两种平台的账户
|
||||
geminiAcc := mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "gemini-route-test",
|
||||
Platform: service.PlatformGemini,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
})
|
||||
|
||||
antigravityAcc := mustCreateAccount(s.T(), s.db, &accountModel{
|
||||
Name: "antigravity-route-test",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
})
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountID int64
|
||||
expectedService string
|
||||
}{
|
||||
{
|
||||
name: "Gemini账户路由到ForwardNative",
|
||||
accountID: geminiAcc.ID,
|
||||
expectedService: "GeminiMessagesCompatService.ForwardNative",
|
||||
},
|
||||
{
|
||||
name: "Antigravity账户路由到ForwardGemini",
|
||||
accountID: antigravityAcc.ID,
|
||||
expectedService: "AntigravityGatewayService.ForwardGemini",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
// 从数据库获取账户
|
||||
account, err := s.accountRepo.GetByID(s.ctx, tt.accountID)
|
||||
s.Require().NoError(err)
|
||||
|
||||
// 模拟 Handler 层的路由决策
|
||||
var routedService string
|
||||
if account.Platform == service.PlatformAntigravity {
|
||||
routedService = "AntigravityGatewayService.ForwardGemini"
|
||||
} else {
|
||||
routedService = "GeminiMessagesCompatService.ForwardNative"
|
||||
}
|
||||
|
||||
s.Require().Equal(tt.expectedService, routedService)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -82,8 +82,9 @@ func (s *GroupRepoSuite) TestList() {
|
||||
|
||||
groups, page, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10})
|
||||
s.Require().NoError(err, "List")
|
||||
s.Require().Len(groups, 2)
|
||||
s.Require().Equal(int64(2), page.Total)
|
||||
// 3 default groups + 2 test groups = 5 total
|
||||
s.Require().Len(groups, 5)
|
||||
s.Require().Equal(int64(5), page.Total)
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestListWithFilters_Platform() {
|
||||
@@ -92,8 +93,12 @@ func (s *GroupRepoSuite) TestListWithFilters_Platform() {
|
||||
|
||||
groups, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, service.PlatformOpenAI, "", nil)
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(groups, 1)
|
||||
s.Require().Equal(service.PlatformOpenAI, groups[0].Platform)
|
||||
// 1 default openai group + 1 test openai group = 2 total
|
||||
s.Require().Len(groups, 2)
|
||||
// Verify all groups are OpenAI platform
|
||||
for _, g := range groups {
|
||||
s.Require().Equal(service.PlatformOpenAI, g.Platform)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestListWithFilters_Status() {
|
||||
@@ -151,8 +156,17 @@ func (s *GroupRepoSuite) TestListActive() {
|
||||
|
||||
groups, err := s.repo.ListActive(s.ctx)
|
||||
s.Require().NoError(err, "ListActive")
|
||||
s.Require().Len(groups, 1)
|
||||
s.Require().Equal("active1", groups[0].Name)
|
||||
// 3 default groups (all active) + 1 test active group = 4 total
|
||||
s.Require().Len(groups, 4)
|
||||
// Verify our test group is in the results
|
||||
var found bool
|
||||
for _, g := range groups {
|
||||
if g.Name == "active1" {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
s.Require().True(found, "active1 group should be in results")
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestListActiveByPlatform() {
|
||||
@@ -162,8 +176,17 @@ func (s *GroupRepoSuite) TestListActiveByPlatform() {
|
||||
|
||||
groups, err := s.repo.ListActiveByPlatform(s.ctx, service.PlatformAnthropic)
|
||||
s.Require().NoError(err, "ListActiveByPlatform")
|
||||
s.Require().Len(groups, 1)
|
||||
s.Require().Equal("g1", groups[0].Name)
|
||||
// 1 default anthropic group + 1 test active anthropic group = 2 total
|
||||
s.Require().Len(groups, 2)
|
||||
// Verify our test group is in the results
|
||||
var found bool
|
||||
for _, g := range groups {
|
||||
if g.Name == "g1" {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
s.Require().True(found, "g1 group should be in results")
|
||||
}
|
||||
|
||||
// --- ExistsByName ---
|
||||
|
||||
@@ -93,7 +93,8 @@ func TestMain(m *testing.M) {
|
||||
log.Printf("failed to open gorm db: %v", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if err := AutoMigrate(integrationDB); err != nil {
|
||||
// 使用 simple 模式以便测试默认分组功能
|
||||
if err := AutoMigrate(integrationDB, "simple"); err != nil {
|
||||
log.Printf("failed to automigrate db: %v", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
@@ -119,6 +119,7 @@ func (r *proxyRepository) CountAccountsByProxyID(ctx context.Context, proxyID in
|
||||
var count int64
|
||||
err := r.db.WithContext(ctx).Table("accounts").
|
||||
Where("proxy_id = ?", proxyID).
|
||||
Where("deleted_at IS NULL").
|
||||
Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
@@ -134,6 +135,7 @@ func (r *proxyRepository) GetAccountCountsForProxies(ctx context.Context) (map[i
|
||||
Table("accounts").
|
||||
Select("proxy_id, COUNT(*) as count").
|
||||
Where("proxy_id IS NOT NULL").
|
||||
Where("deleted_at IS NULL").
|
||||
Group("proxy_id").
|
||||
Scan(&results).Error
|
||||
if err != nil {
|
||||
|
||||
@@ -129,51 +129,68 @@ type DashboardStats = usagestats.DashboardStats
|
||||
func (r *usageLogRepository) GetDashboardStats(ctx context.Context) (*DashboardStats, error) {
|
||||
var stats DashboardStats
|
||||
today := timezone.Today()
|
||||
now := time.Now()
|
||||
|
||||
// 总用户数
|
||||
r.db.WithContext(ctx).Model(&userModel{}).Count(&stats.TotalUsers)
|
||||
// 合并用户统计查询
|
||||
var userStats struct {
|
||||
TotalUsers int64 `gorm:"column:total_users"`
|
||||
TodayNewUsers int64 `gorm:"column:today_new_users"`
|
||||
ActiveUsers int64 `gorm:"column:active_users"`
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Raw(`
|
||||
SELECT
|
||||
COUNT(*) as total_users,
|
||||
COUNT(CASE WHEN created_at >= ? THEN 1 END) as today_new_users,
|
||||
(SELECT COUNT(DISTINCT user_id) FROM usage_logs WHERE created_at >= ?) as active_users
|
||||
FROM users
|
||||
`, today, today).Scan(&userStats).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stats.TotalUsers = userStats.TotalUsers
|
||||
stats.TodayNewUsers = userStats.TodayNewUsers
|
||||
stats.ActiveUsers = userStats.ActiveUsers
|
||||
|
||||
// 今日新增用户数
|
||||
r.db.WithContext(ctx).Model(&userModel{}).
|
||||
Where("created_at >= ?", today).
|
||||
Count(&stats.TodayNewUsers)
|
||||
// 合并API Key统计查询
|
||||
var apiKeyStats struct {
|
||||
TotalApiKeys int64 `gorm:"column:total_api_keys"`
|
||||
ActiveApiKeys int64 `gorm:"column:active_api_keys"`
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Raw(`
|
||||
SELECT
|
||||
COUNT(*) as total_api_keys,
|
||||
COUNT(CASE WHEN status = ? THEN 1 END) as active_api_keys
|
||||
FROM api_keys
|
||||
`, service.StatusActive).Scan(&apiKeyStats).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stats.TotalApiKeys = apiKeyStats.TotalApiKeys
|
||||
stats.ActiveApiKeys = apiKeyStats.ActiveApiKeys
|
||||
|
||||
// 今日活跃用户数 (今日有请求的用户)
|
||||
r.db.WithContext(ctx).Model(&usageLogModel{}).
|
||||
Distinct("user_id").
|
||||
Where("created_at >= ?", today).
|
||||
Count(&stats.ActiveUsers)
|
||||
|
||||
// 总 API Key 数
|
||||
r.db.WithContext(ctx).Model(&apiKeyModel{}).Count(&stats.TotalApiKeys)
|
||||
|
||||
// 活跃 API Key 数
|
||||
r.db.WithContext(ctx).Model(&apiKeyModel{}).
|
||||
Where("status = ?", service.StatusActive).
|
||||
Count(&stats.ActiveApiKeys)
|
||||
|
||||
// 总账户数
|
||||
r.db.WithContext(ctx).Model(&accountModel{}).Count(&stats.TotalAccounts)
|
||||
|
||||
// 正常账户数 (schedulable=true, status=active)
|
||||
r.db.WithContext(ctx).Model(&accountModel{}).
|
||||
Where("status = ? AND schedulable = ?", service.StatusActive, true).
|
||||
Count(&stats.NormalAccounts)
|
||||
|
||||
// 异常账户数 (status=error)
|
||||
r.db.WithContext(ctx).Model(&accountModel{}).
|
||||
Where("status = ?", service.StatusError).
|
||||
Count(&stats.ErrorAccounts)
|
||||
|
||||
// 限流账户数
|
||||
r.db.WithContext(ctx).Model(&accountModel{}).
|
||||
Where("rate_limited_at IS NOT NULL AND rate_limit_reset_at > ?", time.Now()).
|
||||
Count(&stats.RateLimitAccounts)
|
||||
|
||||
// 过载账户数
|
||||
r.db.WithContext(ctx).Model(&accountModel{}).
|
||||
Where("overload_until IS NOT NULL AND overload_until > ?", time.Now()).
|
||||
Count(&stats.OverloadAccounts)
|
||||
// 合并账户统计查询
|
||||
var accountStats struct {
|
||||
TotalAccounts int64 `gorm:"column:total_accounts"`
|
||||
NormalAccounts int64 `gorm:"column:normal_accounts"`
|
||||
ErrorAccounts int64 `gorm:"column:error_accounts"`
|
||||
RateLimitAccounts int64 `gorm:"column:ratelimit_accounts"`
|
||||
OverloadAccounts int64 `gorm:"column:overload_accounts"`
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Raw(`
|
||||
SELECT
|
||||
COUNT(*) as total_accounts,
|
||||
COUNT(CASE WHEN status = ? AND schedulable = true THEN 1 END) as normal_accounts,
|
||||
COUNT(CASE WHEN status = ? THEN 1 END) as error_accounts,
|
||||
COUNT(CASE WHEN rate_limited_at IS NOT NULL AND rate_limit_reset_at > ? THEN 1 END) as ratelimit_accounts,
|
||||
COUNT(CASE WHEN overload_until IS NOT NULL AND overload_until > ? THEN 1 END) as overload_accounts
|
||||
FROM accounts
|
||||
WHERE deleted_at IS NULL
|
||||
`, service.StatusActive, service.StatusError, now, now).Scan(&accountStats).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stats.TotalAccounts = accountStats.TotalAccounts
|
||||
stats.NormalAccounts = accountStats.NormalAccounts
|
||||
stats.ErrorAccounts = accountStats.ErrorAccounts
|
||||
stats.RateLimitAccounts = accountStats.RateLimitAccounts
|
||||
stats.OverloadAccounts = accountStats.OverloadAccounts
|
||||
|
||||
// 累计 Token 统计
|
||||
var totalStats struct {
|
||||
@@ -273,6 +290,88 @@ func (r *usageLogRepository) ListByUserAndTimeRange(ctx context.Context, userID
|
||||
return usageLogModelsToService(logs), nil, err
|
||||
}
|
||||
|
||||
// GetUserStatsAggregated returns aggregated usage statistics for a user using database-level aggregation
|
||||
func (r *usageLogRepository) GetUserStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||||
var stats struct {
|
||||
TotalRequests int64 `gorm:"column:total_requests"`
|
||||
TotalInputTokens int64 `gorm:"column:total_input_tokens"`
|
||||
TotalOutputTokens int64 `gorm:"column:total_output_tokens"`
|
||||
TotalCacheTokens int64 `gorm:"column:total_cache_tokens"`
|
||||
TotalCost float64 `gorm:"column:total_cost"`
|
||||
TotalActualCost float64 `gorm:"column:total_actual_cost"`
|
||||
AverageDurationMs float64 `gorm:"column:avg_duration_ms"`
|
||||
}
|
||||
|
||||
err := r.db.WithContext(ctx).Model(&usageLogModel{}).
|
||||
Select(`
|
||||
COUNT(*) as total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) as total_input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) as total_output_tokens,
|
||||
COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens,
|
||||
COALESCE(SUM(total_cost), 0) as total_cost,
|
||||
COALESCE(SUM(actual_cost), 0) as total_actual_cost,
|
||||
COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms
|
||||
`).
|
||||
Where("user_id = ? AND created_at >= ? AND created_at < ?", userID, startTime, endTime).
|
||||
Scan(&stats).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &usagestats.UsageStats{
|
||||
TotalRequests: stats.TotalRequests,
|
||||
TotalInputTokens: stats.TotalInputTokens,
|
||||
TotalOutputTokens: stats.TotalOutputTokens,
|
||||
TotalCacheTokens: stats.TotalCacheTokens,
|
||||
TotalTokens: stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens,
|
||||
TotalCost: stats.TotalCost,
|
||||
TotalActualCost: stats.TotalActualCost,
|
||||
AverageDurationMs: stats.AverageDurationMs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetApiKeyStatsAggregated returns aggregated usage statistics for an API key using database-level aggregation
|
||||
func (r *usageLogRepository) GetApiKeyStatsAggregated(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||||
var stats struct {
|
||||
TotalRequests int64 `gorm:"column:total_requests"`
|
||||
TotalInputTokens int64 `gorm:"column:total_input_tokens"`
|
||||
TotalOutputTokens int64 `gorm:"column:total_output_tokens"`
|
||||
TotalCacheTokens int64 `gorm:"column:total_cache_tokens"`
|
||||
TotalCost float64 `gorm:"column:total_cost"`
|
||||
TotalActualCost float64 `gorm:"column:total_actual_cost"`
|
||||
AverageDurationMs float64 `gorm:"column:avg_duration_ms"`
|
||||
}
|
||||
|
||||
err := r.db.WithContext(ctx).Model(&usageLogModel{}).
|
||||
Select(`
|
||||
COUNT(*) as total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) as total_input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) as total_output_tokens,
|
||||
COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens,
|
||||
COALESCE(SUM(total_cost), 0) as total_cost,
|
||||
COALESCE(SUM(actual_cost), 0) as total_actual_cost,
|
||||
COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms
|
||||
`).
|
||||
Where("api_key_id = ? AND created_at >= ? AND created_at < ?", apiKeyID, startTime, endTime).
|
||||
Scan(&stats).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &usagestats.UsageStats{
|
||||
TotalRequests: stats.TotalRequests,
|
||||
TotalInputTokens: stats.TotalInputTokens,
|
||||
TotalOutputTokens: stats.TotalOutputTokens,
|
||||
TotalCacheTokens: stats.TotalCacheTokens,
|
||||
TotalTokens: stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens,
|
||||
TotalCost: stats.TotalCost,
|
||||
TotalActualCost: stats.TotalActualCost,
|
||||
AverageDurationMs: stats.AverageDurationMs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *usageLogRepository) ListByApiKeyAndTimeRange(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||||
var logs []usageLogModel
|
||||
err := r.db.WithContext(ctx).
|
||||
@@ -631,6 +730,21 @@ func (r *usageLogRepository) ListWithFilters(ctx context.Context, params paginat
|
||||
if filters.ApiKeyID > 0 {
|
||||
db = db.Where("api_key_id = ?", filters.ApiKeyID)
|
||||
}
|
||||
if filters.AccountID > 0 {
|
||||
db = db.Where("account_id = ?", filters.AccountID)
|
||||
}
|
||||
if filters.GroupID > 0 {
|
||||
db = db.Where("group_id = ?", filters.GroupID)
|
||||
}
|
||||
if filters.Model != "" {
|
||||
db = db.Where("model = ?", filters.Model)
|
||||
}
|
||||
if filters.Stream != nil {
|
||||
db = db.Where("stream = ?", *filters.Stream)
|
||||
}
|
||||
if filters.BillingType != nil {
|
||||
db = db.Where("billing_type = ?", *filters.BillingType)
|
||||
}
|
||||
if filters.StartTime != nil {
|
||||
db = db.Where("created_at >= ?", *filters.StartTime)
|
||||
}
|
||||
@@ -642,8 +756,8 @@ func (r *usageLogRepository) ListWithFilters(ctx context.Context, params paginat
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
// Preload user and api_key for display
|
||||
if err := db.Preload("User").Preload("ApiKey").
|
||||
// Preload user, api_key, account, and group for display
|
||||
if err := db.Preload("User").Preload("ApiKey").Preload("Account").Preload("Group").
|
||||
Offset(params.Offset()).Limit(params.Limit()).
|
||||
Order("id DESC").Find(&logs).Error; err != nil {
|
||||
return nil, nil, err
|
||||
|
||||
@@ -198,6 +198,7 @@ type userModel struct {
|
||||
Concurrency int `gorm:"default:5;not null"`
|
||||
Status string `gorm:"size:20;default:active;not null"`
|
||||
AllowedGroups pq.Int64Array `gorm:"type:bigint[]"`
|
||||
TokenVersion int64 `gorm:"default:0;not null"` // Incremented on password change
|
||||
CreatedAt time.Time `gorm:"not null"`
|
||||
UpdatedAt time.Time `gorm:"not null"`
|
||||
DeletedAt gorm.DeletedAt `gorm:"index"`
|
||||
@@ -221,6 +222,7 @@ func userModelToService(m *userModel) *service.User {
|
||||
Concurrency: m.Concurrency,
|
||||
Status: m.Status,
|
||||
AllowedGroups: []int64(m.AllowedGroups),
|
||||
TokenVersion: m.TokenVersion,
|
||||
CreatedAt: m.CreatedAt,
|
||||
UpdatedAt: m.UpdatedAt,
|
||||
}
|
||||
@@ -242,6 +244,7 @@ func userModelFromService(u *service.User) *userModel {
|
||||
Concurrency: u.Concurrency,
|
||||
Status: u.Status,
|
||||
AllowedGroups: pq.Int64Array(u.AllowedGroups),
|
||||
TokenVersion: u.TokenVersion,
|
||||
CreatedAt: u.CreatedAt,
|
||||
UpdatedAt: u.UpdatedAt,
|
||||
}
|
||||
@@ -252,6 +255,7 @@ func applyUserModelToService(dst *service.User, src *userModel) {
|
||||
return
|
||||
}
|
||||
dst.ID = src.ID
|
||||
dst.TokenVersion = src.TokenVersion
|
||||
dst.CreatedAt = src.CreatedAt
|
||||
dst.UpdatedAt = src.UpdatedAt
|
||||
}
|
||||
|
||||
@@ -59,7 +59,8 @@ func TestAPIContracts(t *testing.T) {
|
||||
"status": "active",
|
||||
"allowed_groups": null,
|
||||
"created_at": "2025-01-02T03:04:05Z",
|
||||
"updated_at": "2025-01-02T03:04:05Z"
|
||||
"updated_at": "2025-01-02T03:04:05Z",
|
||||
"run_mode": "standard"
|
||||
}
|
||||
}`,
|
||||
},
|
||||
@@ -369,6 +370,7 @@ func newContractDeps(t *testing.T) *contractDeps {
|
||||
Default: config.DefaultConfig{
|
||||
ApiKeyPrefix: "sk-",
|
||||
},
|
||||
RunMode: config.RunModeStandard,
|
||||
}
|
||||
|
||||
userService := service.NewUserService(userRepo)
|
||||
@@ -380,7 +382,7 @@ func newContractDeps(t *testing.T) *contractDeps {
|
||||
settingRepo := newStubSettingRepo()
|
||||
settingService := service.NewSettingService(settingRepo, cfg)
|
||||
|
||||
authHandler := handler.NewAuthHandler(nil, userService)
|
||||
authHandler := handler.NewAuthHandler(cfg, nil, userService)
|
||||
apiKeyHandler := handler.NewAPIKeyHandler(apiKeyService)
|
||||
usageHandler := handler.NewUsageHandler(usageService, apiKeyService)
|
||||
adminSettingHandler := adminhandler.NewSettingHandler(settingService, nil)
|
||||
@@ -788,6 +790,25 @@ func (r *stubApiKeyRepo) ListByUserID(ctx context.Context, userID int64, params
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
|
||||
if len(apiKeyIDs) == 0 {
|
||||
return []int64{}, nil
|
||||
}
|
||||
seen := make(map[int64]struct{}, len(apiKeyIDs))
|
||||
out := make([]int64, 0, len(apiKeyIDs))
|
||||
for _, id := range apiKeyIDs {
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
key, ok := r.byID[id]
|
||||
if ok && key.UserID == userID {
|
||||
out = append(out, id)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) CountByUserID(ctx context.Context, userID int64) (int64, error) {
|
||||
var count int64
|
||||
for _, key := range r.byID {
|
||||
@@ -903,6 +924,55 @@ func (r *stubUsageLogRepo) GetUserUsageTrend(ctx context.Context, startTime, end
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUsageLogRepo) GetUserStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||||
logs := r.userLogs[userID]
|
||||
if len(logs) == 0 {
|
||||
return &usagestats.UsageStats{}, nil
|
||||
}
|
||||
|
||||
var totalRequests int64
|
||||
var totalInputTokens int64
|
||||
var totalOutputTokens int64
|
||||
var totalCacheTokens int64
|
||||
var totalCost float64
|
||||
var totalActualCost float64
|
||||
var totalDuration int64
|
||||
var durationCount int64
|
||||
|
||||
for _, log := range logs {
|
||||
totalRequests++
|
||||
totalInputTokens += int64(log.InputTokens)
|
||||
totalOutputTokens += int64(log.OutputTokens)
|
||||
totalCacheTokens += int64(log.CacheCreationTokens + log.CacheReadTokens)
|
||||
totalCost += log.TotalCost
|
||||
totalActualCost += log.ActualCost
|
||||
if log.DurationMs != nil {
|
||||
totalDuration += int64(*log.DurationMs)
|
||||
durationCount++
|
||||
}
|
||||
}
|
||||
|
||||
var avgDuration float64
|
||||
if durationCount > 0 {
|
||||
avgDuration = float64(totalDuration) / float64(durationCount)
|
||||
}
|
||||
|
||||
return &usagestats.UsageStats{
|
||||
TotalRequests: totalRequests,
|
||||
TotalInputTokens: totalInputTokens,
|
||||
TotalOutputTokens: totalOutputTokens,
|
||||
TotalCacheTokens: totalCacheTokens,
|
||||
TotalTokens: totalInputTokens + totalOutputTokens + totalCacheTokens,
|
||||
TotalCost: totalCost,
|
||||
TotalActualCost: totalActualCost,
|
||||
AverageDurationMs: avgDuration,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *stubUsageLogRepo) GetApiKeyStatsAggregated(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUsageLogRepo) GetBatchUserUsageStats(ctx context.Context, userIDs []int64) (map[int64]*usagestats.BatchUserUsageStats, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
@@ -924,7 +994,40 @@ func (r *stubUsageLogRepo) GetUserModelStats(ctx context.Context, userID int64,
|
||||
}
|
||||
|
||||
func (r *stubUsageLogRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters usagestats.UsageLogFilters) ([]service.UsageLog, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
logs := r.userLogs[filters.UserID]
|
||||
|
||||
// Apply filters
|
||||
var filtered []service.UsageLog
|
||||
for _, log := range logs {
|
||||
// Apply ApiKeyID filter
|
||||
if filters.ApiKeyID > 0 && log.ApiKeyID != filters.ApiKeyID {
|
||||
continue
|
||||
}
|
||||
// Apply Model filter
|
||||
if filters.Model != "" && log.Model != filters.Model {
|
||||
continue
|
||||
}
|
||||
// Apply Stream filter
|
||||
if filters.Stream != nil && log.Stream != *filters.Stream {
|
||||
continue
|
||||
}
|
||||
// Apply BillingType filter
|
||||
if filters.BillingType != nil && log.BillingType != *filters.BillingType {
|
||||
continue
|
||||
}
|
||||
// Apply time range filters
|
||||
if filters.StartTime != nil && log.CreatedAt.Before(*filters.StartTime) {
|
||||
continue
|
||||
}
|
||||
if filters.EndTime != nil && log.CreatedAt.After(*filters.EndTime) {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, log)
|
||||
}
|
||||
|
||||
total := int64(len(filtered))
|
||||
out := paginateLogs(filtered, params)
|
||||
return out, paginationResult(total, params), nil
|
||||
}
|
||||
|
||||
func (r *stubUsageLogRepo) GetGlobalStats(ctx context.Context, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
|
||||
|
||||
@@ -36,7 +36,7 @@ func ProvideRouter(
|
||||
r := gin.New()
|
||||
r.Use(middleware2.Recovery())
|
||||
|
||||
return SetupRouter(r, handlers, jwtAuth, adminAuth, apiKeyAuth, apiKeyService, subscriptionService)
|
||||
return SetupRouter(r, handlers, jwtAuth, adminAuth, apiKeyAuth, apiKeyService, subscriptionService, cfg)
|
||||
}
|
||||
|
||||
// ProvideHTTPServer 提供 HTTP 服务器
|
||||
|
||||
@@ -5,18 +5,19 @@ import (
|
||||
"log"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// NewApiKeyAuthMiddleware 创建 API Key 认证中间件
|
||||
func NewApiKeyAuthMiddleware(apiKeyService *service.ApiKeyService, subscriptionService *service.SubscriptionService) ApiKeyAuthMiddleware {
|
||||
return ApiKeyAuthMiddleware(apiKeyAuthWithSubscription(apiKeyService, subscriptionService))
|
||||
func NewApiKeyAuthMiddleware(apiKeyService *service.ApiKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) ApiKeyAuthMiddleware {
|
||||
return ApiKeyAuthMiddleware(apiKeyAuthWithSubscription(apiKeyService, subscriptionService, cfg))
|
||||
}
|
||||
|
||||
// apiKeyAuthWithSubscription API Key认证中间件(支持订阅验证)
|
||||
func apiKeyAuthWithSubscription(apiKeyService *service.ApiKeyService, subscriptionService *service.SubscriptionService) gin.HandlerFunc {
|
||||
func apiKeyAuthWithSubscription(apiKeyService *service.ApiKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// 尝试从Authorization header中提取API key (Bearer scheme)
|
||||
authHeader := c.GetHeader("Authorization")
|
||||
@@ -85,6 +86,18 @@ func apiKeyAuthWithSubscription(apiKeyService *service.ApiKeyService, subscripti
|
||||
return
|
||||
}
|
||||
|
||||
if cfg.RunMode == config.RunModeSimple {
|
||||
// 简易模式:跳过余额和订阅检查,但仍需设置必要的上下文
|
||||
c.Set(string(ContextKeyApiKey), apiKey)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: apiKey.User.ID,
|
||||
Concurrency: apiKey.User.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// 判断计费方式:订阅模式 vs 余额模式
|
||||
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
|
||||
|
||||
|
||||
@@ -4,23 +4,23 @@ import (
|
||||
"errors"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ApiKeyAuthGoogle is a Google-style error wrapper for API key auth.
|
||||
func ApiKeyAuthGoogle(apiKeyService *service.ApiKeyService) gin.HandlerFunc {
|
||||
return ApiKeyAuthWithSubscriptionGoogle(apiKeyService, nil)
|
||||
func ApiKeyAuthGoogle(apiKeyService *service.ApiKeyService, cfg *config.Config) gin.HandlerFunc {
|
||||
return ApiKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg)
|
||||
}
|
||||
|
||||
// ApiKeyAuthWithSubscriptionGoogle behaves like ApiKeyAuthWithSubscription but returns Google-style errors:
|
||||
// {"error":{"code":401,"message":"...","status":"UNAUTHENTICATED"}}
|
||||
//
|
||||
// It is intended for Gemini native endpoints (/v1beta) to match Gemini SDK expectations.
|
||||
func ApiKeyAuthWithSubscriptionGoogle(apiKeyService *service.ApiKeyService, subscriptionService *service.SubscriptionService) gin.HandlerFunc {
|
||||
func ApiKeyAuthWithSubscriptionGoogle(apiKeyService *service.ApiKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
apiKeyString := extractAPIKeyFromRequest(c)
|
||||
if apiKeyString == "" {
|
||||
@@ -30,7 +30,7 @@ func ApiKeyAuthWithSubscriptionGoogle(apiKeyService *service.ApiKeyService, subs
|
||||
|
||||
apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
if errors.Is(err, service.ErrApiKeyNotFound) {
|
||||
abortWithGoogleError(c, 401, "Invalid API key")
|
||||
return
|
||||
}
|
||||
@@ -51,6 +51,18 @@ func ApiKeyAuthWithSubscriptionGoogle(apiKeyService *service.ApiKeyService, subs
|
||||
return
|
||||
}
|
||||
|
||||
// 简易模式:跳过余额和订阅检查
|
||||
if cfg.RunMode == config.RunModeSimple {
|
||||
c.Set(string(ContextKeyApiKey), apiKey)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: apiKey.User.ID,
|
||||
Concurrency: apiKey.User.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
|
||||
if isSubscriptionType && subscriptionService != nil {
|
||||
subscription, err := subscriptionService.GetActiveSubscription(
|
||||
@@ -81,7 +93,11 @@ func ApiKeyAuthWithSubscriptionGoogle(apiKeyService *service.ApiKeyService, subs
|
||||
}
|
||||
|
||||
c.Set(string(ContextKeyApiKey), apiKey)
|
||||
c.Set(string(ContextKeyUser), apiKey.User)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: apiKey.User.ID,
|
||||
Concurrency: apiKey.User.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
286
backend/internal/server/middleware/api_key_auth_test.go
Normal file
286
backend/internal/server/middleware/api_key_auth_test.go
Normal file
@@ -0,0 +1,286 @@
|
||||
//go:build unit
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSimpleModeBypassesQuotaCheck(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
limit := 1.0
|
||||
group := &service.Group{
|
||||
ID: 42,
|
||||
Name: "sub",
|
||||
Status: service.StatusActive,
|
||||
SubscriptionType: service.SubscriptionTypeSubscription,
|
||||
DailyLimitUSD: &limit,
|
||||
}
|
||||
user := &service.User{
|
||||
ID: 7,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.ApiKey{
|
||||
ID: 100,
|
||||
UserID: user.ID,
|
||||
Key: "test-key",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
Group: group,
|
||||
}
|
||||
apiKey.GroupID = &group.ID
|
||||
|
||||
apiKeyRepo := &stubApiKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.ApiKey, error) {
|
||||
if key != apiKey.Key {
|
||||
return nil, service.ErrApiKeyNotFound
|
||||
}
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
}
|
||||
|
||||
t.Run("simple_mode_bypasses_quota_check", func(t *testing.T) {
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
apiKeyService := service.NewApiKeyService(apiKeyRepo, nil, nil, nil, nil, cfg)
|
||||
subscriptionService := service.NewSubscriptionService(nil, &stubUserSubscriptionRepo{}, nil)
|
||||
router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set("x-api-key", apiKey.Key)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
})
|
||||
|
||||
t.Run("standard_mode_enforces_quota_check", func(t *testing.T) {
|
||||
cfg := &config.Config{RunMode: config.RunModeStandard}
|
||||
apiKeyService := service.NewApiKeyService(apiKeyRepo, nil, nil, nil, nil, cfg)
|
||||
|
||||
now := time.Now()
|
||||
sub := &service.UserSubscription{
|
||||
ID: 55,
|
||||
UserID: user.ID,
|
||||
GroupID: group.ID,
|
||||
Status: service.SubscriptionStatusActive,
|
||||
ExpiresAt: now.Add(24 * time.Hour),
|
||||
DailyWindowStart: &now,
|
||||
DailyUsageUSD: 10,
|
||||
}
|
||||
subscriptionRepo := &stubUserSubscriptionRepo{
|
||||
getActive: func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
||||
if userID != sub.UserID || groupID != sub.GroupID {
|
||||
return nil, service.ErrSubscriptionNotFound
|
||||
}
|
||||
clone := *sub
|
||||
return &clone, nil
|
||||
},
|
||||
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
|
||||
activateWindow: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
resetDaily: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
}
|
||||
subscriptionService := service.NewSubscriptionService(nil, subscriptionRepo, nil)
|
||||
router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set("x-api-key", apiKey.Key)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, w.Code)
|
||||
require.Contains(t, w.Body.String(), "USAGE_LIMIT_EXCEEDED")
|
||||
})
|
||||
}
|
||||
|
||||
func newAuthTestRouter(apiKeyService *service.ApiKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) *gin.Engine {
|
||||
router := gin.New()
|
||||
router.Use(gin.HandlerFunc(NewApiKeyAuthMiddleware(apiKeyService, subscriptionService, cfg)))
|
||||
router.GET("/t", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
return router
|
||||
}
|
||||
|
||||
type stubApiKeyRepo struct {
|
||||
getByKey func(ctx context.Context, key string) (*service.ApiKey, error)
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) Create(ctx context.Context, key *service.ApiKey) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) GetByID(ctx context.Context, id int64) (*service.ApiKey, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) GetByKey(ctx context.Context, key string) (*service.ApiKey, error) {
|
||||
if r.getByKey != nil {
|
||||
return r.getByKey(ctx, key)
|
||||
}
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) Update(ctx context.Context, key *service.ApiKey) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) Delete(ctx context.Context, id int64) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams) ([]service.ApiKey, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) CountByUserID(ctx context.Context, userID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) ExistsByKey(ctx context.Context, key string) (bool, error) {
|
||||
return false, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.ApiKey, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) SearchApiKeys(ctx context.Context, userID int64, keyword string, limit int) ([]service.ApiKey, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
type stubUserSubscriptionRepo struct {
|
||||
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
|
||||
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
|
||||
activateWindow func(ctx context.Context, id int64, start time.Time) error
|
||||
resetDaily func(ctx context.Context, id int64, start time.Time) error
|
||||
resetWeekly func(ctx context.Context, id int64, start time.Time) error
|
||||
resetMonthly func(ctx context.Context, id int64, start time.Time) error
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) Create(ctx context.Context, sub *service.UserSubscription) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) GetByID(ctx context.Context, id int64) (*service.UserSubscription, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) GetByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) GetActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
||||
if r.getActive != nil {
|
||||
return r.getActive(ctx, userID, groupID)
|
||||
}
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) Update(ctx context.Context, sub *service.UserSubscription) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) Delete(ctx context.Context, id int64) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ListByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ListActiveByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ExistsByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
|
||||
return false, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ExtendExpiry(ctx context.Context, subscriptionID int64, newExpiresAt time.Time) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) UpdateStatus(ctx context.Context, subscriptionID int64, status string) error {
|
||||
if r.updateStatus != nil {
|
||||
return r.updateStatus(ctx, subscriptionID, status)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID int64, notes string) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, start time.Time) error {
|
||||
if r.activateWindow != nil {
|
||||
return r.activateWindow(ctx, id, start)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
|
||||
if r.resetDaily != nil {
|
||||
return r.resetDaily(ctx, id, newWindowStart)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
|
||||
if r.resetWeekly != nil {
|
||||
return r.resetWeekly(ctx, id, newWindowStart)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
|
||||
if r.resetMonthly != nil {
|
||||
return r.resetMonthly(ctx, id, newWindowStart)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) IncrementUsage(ctx context.Context, id int64, costUSD float64) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) BatchUpdateExpiredStatus(ctx context.Context) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
@@ -61,6 +61,13 @@ func jwtAuth(authService *service.AuthService, userService *service.UserService)
|
||||
return
|
||||
}
|
||||
|
||||
// Security: Validate TokenVersion to ensure token hasn't been invalidated
|
||||
// This check ensures tokens issued before a password change are rejected
|
||||
if claims.TokenVersion != user.TokenVersion {
|
||||
AbortWithError(c, 401, "TOKEN_REVOKED", "Token has been revoked (password changed)")
|
||||
return
|
||||
}
|
||||
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: user.ID,
|
||||
Concurrency: user.Concurrency,
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
package middleware
|
||||
|
||||
import "github.com/gin-gonic/gin"
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// ContextKey 定义上下文键类型
|
||||
type ContextKey string
|
||||
@@ -14,8 +19,39 @@ const (
|
||||
ContextKeyApiKey ContextKey = "api_key"
|
||||
// ContextKeySubscription 订阅上下文键
|
||||
ContextKeySubscription ContextKey = "subscription"
|
||||
// ContextKeyForcePlatform 强制平台(用于 /antigravity 路由)
|
||||
ContextKeyForcePlatform ContextKey = "force_platform"
|
||||
)
|
||||
|
||||
// ForcePlatform 返回设置强制平台的中间件
|
||||
// 同时设置 request.Context(供 Service 使用)和 gin.Context(供 Handler 快速检查)
|
||||
func ForcePlatform(platform string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// 设置到 request.Context,使用 ctxkey.ForcePlatform 供 Service 层读取
|
||||
ctx := context.WithValue(c.Request.Context(), ctxkey.ForcePlatform, platform)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
// 同时设置到 gin.Context,供 Handler 快速检查
|
||||
c.Set(string(ContextKeyForcePlatform), platform)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// HasForcePlatform 检查是否有强制平台(用于 Handler 跳过分组检查)
|
||||
func HasForcePlatform(c *gin.Context) bool {
|
||||
_, exists := c.Get(string(ContextKeyForcePlatform))
|
||||
return exists
|
||||
}
|
||||
|
||||
// GetForcePlatformFromContext 从 gin.Context 获取强制平台
|
||||
func GetForcePlatformFromContext(c *gin.Context) (string, bool) {
|
||||
value, exists := c.Get(string(ContextKeyForcePlatform))
|
||||
if !exists {
|
||||
return "", false
|
||||
}
|
||||
platform, ok := value.(string)
|
||||
return platform, ok
|
||||
}
|
||||
|
||||
// ErrorResponse 标准错误响应结构
|
||||
type ErrorResponse struct {
|
||||
Code string `json:"code"`
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/routes"
|
||||
@@ -19,6 +20,7 @@ func SetupRouter(
|
||||
apiKeyAuth middleware2.ApiKeyAuthMiddleware,
|
||||
apiKeyService *service.ApiKeyService,
|
||||
subscriptionService *service.SubscriptionService,
|
||||
cfg *config.Config,
|
||||
) *gin.Engine {
|
||||
// 应用中间件
|
||||
r.Use(middleware2.Logger())
|
||||
@@ -30,7 +32,7 @@ func SetupRouter(
|
||||
}
|
||||
|
||||
// 注册路由
|
||||
registerRoutes(r, handlers, jwtAuth, adminAuth, apiKeyAuth, apiKeyService, subscriptionService)
|
||||
registerRoutes(r, handlers, jwtAuth, adminAuth, apiKeyAuth, apiKeyService, subscriptionService, cfg)
|
||||
|
||||
return r
|
||||
}
|
||||
@@ -44,6 +46,7 @@ func registerRoutes(
|
||||
apiKeyAuth middleware2.ApiKeyAuthMiddleware,
|
||||
apiKeyService *service.ApiKeyService,
|
||||
subscriptionService *service.SubscriptionService,
|
||||
cfg *config.Config,
|
||||
) {
|
||||
// 通用路由(健康检查、状态等)
|
||||
routes.RegisterCommonRoutes(r)
|
||||
@@ -55,5 +58,5 @@ func registerRoutes(
|
||||
routes.RegisterAuthRoutes(v1, h, jwtAuth)
|
||||
routes.RegisterUserRoutes(v1, h, jwtAuth)
|
||||
routes.RegisterAdminRoutes(v1, h, adminAuth)
|
||||
routes.RegisterGatewayRoutes(r, h, apiKeyAuth, apiKeyService, subscriptionService)
|
||||
routes.RegisterGatewayRoutes(r, h, apiKeyAuth, apiKeyService, subscriptionService, cfg)
|
||||
}
|
||||
|
||||
@@ -34,6 +34,9 @@ func RegisterAdminRoutes(
|
||||
// Gemini OAuth
|
||||
registerGeminiOAuthRoutes(admin, h)
|
||||
|
||||
// Antigravity OAuth
|
||||
registerAntigravityOAuthRoutes(admin, h)
|
||||
|
||||
// 代理管理
|
||||
registerProxyRoutes(admin, h)
|
||||
|
||||
@@ -148,6 +151,14 @@ func registerGeminiOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
}
|
||||
}
|
||||
|
||||
func registerAntigravityOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
antigravity := admin.Group("/antigravity")
|
||||
{
|
||||
antigravity.POST("/oauth/auth-url", h.Admin.AntigravityOAuth.GenerateAuthURL)
|
||||
antigravity.POST("/oauth/exchange-code", h.Admin.AntigravityOAuth.ExchangeCode)
|
||||
}
|
||||
}
|
||||
|
||||
func registerProxyRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
proxies := admin.Group("/proxies")
|
||||
{
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -15,6 +16,7 @@ func RegisterGatewayRoutes(
|
||||
apiKeyAuth middleware.ApiKeyAuthMiddleware,
|
||||
apiKeyService *service.ApiKeyService,
|
||||
subscriptionService *service.SubscriptionService,
|
||||
cfg *config.Config,
|
||||
) {
|
||||
// API网关(Claude API兼容)
|
||||
gateway := r.Group("/v1")
|
||||
@@ -30,7 +32,7 @@ func RegisterGatewayRoutes(
|
||||
|
||||
// Gemini 原生 API 兼容层(Gemini SDK/CLI 直连)
|
||||
gemini := r.Group("/v1beta")
|
||||
gemini.Use(middleware.ApiKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService))
|
||||
gemini.Use(middleware.ApiKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg))
|
||||
{
|
||||
gemini.GET("/models", h.Gateway.GeminiV1BetaListModels)
|
||||
gemini.GET("/models/:model", h.Gateway.GeminiV1BetaGetModel)
|
||||
@@ -40,4 +42,24 @@ func RegisterGatewayRoutes(
|
||||
|
||||
// OpenAI Responses API(不带v1前缀的别名)
|
||||
r.POST("/responses", gin.HandlerFunc(apiKeyAuth), h.OpenAIGateway.Responses)
|
||||
|
||||
// Antigravity 专用路由(仅使用 antigravity 账户,不混合调度)
|
||||
antigravityV1 := r.Group("/antigravity/v1")
|
||||
antigravityV1.Use(middleware.ForcePlatform(service.PlatformAntigravity))
|
||||
antigravityV1.Use(gin.HandlerFunc(apiKeyAuth))
|
||||
{
|
||||
antigravityV1.POST("/messages", h.Gateway.Messages)
|
||||
antigravityV1.POST("/messages/count_tokens", h.Gateway.CountTokens)
|
||||
antigravityV1.GET("/models", h.Gateway.Models)
|
||||
antigravityV1.GET("/usage", h.Gateway.Usage)
|
||||
}
|
||||
|
||||
antigravityV1Beta := r.Group("/antigravity/v1beta")
|
||||
antigravityV1Beta.Use(middleware.ForcePlatform(service.PlatformAntigravity))
|
||||
antigravityV1Beta.Use(middleware.ApiKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg))
|
||||
{
|
||||
antigravityV1Beta.GET("/models", h.Gateway.GeminiV1BetaListModels)
|
||||
antigravityV1Beta.GET("/models/:model", h.Gateway.GeminiV1BetaGetModel)
|
||||
antigravityV1Beta.POST("/models/*modelAction", h.Gateway.GeminiV1BetaModels)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
package service
|
||||
|
||||
import "time"
|
||||
import (
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Account struct {
|
||||
ID int64
|
||||
@@ -82,12 +86,28 @@ func (a *Account) GetCredential(key string) string {
|
||||
if a.Credentials == nil {
|
||||
return ""
|
||||
}
|
||||
if v, ok := a.Credentials[key]; ok {
|
||||
if s, ok := v.(string); ok {
|
||||
return s
|
||||
}
|
||||
v, ok := a.Credentials[key]
|
||||
if !ok || v == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
// 支持多种类型(兼容历史数据中 expires_at 等字段可能是数字或字符串)
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
return val
|
||||
case json.Number:
|
||||
// GORM datatypes.JSONMap 使用 UseNumber() 解析,数字类型为 json.Number
|
||||
return val.String()
|
||||
case float64:
|
||||
// JSON 解析后数字默认为 float64
|
||||
return strconv.FormatInt(int64(val), 10)
|
||||
case int64:
|
||||
return strconv.FormatInt(val, 10)
|
||||
case int:
|
||||
return strconv.Itoa(val)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (a *Account) GetModelMapping() map[string]string {
|
||||
@@ -326,3 +346,20 @@ func (a *Account) IsOpenAITokenExpired() bool {
|
||||
}
|
||||
return time.Now().Add(60 * time.Second).After(*expiresAt)
|
||||
}
|
||||
|
||||
// IsMixedSchedulingEnabled 检查 antigravity 账户是否启用混合调度
|
||||
// 启用后可参与 anthropic/gemini 分组的账户调度
|
||||
func (a *Account) IsMixedSchedulingEnabled() bool {
|
||||
if a.Platform != PlatformAntigravity {
|
||||
return false
|
||||
}
|
||||
if a.Extra == nil {
|
||||
return false
|
||||
}
|
||||
if v, ok := a.Extra["mixed_scheduling"]; ok {
|
||||
if enabled, ok := v.(bool); ok {
|
||||
return enabled
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -29,6 +29,7 @@ type AccountRepository interface {
|
||||
ListByPlatform(ctx context.Context, platform string) ([]Account, error)
|
||||
|
||||
UpdateLastUsed(ctx context.Context, id int64) error
|
||||
BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error
|
||||
SetError(ctx context.Context, id int64, errorMsg string) error
|
||||
SetSchedulable(ctx context.Context, id int64, schedulable bool) error
|
||||
BindGroups(ctx context.Context, accountID int64, groupIDs []int64) error
|
||||
@@ -37,6 +38,8 @@ type AccountRepository interface {
|
||||
ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error)
|
||||
ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error)
|
||||
ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error)
|
||||
ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]Account, error)
|
||||
ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error)
|
||||
|
||||
SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error
|
||||
SetOverloaded(ctx context.Context, id int64, until time.Time) error
|
||||
@@ -208,20 +211,23 @@ func (s *AccountService) Update(ctx context.Context, id int64, req UpdateAccount
|
||||
account.Status = *req.Status
|
||||
}
|
||||
|
||||
if err := s.accountRepo.Update(ctx, account); err != nil {
|
||||
return nil, fmt.Errorf("update account: %w", err)
|
||||
}
|
||||
|
||||
// 更新分组绑定
|
||||
// 先验证分组是否存在(在任何写操作之前)
|
||||
if req.GroupIDs != nil {
|
||||
// 验证分组是否存在
|
||||
for _, groupID := range *req.GroupIDs {
|
||||
_, err := s.groupRepo.GetByID(ctx, groupID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get group: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 执行更新
|
||||
if err := s.accountRepo.Update(ctx, account); err != nil {
|
||||
return nil, fmt.Errorf("update account: %w", err)
|
||||
}
|
||||
|
||||
// 绑定分组
|
||||
if req.GroupIDs != nil {
|
||||
if err := s.accountRepo.BindGroups(ctx, account.ID, *req.GroupIDs); err != nil {
|
||||
return nil, fmt.Errorf("bind groups: %w", err)
|
||||
}
|
||||
|
||||
@@ -48,17 +48,29 @@ type UsageLogRepository interface {
|
||||
|
||||
// Account stats
|
||||
GetAccountUsageStats(ctx context.Context, accountID int64, startTime, endTime time.Time) (*usagestats.AccountUsageStatsResponse, error)
|
||||
|
||||
// Aggregated stats (optimized)
|
||||
GetUserStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error)
|
||||
GetApiKeyStatsAggregated(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error)
|
||||
}
|
||||
|
||||
// usageCache 用于缓存usage数据
|
||||
type usageCache struct {
|
||||
data *UsageInfo
|
||||
// apiUsageCache 缓存从 Anthropic API 获取的使用率数据(utilization, resets_at)
|
||||
type apiUsageCache struct {
|
||||
response *ClaudeUsageResponse
|
||||
timestamp time.Time
|
||||
}
|
||||
|
||||
// windowStatsCache 缓存从本地数据库查询的窗口统计(requests, tokens, cost)
|
||||
type windowStatsCache struct {
|
||||
stats *WindowStats
|
||||
timestamp time.Time
|
||||
}
|
||||
|
||||
var (
|
||||
usageCacheMap = sync.Map{}
|
||||
cacheTTL = 10 * time.Minute
|
||||
apiCacheMap = sync.Map{} // 缓存 API 响应
|
||||
windowStatsCacheMap = sync.Map{} // 缓存窗口统计
|
||||
apiCacheTTL = 10 * time.Minute
|
||||
windowStatsCacheTTL = 1 * time.Minute
|
||||
)
|
||||
|
||||
// WindowStats 窗口期统计
|
||||
@@ -122,7 +134,7 @@ func NewAccountUsageService(accountRepo AccountRepository, usageLogRepo UsageLog
|
||||
}
|
||||
|
||||
// GetUsage 获取账号使用量
|
||||
// OAuth账号: 调用Anthropic API获取真实数据(需要profile scope),缓存10分钟
|
||||
// OAuth账号: 调用Anthropic API获取真实数据(需要profile scope),API响应缓存10分钟,窗口统计缓存1分钟
|
||||
// Setup Token账号: 根据session_window推算5h窗口,7d数据不可用(没有profile scope)
|
||||
// API Key账号: 不支持usage查询
|
||||
func (s *AccountUsageService) GetUsage(ctx context.Context, accountID int64) (*UsageInfo, error) {
|
||||
@@ -133,30 +145,34 @@ func (s *AccountUsageService) GetUsage(ctx context.Context, accountID int64) (*U
|
||||
|
||||
// 只有oauth类型账号可以通过API获取usage(有profile scope)
|
||||
if account.CanGetUsage() {
|
||||
// 检查缓存
|
||||
if cached, ok := usageCacheMap.Load(accountID); ok {
|
||||
cache, ok := cached.(*usageCache)
|
||||
if !ok {
|
||||
usageCacheMap.Delete(accountID)
|
||||
} else if time.Since(cache.timestamp) < cacheTTL {
|
||||
return cache.data, nil
|
||||
var apiResp *ClaudeUsageResponse
|
||||
|
||||
// 1. 检查 API 缓存(10 分钟)
|
||||
if cached, ok := apiCacheMap.Load(accountID); ok {
|
||||
if cache, ok := cached.(*apiUsageCache); ok && time.Since(cache.timestamp) < apiCacheTTL {
|
||||
apiResp = cache.response
|
||||
}
|
||||
}
|
||||
|
||||
// 从API获取数据
|
||||
usage, err := s.fetchOAuthUsage(ctx, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
// 2. 如果没有缓存,从 API 获取
|
||||
if apiResp == nil {
|
||||
apiResp, err = s.fetchOAuthUsageRaw(ctx, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 缓存 API 响应
|
||||
apiCacheMap.Store(accountID, &apiUsageCache{
|
||||
response: apiResp,
|
||||
timestamp: time.Now(),
|
||||
})
|
||||
}
|
||||
|
||||
// 添加5h窗口统计数据
|
||||
s.addWindowStats(ctx, account, usage)
|
||||
// 3. 构建 UsageInfo(每次都重新计算 RemainingSeconds)
|
||||
now := time.Now()
|
||||
usage := s.buildUsageInfo(apiResp, &now)
|
||||
|
||||
// 缓存结果
|
||||
usageCacheMap.Store(accountID, &usageCache{
|
||||
data: usage,
|
||||
timestamp: time.Now(),
|
||||
})
|
||||
// 4. 添加窗口统计(有独立缓存,1 分钟)
|
||||
s.addWindowStats(ctx, account, usage)
|
||||
|
||||
return usage, nil
|
||||
}
|
||||
@@ -173,31 +189,54 @@ func (s *AccountUsageService) GetUsage(ctx context.Context, accountID int64) (*U
|
||||
return nil, fmt.Errorf("account type %s does not support usage query", account.Type)
|
||||
}
|
||||
|
||||
// addWindowStats 为usage数据添加窗口期统计
|
||||
// addWindowStats 为 usage 数据添加窗口期统计
|
||||
// 使用独立缓存(1 分钟),与 API 缓存分离
|
||||
func (s *AccountUsageService) addWindowStats(ctx context.Context, account *Account, usage *UsageInfo) {
|
||||
if usage.FiveHour == nil {
|
||||
// 修复:即使 FiveHour 为 nil,也要尝试获取统计数据
|
||||
// 因为 SevenDay/SevenDaySonnet 可能需要
|
||||
if usage.FiveHour == nil && usage.SevenDay == nil && usage.SevenDaySonnet == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// 使用session_window_start作为统计起始时间
|
||||
var startTime time.Time
|
||||
if account.SessionWindowStart != nil {
|
||||
startTime = *account.SessionWindowStart
|
||||
} else {
|
||||
// 如果没有窗口信息,使用5小时前作为默认
|
||||
startTime = time.Now().Add(-5 * time.Hour)
|
||||
// 检查窗口统计缓存(1 分钟)
|
||||
var windowStats *WindowStats
|
||||
if cached, ok := windowStatsCacheMap.Load(account.ID); ok {
|
||||
if cache, ok := cached.(*windowStatsCache); ok && time.Since(cache.timestamp) < windowStatsCacheTTL {
|
||||
windowStats = cache.stats
|
||||
}
|
||||
}
|
||||
|
||||
stats, err := s.usageLogRepo.GetAccountWindowStats(ctx, account.ID, startTime)
|
||||
if err != nil {
|
||||
log.Printf("Failed to get window stats for account %d: %v", account.ID, err)
|
||||
return
|
||||
// 如果没有缓存,从数据库查询
|
||||
if windowStats == nil {
|
||||
var startTime time.Time
|
||||
if account.SessionWindowStart != nil {
|
||||
startTime = *account.SessionWindowStart
|
||||
} else {
|
||||
startTime = time.Now().Add(-5 * time.Hour)
|
||||
}
|
||||
|
||||
stats, err := s.usageLogRepo.GetAccountWindowStats(ctx, account.ID, startTime)
|
||||
if err != nil {
|
||||
log.Printf("Failed to get window stats for account %d: %v", account.ID, err)
|
||||
return
|
||||
}
|
||||
|
||||
windowStats = &WindowStats{
|
||||
Requests: stats.Requests,
|
||||
Tokens: stats.Tokens,
|
||||
Cost: stats.Cost,
|
||||
}
|
||||
|
||||
// 缓存窗口统计(1 分钟)
|
||||
windowStatsCacheMap.Store(account.ID, &windowStatsCache{
|
||||
stats: windowStats,
|
||||
timestamp: time.Now(),
|
||||
})
|
||||
}
|
||||
|
||||
usage.FiveHour.WindowStats = &WindowStats{
|
||||
Requests: stats.Requests,
|
||||
Tokens: stats.Tokens,
|
||||
Cost: stats.Cost,
|
||||
// 为 FiveHour 添加 WindowStats(5h 窗口统计)
|
||||
if usage.FiveHour != nil {
|
||||
usage.FiveHour.WindowStats = windowStats
|
||||
}
|
||||
}
|
||||
|
||||
@@ -223,8 +262,8 @@ func (s *AccountUsageService) GetAccountUsageStats(ctx context.Context, accountI
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
// fetchOAuthUsage 从Anthropic API获取OAuth账号的使用量
|
||||
func (s *AccountUsageService) fetchOAuthUsage(ctx context.Context, account *Account) (*UsageInfo, error) {
|
||||
// fetchOAuthUsageRaw 从 Anthropic API 获取原始响应(不构建 UsageInfo)
|
||||
func (s *AccountUsageService) fetchOAuthUsageRaw(ctx context.Context, account *Account) (*ClaudeUsageResponse, error) {
|
||||
accessToken := account.GetCredential("access_token")
|
||||
if accessToken == "" {
|
||||
return nil, fmt.Errorf("no access token available")
|
||||
@@ -235,13 +274,7 @@ func (s *AccountUsageService) fetchOAuthUsage(ctx context.Context, account *Acco
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
usageResp, err := s.usageFetcher.FetchUsage(ctx, accessToken, proxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
return s.buildUsageInfo(usageResp, &now), nil
|
||||
return s.usageFetcher.FetchUsage(ctx, accessToken, proxyURL)
|
||||
}
|
||||
|
||||
// parseTime 尝试多种格式解析时间
|
||||
@@ -266,20 +299,16 @@ func (s *AccountUsageService) buildUsageInfo(resp *ClaudeUsageResponse, updatedA
|
||||
UpdatedAt: updatedAt,
|
||||
}
|
||||
|
||||
// 5小时窗口
|
||||
// 5小时窗口 - 始终创建对象(即使 ResetsAt 为空)
|
||||
info.FiveHour = &UsageProgress{
|
||||
Utilization: resp.FiveHour.Utilization,
|
||||
}
|
||||
if resp.FiveHour.ResetsAt != "" {
|
||||
if fiveHourReset, err := parseTime(resp.FiveHour.ResetsAt); err == nil {
|
||||
info.FiveHour = &UsageProgress{
|
||||
Utilization: resp.FiveHour.Utilization,
|
||||
ResetsAt: &fiveHourReset,
|
||||
RemainingSeconds: int(time.Until(fiveHourReset).Seconds()),
|
||||
}
|
||||
info.FiveHour.ResetsAt = &fiveHourReset
|
||||
info.FiveHour.RemainingSeconds = int(time.Until(fiveHourReset).Seconds())
|
||||
} else {
|
||||
log.Printf("Failed to parse FiveHour.ResetsAt: %s, error: %v", resp.FiveHour.ResetsAt, err)
|
||||
// 即使解析失败也返回utilization
|
||||
info.FiveHour = &UsageProgress{
|
||||
Utilization: resp.FiveHour.Utilization,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -609,12 +609,30 @@ func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccou
|
||||
if err := s.accountRepo.Create(ctx, account); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 绑定分组
|
||||
if len(input.GroupIDs) > 0 {
|
||||
if err := s.accountRepo.BindGroups(ctx, account.ID, input.GroupIDs); err != nil {
|
||||
groupIDs := input.GroupIDs
|
||||
// 如果没有指定分组,自动绑定对应平台的默认分组
|
||||
if len(groupIDs) == 0 {
|
||||
defaultGroupName := input.Platform + "-default"
|
||||
groups, err := s.groupRepo.ListActiveByPlatform(ctx, input.Platform)
|
||||
if err == nil {
|
||||
for _, g := range groups {
|
||||
if g.Name == defaultGroupName {
|
||||
groupIDs = []int64{g.ID}
|
||||
log.Printf("[CreateAccount] Auto-binding account %d to default group %s (ID: %d)", account.ID, defaultGroupName, g.ID)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(groupIDs) > 0 {
|
||||
if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return account, nil
|
||||
}
|
||||
|
||||
@@ -652,11 +670,20 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
|
||||
account.Status = input.Status
|
||||
}
|
||||
|
||||
// 先验证分组是否存在(在任何写操作之前)
|
||||
if input.GroupIDs != nil {
|
||||
for _, groupID := range *input.GroupIDs {
|
||||
if _, err := s.groupRepo.GetByID(ctx, groupID); err != nil {
|
||||
return nil, fmt.Errorf("get group: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err := s.accountRepo.Update(ctx, account); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 更新分组绑定
|
||||
// 绑定分组
|
||||
if input.GroupIDs != nil {
|
||||
if err := s.accountRepo.BindGroups(ctx, account.ID, *input.GroupIDs); err != nil {
|
||||
return nil, err
|
||||
|
||||
823
backend/internal/service/antigravity_gateway_service.go
Normal file
823
backend/internal/service/antigravity_gateway_service.go
Normal file
@@ -0,0 +1,823 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
antigravityStickySessionTTL = time.Hour
|
||||
antigravityMaxRetries = 5
|
||||
antigravityRetryBaseDelay = 1 * time.Second
|
||||
antigravityRetryMaxDelay = 16 * time.Second
|
||||
)
|
||||
|
||||
// Antigravity 直接支持的模型
|
||||
var antigravitySupportedModels = map[string]bool{
|
||||
"claude-opus-4-5-thinking": true,
|
||||
"claude-sonnet-4-5": true,
|
||||
"claude-sonnet-4-5-thinking": true,
|
||||
"gemini-2.5-flash": true,
|
||||
"gemini-2.5-flash-lite": true,
|
||||
"gemini-2.5-flash-thinking": true,
|
||||
"gemini-3-flash": true,
|
||||
"gemini-3-pro-low": true,
|
||||
"gemini-3-pro-high": true,
|
||||
"gemini-3-pro-preview": true,
|
||||
"gemini-3-pro-image": true,
|
||||
}
|
||||
|
||||
// Antigravity 系统默认模型映射表(不支持 → 支持)
|
||||
var antigravityModelMapping = map[string]string{
|
||||
"claude-3-5-sonnet-20241022": "claude-sonnet-4-5",
|
||||
"claude-3-5-sonnet-20240620": "claude-sonnet-4-5",
|
||||
"claude-sonnet-4-5-20250929": "claude-sonnet-4-5-thinking",
|
||||
"claude-opus-4": "claude-opus-4-5-thinking",
|
||||
"claude-opus-4-5-20251101": "claude-opus-4-5-thinking",
|
||||
"claude-haiku-4": "gemini-3-flash",
|
||||
"claude-haiku-4-5": "gemini-3-flash",
|
||||
"claude-3-haiku-20240307": "gemini-3-flash",
|
||||
"claude-haiku-4-5-20251001": "gemini-3-flash",
|
||||
// 生图模型:官方名 → Antigravity 内部名
|
||||
"gemini-3-pro-image-preview": "gemini-3-pro-image",
|
||||
}
|
||||
|
||||
// AntigravityGatewayService 处理 Antigravity 平台的 API 转发
|
||||
type AntigravityGatewayService struct {
|
||||
accountRepo AccountRepository
|
||||
tokenProvider *AntigravityTokenProvider
|
||||
rateLimitService *RateLimitService
|
||||
httpUpstream HTTPUpstream
|
||||
}
|
||||
|
||||
func NewAntigravityGatewayService(
|
||||
accountRepo AccountRepository,
|
||||
_ GatewayCache,
|
||||
tokenProvider *AntigravityTokenProvider,
|
||||
rateLimitService *RateLimitService,
|
||||
httpUpstream HTTPUpstream,
|
||||
) *AntigravityGatewayService {
|
||||
return &AntigravityGatewayService{
|
||||
accountRepo: accountRepo,
|
||||
tokenProvider: tokenProvider,
|
||||
rateLimitService: rateLimitService,
|
||||
httpUpstream: httpUpstream,
|
||||
}
|
||||
}
|
||||
|
||||
// GetTokenProvider 返回 token provider
|
||||
func (s *AntigravityGatewayService) GetTokenProvider() *AntigravityTokenProvider {
|
||||
return s.tokenProvider
|
||||
}
|
||||
|
||||
// getMappedModel 获取映射后的模型名
|
||||
func (s *AntigravityGatewayService) getMappedModel(account *Account, requestedModel string) string {
|
||||
// 1. 优先使用账户级映射(复用现有方法)
|
||||
if mapped := account.GetMappedModel(requestedModel); mapped != requestedModel {
|
||||
return mapped
|
||||
}
|
||||
|
||||
// 2. 系统默认映射
|
||||
if mapped, ok := antigravityModelMapping[requestedModel]; ok {
|
||||
return mapped
|
||||
}
|
||||
|
||||
// 3. Gemini 模型透传
|
||||
if strings.HasPrefix(requestedModel, "gemini-") {
|
||||
return requestedModel
|
||||
}
|
||||
|
||||
// 4. Claude 前缀透传直接支持的模型
|
||||
if antigravitySupportedModels[requestedModel] {
|
||||
return requestedModel
|
||||
}
|
||||
|
||||
// 5. 默认值
|
||||
return "claude-sonnet-4-5"
|
||||
}
|
||||
|
||||
// IsModelSupported 检查模型是否被支持
|
||||
func (s *AntigravityGatewayService) IsModelSupported(requestedModel string) bool {
|
||||
// 直接支持的模型
|
||||
if antigravitySupportedModels[requestedModel] {
|
||||
return true
|
||||
}
|
||||
// 可映射的模型
|
||||
if _, ok := antigravityModelMapping[requestedModel]; ok {
|
||||
return true
|
||||
}
|
||||
// Gemini 前缀透传
|
||||
if strings.HasPrefix(requestedModel, "gemini-") {
|
||||
return true
|
||||
}
|
||||
// Claude 模型支持(通过默认映射)
|
||||
if strings.HasPrefix(requestedModel, "claude-") {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// wrapV1InternalRequest 包装请求为 v1internal 格式
|
||||
func (s *AntigravityGatewayService) wrapV1InternalRequest(projectID, model string, originalBody []byte) ([]byte, error) {
|
||||
var request any
|
||||
if err := json.Unmarshal(originalBody, &request); err != nil {
|
||||
return nil, fmt.Errorf("解析请求体失败: %w", err)
|
||||
}
|
||||
|
||||
wrapped := map[string]any{
|
||||
"project": projectID,
|
||||
"requestId": "agent-" + uuid.New().String(),
|
||||
"userAgent": "sub2api",
|
||||
"requestType": "agent",
|
||||
"model": model,
|
||||
"request": request,
|
||||
}
|
||||
|
||||
return json.Marshal(wrapped)
|
||||
}
|
||||
|
||||
// unwrapV1InternalResponse 解包 v1internal 响应
|
||||
func (s *AntigravityGatewayService) unwrapV1InternalResponse(body []byte) ([]byte, error) {
|
||||
var outer map[string]any
|
||||
if err := json.Unmarshal(body, &outer); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp, ok := outer["response"]; ok {
|
||||
return json.Marshal(resp)
|
||||
}
|
||||
|
||||
return body, nil
|
||||
}
|
||||
|
||||
// Forward 转发 Claude 协议请求(Claude → Gemini 转换)
|
||||
func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
// 解析 Claude 请求
|
||||
var claudeReq antigravity.ClaudeRequest
|
||||
if err := json.Unmarshal(body, &claudeReq); err != nil {
|
||||
return nil, fmt.Errorf("parse claude request: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(claudeReq.Model) == "" {
|
||||
return nil, fmt.Errorf("missing model")
|
||||
}
|
||||
|
||||
originalModel := claudeReq.Model
|
||||
mappedModel := s.getMappedModel(account, claudeReq.Model)
|
||||
if mappedModel != claudeReq.Model {
|
||||
log.Printf("Antigravity model mapping: %s -> %s (account: %s)", claudeReq.Model, mappedModel, account.Name)
|
||||
}
|
||||
|
||||
// 获取 access_token
|
||||
if s.tokenProvider == nil {
|
||||
return nil, errors.New("antigravity token provider not configured")
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("获取 access_token 失败: %w", err)
|
||||
}
|
||||
|
||||
// 获取 project_id
|
||||
projectID := strings.TrimSpace(account.GetCredential("project_id"))
|
||||
if projectID == "" {
|
||||
return nil, errors.New("project_id not found in credentials")
|
||||
}
|
||||
|
||||
// 代理 URL
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
// 转换 Claude 请求为 Gemini 格式
|
||||
geminiBody, err := antigravity.TransformClaudeToGemini(&claudeReq, projectID, mappedModel)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("transform request: %w", err)
|
||||
}
|
||||
|
||||
// 构建上游 URL
|
||||
action := "generateContent"
|
||||
if claudeReq.Stream {
|
||||
action = "streamGenerateContent"
|
||||
}
|
||||
fullURL := fmt.Sprintf("%s/v1internal:%s", antigravity.BaseURL, action)
|
||||
if claudeReq.Stream {
|
||||
fullURL += "?alt=sse"
|
||||
}
|
||||
|
||||
// 重试循环
|
||||
var resp *http.Response
|
||||
for attempt := 1; attempt <= antigravityMaxRetries; attempt++ {
|
||||
upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(geminiBody))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
upstreamReq.Header.Set("Content-Type", "application/json")
|
||||
upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
upstreamReq.Header.Set("User-Agent", antigravity.UserAgent)
|
||||
|
||||
resp, err = s.httpUpstream.Do(upstreamReq, proxyURL)
|
||||
if err != nil {
|
||||
if attempt < antigravityMaxRetries {
|
||||
log.Printf("Antigravity account %d: upstream request failed, retry %d/%d: %v", account.ID, attempt, antigravityMaxRetries, err)
|
||||
sleepAntigravityBackoff(attempt)
|
||||
continue
|
||||
}
|
||||
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries")
|
||||
}
|
||||
|
||||
if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(resp.StatusCode) {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
_ = resp.Body.Close()
|
||||
|
||||
if attempt < antigravityMaxRetries {
|
||||
log.Printf("Antigravity account %d: upstream status %d, retry %d/%d", account.ID, resp.StatusCode, attempt, antigravityMaxRetries)
|
||||
sleepAntigravityBackoff(attempt)
|
||||
continue
|
||||
}
|
||||
// 所有重试都失败,标记限流状态
|
||||
if resp.StatusCode == 429 {
|
||||
s.handleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
}
|
||||
// 最后一次尝试也失败
|
||||
resp = &http.Response{
|
||||
StatusCode: resp.StatusCode,
|
||||
Header: resp.Header.Clone(),
|
||||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
break
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
// 处理错误响应
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
s.handleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
|
||||
if s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode}
|
||||
}
|
||||
|
||||
return nil, s.writeMappedClaudeError(c, resp.StatusCode, respBody)
|
||||
}
|
||||
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
if requestID != "" {
|
||||
c.Header("x-request-id", requestID)
|
||||
}
|
||||
|
||||
var usage *ClaudeUsage
|
||||
var firstTokenMs *int
|
||||
if claudeReq.Stream {
|
||||
streamRes, err := s.handleClaudeStreamingResponse(c, resp, startTime, originalModel)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
usage = streamRes.usage
|
||||
firstTokenMs = streamRes.firstTokenMs
|
||||
} else {
|
||||
usage, err = s.handleClaudeNonStreamingResponse(c, resp, originalModel)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel, // 使用原始模型用于计费和日志
|
||||
Stream: claudeReq.Stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ForwardGemini 转发 Gemini 协议请求
|
||||
func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte) (*ForwardResult, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
if strings.TrimSpace(originalModel) == "" {
|
||||
return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing model in URL")
|
||||
}
|
||||
if strings.TrimSpace(action) == "" {
|
||||
return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing action in URL")
|
||||
}
|
||||
if len(body) == 0 {
|
||||
return nil, s.writeGoogleError(c, http.StatusBadRequest, "Request body is empty")
|
||||
}
|
||||
|
||||
switch action {
|
||||
case "generateContent", "streamGenerateContent", "countTokens":
|
||||
// ok
|
||||
default:
|
||||
return nil, s.writeGoogleError(c, http.StatusNotFound, "Unsupported action: "+action)
|
||||
}
|
||||
|
||||
mappedModel := s.getMappedModel(account, originalModel)
|
||||
|
||||
// 获取 access_token
|
||||
if s.tokenProvider == nil {
|
||||
return nil, errors.New("antigravity token provider not configured")
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("获取 access_token 失败: %w", err)
|
||||
}
|
||||
|
||||
// 获取 project_id
|
||||
projectID := strings.TrimSpace(account.GetCredential("project_id"))
|
||||
if projectID == "" {
|
||||
return nil, errors.New("project_id not found in credentials")
|
||||
}
|
||||
|
||||
// 代理 URL
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
// 包装请求
|
||||
wrappedBody, err := s.wrapV1InternalRequest(projectID, mappedModel, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 构建上游 URL
|
||||
upstreamAction := action
|
||||
if action == "generateContent" && stream {
|
||||
upstreamAction = "streamGenerateContent"
|
||||
}
|
||||
fullURL := fmt.Sprintf("%s/v1internal:%s", antigravity.BaseURL, upstreamAction)
|
||||
if stream || upstreamAction == "streamGenerateContent" {
|
||||
fullURL += "?alt=sse"
|
||||
}
|
||||
|
||||
// 重试循环
|
||||
var resp *http.Response
|
||||
for attempt := 1; attempt <= antigravityMaxRetries; attempt++ {
|
||||
upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(wrappedBody))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
upstreamReq.Header.Set("Content-Type", "application/json")
|
||||
upstreamReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
upstreamReq.Header.Set("User-Agent", antigravity.UserAgent)
|
||||
|
||||
resp, err = s.httpUpstream.Do(upstreamReq, proxyURL)
|
||||
if err != nil {
|
||||
if attempt < antigravityMaxRetries {
|
||||
log.Printf("Antigravity account %d: upstream request failed, retry %d/%d: %v", account.ID, attempt, antigravityMaxRetries, err)
|
||||
sleepAntigravityBackoff(attempt)
|
||||
continue
|
||||
}
|
||||
if action == "countTokens" {
|
||||
estimated := estimateGeminiCountTokens(body)
|
||||
c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
|
||||
return &ForwardResult{
|
||||
RequestID: "",
|
||||
Usage: ClaudeUsage{},
|
||||
Model: originalModel,
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: nil,
|
||||
}, nil
|
||||
}
|
||||
return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries")
|
||||
}
|
||||
|
||||
if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(resp.StatusCode) {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
_ = resp.Body.Close()
|
||||
|
||||
if attempt < antigravityMaxRetries {
|
||||
log.Printf("Antigravity account %d: upstream status %d, retry %d/%d", account.ID, resp.StatusCode, attempt, antigravityMaxRetries)
|
||||
sleepAntigravityBackoff(attempt)
|
||||
continue
|
||||
}
|
||||
// 所有重试都失败,标记限流状态
|
||||
if resp.StatusCode == 429 {
|
||||
s.handleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
}
|
||||
if action == "countTokens" {
|
||||
estimated := estimateGeminiCountTokens(body)
|
||||
c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
|
||||
return &ForwardResult{
|
||||
RequestID: "",
|
||||
Usage: ClaudeUsage{},
|
||||
Model: originalModel,
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: nil,
|
||||
}, nil
|
||||
}
|
||||
resp = &http.Response{
|
||||
StatusCode: resp.StatusCode,
|
||||
Header: resp.Header.Clone(),
|
||||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
break
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
if requestID != "" {
|
||||
c.Header("x-request-id", requestID)
|
||||
}
|
||||
|
||||
// 处理错误响应
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
s.handleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
|
||||
if action == "countTokens" {
|
||||
estimated := estimateGeminiCountTokens(body)
|
||||
c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: ClaudeUsage{},
|
||||
Model: originalModel,
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: nil,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode}
|
||||
}
|
||||
|
||||
// 解包并返回错误
|
||||
unwrapped, _ := s.unwrapV1InternalResponse(respBody)
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
if contentType == "" {
|
||||
contentType = "application/json"
|
||||
}
|
||||
c.Data(resp.StatusCode, contentType, unwrapped)
|
||||
return nil, fmt.Errorf("antigravity upstream error: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var usage *ClaudeUsage
|
||||
var firstTokenMs *int
|
||||
|
||||
if stream || upstreamAction == "streamGenerateContent" {
|
||||
streamRes, err := s.handleGeminiStreamingResponse(c, resp, startTime)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
usage = streamRes.usage
|
||||
firstTokenMs = streamRes.firstTokenMs
|
||||
} else {
|
||||
usageResp, err := s.handleGeminiNonStreamingResponse(c, resp)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
usage = usageResp
|
||||
}
|
||||
|
||||
if usage == nil {
|
||||
usage = &ClaudeUsage{}
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
Stream: stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) shouldRetryUpstreamError(statusCode int) bool {
|
||||
switch statusCode {
|
||||
case 429, 500, 502, 503, 504, 529:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) shouldFailoverUpstreamError(statusCode int) bool {
|
||||
switch statusCode {
|
||||
case 401, 403, 429, 529:
|
||||
return true
|
||||
default:
|
||||
return statusCode >= 500
|
||||
}
|
||||
}
|
||||
|
||||
func sleepAntigravityBackoff(attempt int) {
|
||||
sleepGeminiBackoff(attempt) // 复用 Gemini 的退避逻辑
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, body []byte) {
|
||||
// 429 使用 Gemini 格式解析(从 body 解析重置时间)
|
||||
if statusCode == 429 {
|
||||
resetAt := ParseGeminiRateLimitResetTime(body)
|
||||
if resetAt == nil {
|
||||
// 解析失败:Gemini 有重试时间用 5 分钟,Claude 没有用 1 分钟
|
||||
defaultDur := 1 * time.Minute
|
||||
if bytes.Contains(body, []byte("Please retry in")) || bytes.Contains(body, []byte("retryDelay")) {
|
||||
defaultDur = 5 * time.Minute
|
||||
}
|
||||
ra := time.Now().Add(defaultDur)
|
||||
_ = s.accountRepo.SetRateLimited(ctx, account.ID, ra)
|
||||
return
|
||||
}
|
||||
_ = s.accountRepo.SetRateLimited(ctx, account.ID, time.Unix(*resetAt, 0))
|
||||
return
|
||||
}
|
||||
// 其他错误码继续使用 rateLimitService
|
||||
if s.rateLimitService == nil {
|
||||
return
|
||||
}
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, headers, body)
|
||||
}
|
||||
|
||||
type antigravityStreamResult struct {
|
||||
usage *ClaudeUsage
|
||||
firstTokenMs *int
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) {
|
||||
c.Status(resp.StatusCode)
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
if contentType == "" {
|
||||
contentType = "text/event-stream; charset=utf-8"
|
||||
}
|
||||
c.Header("Content-Type", contentType)
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
return nil, errors.New("streaming not supported")
|
||||
}
|
||||
|
||||
reader := bufio.NewReader(resp.Body)
|
||||
usage := &ClaudeUsage{}
|
||||
var firstTokenMs *int
|
||||
|
||||
for {
|
||||
line, err := reader.ReadString('\n')
|
||||
if len(line) > 0 {
|
||||
trimmed := strings.TrimRight(line, "\r\n")
|
||||
if strings.HasPrefix(trimmed, "data:") {
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
|
||||
if payload == "" || payload == "[DONE]" {
|
||||
_, _ = io.WriteString(c.Writer, line)
|
||||
flusher.Flush()
|
||||
} else {
|
||||
// 解包 v1internal 响应
|
||||
inner, parseErr := s.unwrapV1InternalResponse([]byte(payload))
|
||||
if parseErr == nil && inner != nil {
|
||||
payload = string(inner)
|
||||
}
|
||||
|
||||
// 解析 usage
|
||||
var parsed map[string]any
|
||||
if json.Unmarshal(inner, &parsed) == nil {
|
||||
if u := extractGeminiUsage(parsed); u != nil {
|
||||
usage = u
|
||||
}
|
||||
}
|
||||
|
||||
if firstTokenMs == nil {
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
firstTokenMs = &ms
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", payload)
|
||||
flusher.Flush()
|
||||
}
|
||||
} else {
|
||||
_, _ = io.WriteString(c.Writer, line)
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleGeminiNonStreamingResponse(c *gin.Context, resp *http.Response) (*ClaudeUsage, error) {
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 解包 v1internal 响应
|
||||
unwrapped, _ := s.unwrapV1InternalResponse(respBody)
|
||||
|
||||
var parsed map[string]any
|
||||
if json.Unmarshal(unwrapped, &parsed) == nil {
|
||||
if u := extractGeminiUsage(parsed); u != nil {
|
||||
c.Data(resp.StatusCode, "application/json", unwrapped)
|
||||
return u, nil
|
||||
}
|
||||
}
|
||||
|
||||
c.Data(resp.StatusCode, "application/json", unwrapped)
|
||||
return &ClaudeUsage{}, nil
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) writeClaudeError(c *gin.Context, status int, errType, message string) error {
|
||||
c.JSON(status, gin.H{
|
||||
"type": "error",
|
||||
"error": gin.H{"type": errType, "message": message},
|
||||
})
|
||||
return fmt.Errorf("%s", message)
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) writeMappedClaudeError(c *gin.Context, upstreamStatus int, body []byte) error {
|
||||
// 记录上游错误详情便于调试
|
||||
log.Printf("Antigravity upstream error %d: %s", upstreamStatus, string(body))
|
||||
|
||||
var statusCode int
|
||||
var errType, errMsg string
|
||||
|
||||
switch upstreamStatus {
|
||||
case 400:
|
||||
statusCode = http.StatusBadRequest
|
||||
errType = "invalid_request_error"
|
||||
errMsg = "Invalid request"
|
||||
case 401:
|
||||
statusCode = http.StatusBadGateway
|
||||
errType = "authentication_error"
|
||||
errMsg = "Upstream authentication failed"
|
||||
case 403:
|
||||
statusCode = http.StatusBadGateway
|
||||
errType = "permission_error"
|
||||
errMsg = "Upstream access forbidden"
|
||||
case 429:
|
||||
statusCode = http.StatusTooManyRequests
|
||||
errType = "rate_limit_error"
|
||||
errMsg = "Upstream rate limit exceeded"
|
||||
case 529:
|
||||
statusCode = http.StatusServiceUnavailable
|
||||
errType = "overloaded_error"
|
||||
errMsg = "Upstream service overloaded"
|
||||
default:
|
||||
statusCode = http.StatusBadGateway
|
||||
errType = "upstream_error"
|
||||
errMsg = "Upstream request failed"
|
||||
}
|
||||
|
||||
c.JSON(statusCode, gin.H{
|
||||
"type": "error",
|
||||
"error": gin.H{"type": errType, "message": errMsg},
|
||||
})
|
||||
return fmt.Errorf("upstream error: %d", upstreamStatus)
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int, message string) error {
|
||||
statusStr := "UNKNOWN"
|
||||
switch status {
|
||||
case 400:
|
||||
statusStr = "INVALID_ARGUMENT"
|
||||
case 404:
|
||||
statusStr = "NOT_FOUND"
|
||||
case 429:
|
||||
statusStr = "RESOURCE_EXHAUSTED"
|
||||
case 500:
|
||||
statusStr = "INTERNAL"
|
||||
case 502, 503:
|
||||
statusStr = "UNAVAILABLE"
|
||||
}
|
||||
|
||||
c.JSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"code": status,
|
||||
"message": message,
|
||||
"status": statusStr,
|
||||
},
|
||||
})
|
||||
return fmt.Errorf("%s", message)
|
||||
}
|
||||
|
||||
// handleClaudeNonStreamingResponse 处理 Claude 非流式响应(Gemini → Claude 转换)
|
||||
func (s *AntigravityGatewayService) handleClaudeNonStreamingResponse(c *gin.Context, resp *http.Response, originalModel string) (*ClaudeUsage, error) {
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
|
||||
if err != nil {
|
||||
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to read upstream response")
|
||||
}
|
||||
|
||||
// 转换 Gemini 响应为 Claude 格式
|
||||
claudeResp, agUsage, err := antigravity.TransformGeminiToClaude(body, originalModel)
|
||||
if err != nil {
|
||||
log.Printf("Transform Gemini to Claude failed: %v, body: %s", err, string(body))
|
||||
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
|
||||
}
|
||||
|
||||
c.Data(http.StatusOK, "application/json", claudeResp)
|
||||
|
||||
// 转换为 service.ClaudeUsage
|
||||
usage := &ClaudeUsage{
|
||||
InputTokens: agUsage.InputTokens,
|
||||
OutputTokens: agUsage.OutputTokens,
|
||||
CacheCreationInputTokens: agUsage.CacheCreationInputTokens,
|
||||
CacheReadInputTokens: agUsage.CacheReadInputTokens,
|
||||
}
|
||||
return usage, nil
|
||||
}
|
||||
|
||||
// handleClaudeStreamingResponse 处理 Claude 流式响应(Gemini SSE → Claude SSE 转换)
|
||||
func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
c.Status(http.StatusOK)
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
return nil, errors.New("streaming not supported")
|
||||
}
|
||||
|
||||
processor := antigravity.NewStreamingProcessor(originalModel)
|
||||
var firstTokenMs *int
|
||||
reader := bufio.NewReader(resp.Body)
|
||||
|
||||
// 辅助函数:转换 antigravity.ClaudeUsage 到 service.ClaudeUsage
|
||||
convertUsage := func(agUsage *antigravity.ClaudeUsage) *ClaudeUsage {
|
||||
if agUsage == nil {
|
||||
return &ClaudeUsage{}
|
||||
}
|
||||
return &ClaudeUsage{
|
||||
InputTokens: agUsage.InputTokens,
|
||||
OutputTokens: agUsage.OutputTokens,
|
||||
CacheCreationInputTokens: agUsage.CacheCreationInputTokens,
|
||||
CacheReadInputTokens: agUsage.CacheReadInputTokens,
|
||||
}
|
||||
}
|
||||
|
||||
for {
|
||||
line, err := reader.ReadString('\n')
|
||||
if err != nil && !errors.Is(err, io.EOF) {
|
||||
return nil, fmt.Errorf("stream read error: %w", err)
|
||||
}
|
||||
|
||||
if len(line) > 0 {
|
||||
// 处理 SSE 行,转换为 Claude 格式
|
||||
claudeEvents := processor.ProcessLine(strings.TrimRight(line, "\r\n"))
|
||||
|
||||
if len(claudeEvents) > 0 {
|
||||
if firstTokenMs == nil {
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
firstTokenMs = &ms
|
||||
}
|
||||
|
||||
if _, writeErr := c.Writer.Write(claudeEvents); writeErr != nil {
|
||||
finalEvents, agUsage := processor.Finish()
|
||||
if len(finalEvents) > 0 {
|
||||
_, _ = c.Writer.Write(finalEvents)
|
||||
}
|
||||
return &antigravityStreamResult{usage: convertUsage(agUsage), firstTokenMs: firstTokenMs}, writeErr
|
||||
}
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// 发送结束事件
|
||||
finalEvents, agUsage := processor.Finish()
|
||||
if len(finalEvents) > 0 {
|
||||
_, _ = c.Writer.Write(finalEvents)
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
return &antigravityStreamResult{usage: convertUsage(agUsage), firstTokenMs: firstTokenMs}, nil
|
||||
}
|
||||
269
backend/internal/service/antigravity_model_mapping_test.go
Normal file
269
backend/internal/service/antigravity_model_mapping_test.go
Normal file
@@ -0,0 +1,269 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestIsAntigravityModelSupported(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
model string
|
||||
expected bool
|
||||
}{
|
||||
// 直接支持的模型
|
||||
{"直接支持 - claude-sonnet-4-5", "claude-sonnet-4-5", true},
|
||||
{"直接支持 - claude-opus-4-5-thinking", "claude-opus-4-5-thinking", true},
|
||||
{"直接支持 - claude-sonnet-4-5-thinking", "claude-sonnet-4-5-thinking", true},
|
||||
{"直接支持 - gemini-2.5-flash", "gemini-2.5-flash", true},
|
||||
{"直接支持 - gemini-2.5-flash-lite", "gemini-2.5-flash-lite", true},
|
||||
{"直接支持 - gemini-3-pro-high", "gemini-3-pro-high", true},
|
||||
|
||||
// 可映射的模型
|
||||
{"可映射 - claude-3-5-sonnet-20241022", "claude-3-5-sonnet-20241022", true},
|
||||
{"可映射 - claude-3-5-sonnet-20240620", "claude-3-5-sonnet-20240620", true},
|
||||
{"可映射 - claude-opus-4", "claude-opus-4", true},
|
||||
{"可映射 - claude-haiku-4", "claude-haiku-4", true},
|
||||
{"可映射 - claude-3-haiku-20240307", "claude-3-haiku-20240307", true},
|
||||
|
||||
// Gemini 前缀透传
|
||||
{"Gemini前缀 - gemini-1.5-pro", "gemini-1.5-pro", true},
|
||||
{"Gemini前缀 - gemini-unknown-model", "gemini-unknown-model", true},
|
||||
{"Gemini前缀 - gemini-future-version", "gemini-future-version", true},
|
||||
|
||||
// Claude 前缀兜底
|
||||
{"Claude前缀 - claude-unknown-model", "claude-unknown-model", true},
|
||||
{"Claude前缀 - claude-3-opus-20240229", "claude-3-opus-20240229", true},
|
||||
{"Claude前缀 - claude-future-version", "claude-future-version", true},
|
||||
|
||||
// 不支持的模型
|
||||
{"不支持 - gpt-4", "gpt-4", false},
|
||||
{"不支持 - gpt-4o", "gpt-4o", false},
|
||||
{"不支持 - llama-3", "llama-3", false},
|
||||
{"不支持 - mistral-7b", "mistral-7b", false},
|
||||
{"不支持 - 空字符串", "", false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := IsAntigravityModelSupported(tt.model)
|
||||
require.Equal(t, tt.expected, got, "model: %s", tt.model)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityGatewayService_GetMappedModel(t *testing.T) {
|
||||
svc := &AntigravityGatewayService{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
requestedModel string
|
||||
accountMapping map[string]string
|
||||
expected string
|
||||
}{
|
||||
// 1. 账户级映射优先(注意:model_mapping 在 credentials 中存储为 map[string]any)
|
||||
{
|
||||
name: "账户映射优先",
|
||||
requestedModel: "claude-3-5-sonnet-20241022",
|
||||
accountMapping: map[string]string{"claude-3-5-sonnet-20241022": "custom-model"},
|
||||
expected: "custom-model",
|
||||
},
|
||||
{
|
||||
name: "账户映射覆盖系统映射",
|
||||
requestedModel: "claude-opus-4",
|
||||
accountMapping: map[string]string{"claude-opus-4": "my-opus"},
|
||||
expected: "my-opus",
|
||||
},
|
||||
|
||||
// 2. 系统默认映射
|
||||
{
|
||||
name: "系统映射 - claude-3-5-sonnet-20241022",
|
||||
requestedModel: "claude-3-5-sonnet-20241022",
|
||||
accountMapping: nil,
|
||||
expected: "claude-sonnet-4-5",
|
||||
},
|
||||
{
|
||||
name: "系统映射 - claude-3-5-sonnet-20240620",
|
||||
requestedModel: "claude-3-5-sonnet-20240620",
|
||||
accountMapping: nil,
|
||||
expected: "claude-sonnet-4-5",
|
||||
},
|
||||
{
|
||||
name: "系统映射 - claude-opus-4",
|
||||
requestedModel: "claude-opus-4",
|
||||
accountMapping: nil,
|
||||
expected: "claude-opus-4-5-thinking",
|
||||
},
|
||||
{
|
||||
name: "系统映射 - claude-opus-4-5-20251101",
|
||||
requestedModel: "claude-opus-4-5-20251101",
|
||||
accountMapping: nil,
|
||||
expected: "claude-opus-4-5-thinking",
|
||||
},
|
||||
{
|
||||
name: "系统映射 - claude-haiku-4 → gemini-3-flash",
|
||||
requestedModel: "claude-haiku-4",
|
||||
accountMapping: nil,
|
||||
expected: "gemini-3-flash",
|
||||
},
|
||||
{
|
||||
name: "系统映射 - claude-haiku-4-5 → gemini-3-flash",
|
||||
requestedModel: "claude-haiku-4-5",
|
||||
accountMapping: nil,
|
||||
expected: "gemini-3-flash",
|
||||
},
|
||||
{
|
||||
name: "系统映射 - claude-3-haiku-20240307 → gemini-3-flash",
|
||||
requestedModel: "claude-3-haiku-20240307",
|
||||
accountMapping: nil,
|
||||
expected: "gemini-3-flash",
|
||||
},
|
||||
{
|
||||
name: "系统映射 - claude-haiku-4-5-20251001 → gemini-3-flash",
|
||||
requestedModel: "claude-haiku-4-5-20251001",
|
||||
accountMapping: nil,
|
||||
expected: "gemini-3-flash",
|
||||
},
|
||||
{
|
||||
name: "系统映射 - claude-sonnet-4-5-20250929",
|
||||
requestedModel: "claude-sonnet-4-5-20250929",
|
||||
accountMapping: nil,
|
||||
expected: "claude-sonnet-4-5-thinking",
|
||||
},
|
||||
|
||||
// 3. Gemini 透传
|
||||
{
|
||||
name: "Gemini透传 - gemini-2.5-flash",
|
||||
requestedModel: "gemini-2.5-flash",
|
||||
accountMapping: nil,
|
||||
expected: "gemini-2.5-flash",
|
||||
},
|
||||
{
|
||||
name: "Gemini透传 - gemini-1.5-pro",
|
||||
requestedModel: "gemini-1.5-pro",
|
||||
accountMapping: nil,
|
||||
expected: "gemini-1.5-pro",
|
||||
},
|
||||
{
|
||||
name: "Gemini透传 - gemini-future-model",
|
||||
requestedModel: "gemini-future-model",
|
||||
accountMapping: nil,
|
||||
expected: "gemini-future-model",
|
||||
},
|
||||
|
||||
// 4. 直接支持的模型
|
||||
{
|
||||
name: "直接支持 - claude-sonnet-4-5",
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
accountMapping: nil,
|
||||
expected: "claude-sonnet-4-5",
|
||||
},
|
||||
{
|
||||
name: "直接支持 - claude-opus-4-5-thinking",
|
||||
requestedModel: "claude-opus-4-5-thinking",
|
||||
accountMapping: nil,
|
||||
expected: "claude-opus-4-5-thinking",
|
||||
},
|
||||
{
|
||||
name: "直接支持 - claude-sonnet-4-5-thinking",
|
||||
requestedModel: "claude-sonnet-4-5-thinking",
|
||||
accountMapping: nil,
|
||||
expected: "claude-sonnet-4-5-thinking",
|
||||
},
|
||||
|
||||
// 5. 默认值 fallback(未知 claude 模型)
|
||||
{
|
||||
name: "默认值 - claude-unknown",
|
||||
requestedModel: "claude-unknown",
|
||||
accountMapping: nil,
|
||||
expected: "claude-sonnet-4-5",
|
||||
},
|
||||
{
|
||||
name: "默认值 - claude-3-opus-20240229",
|
||||
requestedModel: "claude-3-opus-20240229",
|
||||
accountMapping: nil,
|
||||
expected: "claude-sonnet-4-5",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformAntigravity,
|
||||
}
|
||||
if tt.accountMapping != nil {
|
||||
// GetModelMapping 期望 model_mapping 是 map[string]any 格式
|
||||
mappingAny := make(map[string]any)
|
||||
for k, v := range tt.accountMapping {
|
||||
mappingAny[k] = v
|
||||
}
|
||||
account.Credentials = map[string]any{
|
||||
"model_mapping": mappingAny,
|
||||
}
|
||||
}
|
||||
|
||||
got := svc.getMappedModel(account, tt.requestedModel)
|
||||
require.Equal(t, tt.expected, got, "model: %s", tt.requestedModel)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityGatewayService_GetMappedModel_EdgeCases(t *testing.T) {
|
||||
svc := &AntigravityGatewayService{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
requestedModel string
|
||||
expected string
|
||||
}{
|
||||
// 空字符串回退到默认值
|
||||
{"空字符串", "", "claude-sonnet-4-5"},
|
||||
|
||||
// 非 claude/gemini 前缀回退到默认值
|
||||
{"非claude/gemini前缀 - gpt", "gpt-4", "claude-sonnet-4-5"},
|
||||
{"非claude/gemini前缀 - llama", "llama-3", "claude-sonnet-4-5"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := &Account{Platform: PlatformAntigravity}
|
||||
got := svc.getMappedModel(account, tt.requestedModel)
|
||||
require.Equal(t, tt.expected, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityGatewayService_IsModelSupported(t *testing.T) {
|
||||
svc := &AntigravityGatewayService{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
model string
|
||||
expected bool
|
||||
}{
|
||||
// 直接支持
|
||||
{"直接支持 - claude-sonnet-4-5", "claude-sonnet-4-5", true},
|
||||
{"直接支持 - gemini-3-flash", "gemini-3-flash", true},
|
||||
|
||||
// 可映射
|
||||
{"可映射 - claude-opus-4", "claude-opus-4", true},
|
||||
|
||||
// 前缀透传
|
||||
{"Gemini前缀", "gemini-unknown", true},
|
||||
{"Claude前缀", "claude-unknown", true},
|
||||
|
||||
// 不支持
|
||||
{"不支持 - gpt-4", "gpt-4", false},
|
||||
{"不支持 - 空字符串", "", false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := svc.IsModelSupported(tt.model)
|
||||
require.Equal(t, tt.expected, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
267
backend/internal/service/antigravity_oauth_service.go
Normal file
267
backend/internal/service/antigravity_oauth_service.go
Normal file
@@ -0,0 +1,267 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
)
|
||||
|
||||
type AntigravityOAuthService struct {
|
||||
sessionStore *antigravity.SessionStore
|
||||
proxyRepo ProxyRepository
|
||||
}
|
||||
|
||||
func NewAntigravityOAuthService(proxyRepo ProxyRepository) *AntigravityOAuthService {
|
||||
return &AntigravityOAuthService{
|
||||
sessionStore: antigravity.NewSessionStore(),
|
||||
proxyRepo: proxyRepo,
|
||||
}
|
||||
}
|
||||
|
||||
// AntigravityAuthURLResult is the result of generating an authorization URL
|
||||
type AntigravityAuthURLResult struct {
|
||||
AuthURL string `json:"auth_url"`
|
||||
SessionID string `json:"session_id"`
|
||||
State string `json:"state"`
|
||||
}
|
||||
|
||||
// GenerateAuthURL 生成 Google OAuth 授权链接
|
||||
func (s *AntigravityOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64) (*AntigravityAuthURLResult, error) {
|
||||
state, err := antigravity.GenerateState()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("生成 state 失败: %w", err)
|
||||
}
|
||||
|
||||
codeVerifier, err := antigravity.GenerateCodeVerifier()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("生成 code_verifier 失败: %w", err)
|
||||
}
|
||||
|
||||
sessionID, err := antigravity.GenerateSessionID()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("生成 session_id 失败: %w", err)
|
||||
}
|
||||
|
||||
var proxyURL string
|
||||
if proxyID != nil {
|
||||
proxy, err := s.proxyRepo.GetByID(ctx, *proxyID)
|
||||
if err == nil && proxy != nil {
|
||||
proxyURL = proxy.URL()
|
||||
}
|
||||
}
|
||||
|
||||
session := &antigravity.OAuthSession{
|
||||
State: state,
|
||||
CodeVerifier: codeVerifier,
|
||||
ProxyURL: proxyURL,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
s.sessionStore.Set(sessionID, session)
|
||||
|
||||
codeChallenge := antigravity.GenerateCodeChallenge(codeVerifier)
|
||||
authURL := antigravity.BuildAuthorizationURL(state, codeChallenge)
|
||||
|
||||
return &AntigravityAuthURLResult{
|
||||
AuthURL: authURL,
|
||||
SessionID: sessionID,
|
||||
State: state,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// AntigravityExchangeCodeInput 交换 code 的输入
|
||||
type AntigravityExchangeCodeInput struct {
|
||||
SessionID string
|
||||
State string
|
||||
Code string
|
||||
ProxyID *int64
|
||||
}
|
||||
|
||||
// AntigravityTokenInfo token 信息
|
||||
type AntigravityTokenInfo struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
ExpiresIn int64 `json:"expires_in"`
|
||||
ExpiresAt int64 `json:"expires_at"`
|
||||
TokenType string `json:"token_type"`
|
||||
Email string `json:"email,omitempty"`
|
||||
ProjectID string `json:"project_id,omitempty"`
|
||||
}
|
||||
|
||||
// ExchangeCode 用 authorization code 交换 token
|
||||
func (s *AntigravityOAuthService) ExchangeCode(ctx context.Context, input *AntigravityExchangeCodeInput) (*AntigravityTokenInfo, error) {
|
||||
session, ok := s.sessionStore.Get(input.SessionID)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("session 不存在或已过期")
|
||||
}
|
||||
|
||||
if strings.TrimSpace(input.State) == "" || input.State != session.State {
|
||||
return nil, fmt.Errorf("state 无效")
|
||||
}
|
||||
|
||||
// 确定代理 URL
|
||||
proxyURL := session.ProxyURL
|
||||
if input.ProxyID != nil {
|
||||
proxy, err := s.proxyRepo.GetByID(ctx, *input.ProxyID)
|
||||
if err == nil && proxy != nil {
|
||||
proxyURL = proxy.URL()
|
||||
}
|
||||
}
|
||||
|
||||
client := antigravity.NewClient(proxyURL)
|
||||
|
||||
// 交换 token
|
||||
tokenResp, err := client.ExchangeCode(ctx, input.Code, session.CodeVerifier)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("token 交换失败: %w", err)
|
||||
}
|
||||
|
||||
// 删除 session
|
||||
s.sessionStore.Delete(input.SessionID)
|
||||
|
||||
// 计算过期时间(减去 5 分钟安全窗口)
|
||||
expiresAt := time.Now().Unix() + tokenResp.ExpiresIn - 300
|
||||
|
||||
result := &AntigravityTokenInfo{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ExpiresIn: tokenResp.ExpiresIn,
|
||||
ExpiresAt: expiresAt,
|
||||
TokenType: tokenResp.TokenType,
|
||||
}
|
||||
|
||||
// 获取用户信息
|
||||
userInfo, err := client.GetUserInfo(ctx, tokenResp.AccessToken)
|
||||
if err != nil {
|
||||
fmt.Printf("[AntigravityOAuth] 警告: 获取用户信息失败: %v\n", err)
|
||||
} else {
|
||||
result.Email = userInfo.Email
|
||||
}
|
||||
|
||||
// 获取 project_id
|
||||
loadResp, err := client.LoadCodeAssist(ctx, tokenResp.AccessToken)
|
||||
if err != nil {
|
||||
fmt.Printf("[AntigravityOAuth] 警告: 获取 project_id 失败: %v\n", err)
|
||||
} else if loadResp != nil && loadResp.CloudAICompanionProject != "" {
|
||||
result.ProjectID = loadResp.CloudAICompanionProject
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// RefreshToken 刷新 token
|
||||
func (s *AntigravityOAuthService) RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*AntigravityTokenInfo, error) {
|
||||
var lastErr error
|
||||
|
||||
for attempt := 0; attempt <= 3; attempt++ {
|
||||
if attempt > 0 {
|
||||
backoff := time.Duration(1<<uint(attempt-1)) * time.Second
|
||||
if backoff > 30*time.Second {
|
||||
backoff = 30 * time.Second
|
||||
}
|
||||
time.Sleep(backoff)
|
||||
}
|
||||
|
||||
client := antigravity.NewClient(proxyURL)
|
||||
tokenResp, err := client.RefreshToken(ctx, refreshToken)
|
||||
if err == nil {
|
||||
expiresAt := time.Now().Unix() + tokenResp.ExpiresIn - 300
|
||||
return &AntigravityTokenInfo{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ExpiresIn: tokenResp.ExpiresIn,
|
||||
ExpiresAt: expiresAt,
|
||||
TokenType: tokenResp.TokenType,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if isNonRetryableAntigravityOAuthError(err) {
|
||||
return nil, err
|
||||
}
|
||||
lastErr = err
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("token 刷新失败 (重试后): %w", lastErr)
|
||||
}
|
||||
|
||||
func isNonRetryableAntigravityOAuthError(err error) bool {
|
||||
msg := err.Error()
|
||||
nonRetryable := []string{
|
||||
"invalid_grant",
|
||||
"invalid_client",
|
||||
"unauthorized_client",
|
||||
"access_denied",
|
||||
}
|
||||
for _, needle := range nonRetryable {
|
||||
if strings.Contains(msg, needle) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// RefreshAccountToken 刷新账户的 token
|
||||
func (s *AntigravityOAuthService) RefreshAccountToken(ctx context.Context, account *Account) (*AntigravityTokenInfo, error) {
|
||||
if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth {
|
||||
return nil, fmt.Errorf("非 Antigravity OAuth 账户")
|
||||
}
|
||||
|
||||
refreshToken := account.GetCredential("refresh_token")
|
||||
if strings.TrimSpace(refreshToken) == "" {
|
||||
return nil, fmt.Errorf("无可用的 refresh_token")
|
||||
}
|
||||
|
||||
var proxyURL string
|
||||
if account.ProxyID != nil {
|
||||
proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID)
|
||||
if err == nil && proxy != nil {
|
||||
proxyURL = proxy.URL()
|
||||
}
|
||||
}
|
||||
|
||||
tokenInfo, err := s.RefreshToken(ctx, refreshToken, proxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 保留原有的 project_id 和 email
|
||||
existingProjectID := strings.TrimSpace(account.GetCredential("project_id"))
|
||||
if existingProjectID != "" {
|
||||
tokenInfo.ProjectID = existingProjectID
|
||||
}
|
||||
existingEmail := strings.TrimSpace(account.GetCredential("email"))
|
||||
if existingEmail != "" {
|
||||
tokenInfo.Email = existingEmail
|
||||
}
|
||||
|
||||
return tokenInfo, nil
|
||||
}
|
||||
|
||||
// BuildAccountCredentials 构建账户凭证
|
||||
func (s *AntigravityOAuthService) BuildAccountCredentials(tokenInfo *AntigravityTokenInfo) map[string]any {
|
||||
creds := map[string]any{
|
||||
"access_token": tokenInfo.AccessToken,
|
||||
"expires_at": strconv.FormatInt(tokenInfo.ExpiresAt, 10),
|
||||
}
|
||||
if tokenInfo.RefreshToken != "" {
|
||||
creds["refresh_token"] = tokenInfo.RefreshToken
|
||||
}
|
||||
if tokenInfo.TokenType != "" {
|
||||
creds["token_type"] = tokenInfo.TokenType
|
||||
}
|
||||
if tokenInfo.Email != "" {
|
||||
creds["email"] = tokenInfo.Email
|
||||
}
|
||||
if tokenInfo.ProjectID != "" {
|
||||
creds["project_id"] = tokenInfo.ProjectID
|
||||
}
|
||||
return creds
|
||||
}
|
||||
|
||||
// Stop 停止服务
|
||||
func (s *AntigravityOAuthService) Stop() {
|
||||
s.sessionStore.Stop()
|
||||
}
|
||||
225
backend/internal/service/antigravity_quota_refresher.go
Normal file
225
backend/internal/service/antigravity_quota_refresher.go
Normal file
@@ -0,0 +1,225 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
)
|
||||
|
||||
// AntigravityQuotaRefresher 定时刷新 Antigravity 账户的配额信息
|
||||
type AntigravityQuotaRefresher struct {
|
||||
accountRepo AccountRepository
|
||||
proxyRepo ProxyRepository
|
||||
cfg *config.TokenRefreshConfig
|
||||
|
||||
stopCh chan struct{}
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewAntigravityQuotaRefresher 创建配额刷新器
|
||||
func NewAntigravityQuotaRefresher(
|
||||
accountRepo AccountRepository,
|
||||
proxyRepo ProxyRepository,
|
||||
_ *AntigravityOAuthService,
|
||||
cfg *config.Config,
|
||||
) *AntigravityQuotaRefresher {
|
||||
return &AntigravityQuotaRefresher{
|
||||
accountRepo: accountRepo,
|
||||
proxyRepo: proxyRepo,
|
||||
cfg: &cfg.TokenRefresh,
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start 启动后台配额刷新服务
|
||||
func (r *AntigravityQuotaRefresher) Start() {
|
||||
if !r.cfg.Enabled {
|
||||
log.Println("[AntigravityQuota] Service disabled by configuration")
|
||||
return
|
||||
}
|
||||
|
||||
r.wg.Add(1)
|
||||
go r.refreshLoop()
|
||||
|
||||
log.Printf("[AntigravityQuota] Service started (check every %d minutes)", r.cfg.CheckIntervalMinutes)
|
||||
}
|
||||
|
||||
// Stop 停止服务
|
||||
func (r *AntigravityQuotaRefresher) Stop() {
|
||||
close(r.stopCh)
|
||||
r.wg.Wait()
|
||||
log.Println("[AntigravityQuota] Service stopped")
|
||||
}
|
||||
|
||||
// refreshLoop 刷新循环
|
||||
func (r *AntigravityQuotaRefresher) refreshLoop() {
|
||||
defer r.wg.Done()
|
||||
|
||||
checkInterval := time.Duration(r.cfg.CheckIntervalMinutes) * time.Minute
|
||||
if checkInterval < time.Minute {
|
||||
checkInterval = 5 * time.Minute
|
||||
}
|
||||
|
||||
ticker := time.NewTicker(checkInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
// 启动时立即执行一次
|
||||
r.processRefresh()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
r.processRefresh()
|
||||
case <-r.stopCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// processRefresh 执行一次刷新
|
||||
func (r *AntigravityQuotaRefresher) processRefresh() {
|
||||
ctx := context.Background()
|
||||
|
||||
// 查询所有 active 的账户,然后过滤 antigravity 平台
|
||||
allAccounts, err := r.accountRepo.ListActive(ctx)
|
||||
if err != nil {
|
||||
log.Printf("[AntigravityQuota] Failed to list accounts: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 过滤 antigravity 平台账户
|
||||
var accounts []Account
|
||||
for _, acc := range allAccounts {
|
||||
if acc.Platform == PlatformAntigravity {
|
||||
accounts = append(accounts, acc)
|
||||
}
|
||||
}
|
||||
|
||||
if len(accounts) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
refreshed, failed := 0, 0
|
||||
|
||||
for i := range accounts {
|
||||
account := &accounts[i]
|
||||
|
||||
if err := r.refreshAccountQuota(ctx, account); err != nil {
|
||||
log.Printf("[AntigravityQuota] Account %d (%s) failed: %v", account.ID, account.Name, err)
|
||||
failed++
|
||||
} else {
|
||||
refreshed++
|
||||
}
|
||||
}
|
||||
|
||||
log.Printf("[AntigravityQuota] Cycle complete: total=%d, refreshed=%d, failed=%d",
|
||||
len(accounts), refreshed, failed)
|
||||
}
|
||||
|
||||
// refreshAccountQuota 刷新单个账户的配额
|
||||
func (r *AntigravityQuotaRefresher) refreshAccountQuota(ctx context.Context, account *Account) error {
|
||||
accessToken := account.GetCredential("access_token")
|
||||
projectID := account.GetCredential("project_id")
|
||||
|
||||
if accessToken == "" || projectID == "" {
|
||||
return nil // 没有有效凭证,跳过
|
||||
}
|
||||
|
||||
// token 过期则跳过,由 TokenRefreshService 负责刷新
|
||||
if r.isTokenExpired(account) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 获取代理 URL
|
||||
var proxyURL string
|
||||
if account.ProxyID != nil {
|
||||
proxy, err := r.proxyRepo.GetByID(ctx, *account.ProxyID)
|
||||
if err == nil && proxy != nil {
|
||||
proxyURL = proxy.URL()
|
||||
}
|
||||
}
|
||||
|
||||
client := antigravity.NewClient(proxyURL)
|
||||
|
||||
// 获取账户类型(tier)
|
||||
loadResp, _ := client.LoadCodeAssist(ctx, accessToken)
|
||||
if loadResp != nil {
|
||||
r.updateAccountTier(account, loadResp)
|
||||
}
|
||||
|
||||
// 调用 API 获取配额
|
||||
modelsResp, err := client.FetchAvailableModels(ctx, accessToken, projectID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 解析配额数据并更新 extra 字段
|
||||
r.updateAccountQuota(account, modelsResp)
|
||||
|
||||
// 保存到数据库
|
||||
return r.accountRepo.Update(ctx, account)
|
||||
}
|
||||
|
||||
// isTokenExpired 检查 token 是否过期
|
||||
func (r *AntigravityQuotaRefresher) isTokenExpired(account *Account) bool {
|
||||
expiresAt := parseAntigravityExpiresAt(account)
|
||||
if expiresAt == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
// 提前 5 分钟认为过期
|
||||
return time.Now().Add(5 * time.Minute).After(*expiresAt)
|
||||
}
|
||||
|
||||
// updateAccountTier 更新账户类型信息
|
||||
func (r *AntigravityQuotaRefresher) updateAccountTier(account *Account, loadResp *antigravity.LoadCodeAssistResponse) {
|
||||
if account.Extra == nil {
|
||||
account.Extra = make(map[string]any)
|
||||
}
|
||||
|
||||
tier := loadResp.GetTier()
|
||||
if tier != "" {
|
||||
account.Extra["tier"] = tier
|
||||
}
|
||||
|
||||
// 保存不符合条件的原因(如 INELIGIBLE_ACCOUNT)
|
||||
if len(loadResp.IneligibleTiers) > 0 && loadResp.IneligibleTiers[0] != nil {
|
||||
ineligible := loadResp.IneligibleTiers[0]
|
||||
if ineligible.ReasonCode != "" {
|
||||
account.Extra["ineligible_reason_code"] = ineligible.ReasonCode
|
||||
}
|
||||
if ineligible.ReasonMessage != "" {
|
||||
account.Extra["ineligible_reason_message"] = ineligible.ReasonMessage
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// updateAccountQuota 更新账户的配额信息
|
||||
func (r *AntigravityQuotaRefresher) updateAccountQuota(account *Account, modelsResp *antigravity.FetchAvailableModelsResponse) {
|
||||
if account.Extra == nil {
|
||||
account.Extra = make(map[string]any)
|
||||
}
|
||||
|
||||
quota := make(map[string]any)
|
||||
|
||||
for modelName, modelInfo := range modelsResp.Models {
|
||||
if modelInfo.QuotaInfo == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// 转换 remainingFraction (0.0-1.0) 为百分比 (0-100)
|
||||
remaining := int(modelInfo.QuotaInfo.RemainingFraction * 100)
|
||||
|
||||
quota[modelName] = map[string]any{
|
||||
"remaining": remaining,
|
||||
"reset_time": modelInfo.QuotaInfo.ResetTime,
|
||||
}
|
||||
}
|
||||
|
||||
account.Extra["quota"] = quota
|
||||
account.Extra["last_quota_check"] = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
145
backend/internal/service/antigravity_token_provider.go
Normal file
145
backend/internal/service/antigravity_token_provider.go
Normal file
@@ -0,0 +1,145 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
antigravityTokenRefreshSkew = 3 * time.Minute
|
||||
antigravityTokenCacheSkew = 5 * time.Minute
|
||||
)
|
||||
|
||||
// AntigravityTokenCache Token 缓存接口(复用 GeminiTokenCache 接口定义)
|
||||
type AntigravityTokenCache = GeminiTokenCache
|
||||
|
||||
// AntigravityTokenProvider 管理 Antigravity 账户的 access_token
|
||||
type AntigravityTokenProvider struct {
|
||||
accountRepo AccountRepository
|
||||
tokenCache AntigravityTokenCache
|
||||
antigravityOAuthService *AntigravityOAuthService
|
||||
}
|
||||
|
||||
func NewAntigravityTokenProvider(
|
||||
accountRepo AccountRepository,
|
||||
tokenCache AntigravityTokenCache,
|
||||
antigravityOAuthService *AntigravityOAuthService,
|
||||
) *AntigravityTokenProvider {
|
||||
return &AntigravityTokenProvider{
|
||||
accountRepo: accountRepo,
|
||||
tokenCache: tokenCache,
|
||||
antigravityOAuthService: antigravityOAuthService,
|
||||
}
|
||||
}
|
||||
|
||||
// GetAccessToken 获取有效的 access_token
|
||||
func (p *AntigravityTokenProvider) GetAccessToken(ctx context.Context, account *Account) (string, error) {
|
||||
if account == nil {
|
||||
return "", errors.New("account is nil")
|
||||
}
|
||||
if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth {
|
||||
return "", errors.New("not an antigravity oauth account")
|
||||
}
|
||||
|
||||
cacheKey := antigravityTokenCacheKey(account)
|
||||
|
||||
// 1. 先尝试缓存
|
||||
if p.tokenCache != nil {
|
||||
if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil && strings.TrimSpace(token) != "" {
|
||||
return token, nil
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 如果即将过期则刷新
|
||||
expiresAt := parseAntigravityExpiresAt(account)
|
||||
needsRefresh := expiresAt == nil || time.Until(*expiresAt) <= antigravityTokenRefreshSkew
|
||||
if needsRefresh && p.tokenCache != nil {
|
||||
locked, err := p.tokenCache.AcquireRefreshLock(ctx, cacheKey, 30*time.Second)
|
||||
if err == nil && locked {
|
||||
defer func() { _ = p.tokenCache.ReleaseRefreshLock(ctx, cacheKey) }()
|
||||
|
||||
// 拿到锁后再次检查缓存(另一个 worker 可能已刷新)
|
||||
if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil && strings.TrimSpace(token) != "" {
|
||||
return token, nil
|
||||
}
|
||||
|
||||
// 从数据库获取最新账户信息
|
||||
fresh, err := p.accountRepo.GetByID(ctx, account.ID)
|
||||
if err == nil && fresh != nil {
|
||||
account = fresh
|
||||
}
|
||||
expiresAt = parseAntigravityExpiresAt(account)
|
||||
if expiresAt == nil || time.Until(*expiresAt) <= antigravityTokenRefreshSkew {
|
||||
if p.antigravityOAuthService == nil {
|
||||
return "", errors.New("antigravity oauth service not configured")
|
||||
}
|
||||
tokenInfo, err := p.antigravityOAuthService.RefreshAccountToken(ctx, account)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
newCredentials := p.antigravityOAuthService.BuildAccountCredentials(tokenInfo)
|
||||
for k, v := range account.Credentials {
|
||||
if _, exists := newCredentials[k]; !exists {
|
||||
newCredentials[k] = v
|
||||
}
|
||||
}
|
||||
account.Credentials = newCredentials
|
||||
if updateErr := p.accountRepo.Update(ctx, account); updateErr != nil {
|
||||
log.Printf("[AntigravityTokenProvider] Failed to update account credentials: %v", updateErr)
|
||||
}
|
||||
expiresAt = parseAntigravityExpiresAt(account)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
accessToken := account.GetCredential("access_token")
|
||||
if strings.TrimSpace(accessToken) == "" {
|
||||
return "", errors.New("access_token not found in credentials")
|
||||
}
|
||||
|
||||
// 3. 存入缓存
|
||||
if p.tokenCache != nil {
|
||||
ttl := 30 * time.Minute
|
||||
if expiresAt != nil {
|
||||
until := time.Until(*expiresAt)
|
||||
switch {
|
||||
case until > antigravityTokenCacheSkew:
|
||||
ttl = until - antigravityTokenCacheSkew
|
||||
case until > 0:
|
||||
ttl = until
|
||||
default:
|
||||
ttl = time.Minute
|
||||
}
|
||||
}
|
||||
_ = p.tokenCache.SetAccessToken(ctx, cacheKey, accessToken, ttl)
|
||||
}
|
||||
|
||||
return accessToken, nil
|
||||
}
|
||||
|
||||
func antigravityTokenCacheKey(account *Account) string {
|
||||
projectID := strings.TrimSpace(account.GetCredential("project_id"))
|
||||
if projectID != "" {
|
||||
return "ag:" + projectID
|
||||
}
|
||||
return "ag:account:" + strconv.FormatInt(account.ID, 10)
|
||||
}
|
||||
|
||||
func parseAntigravityExpiresAt(account *Account) *time.Time {
|
||||
raw := strings.TrimSpace(account.GetCredential("expires_at"))
|
||||
if raw == "" {
|
||||
return nil
|
||||
}
|
||||
if unixSec, err := strconv.ParseInt(raw, 10, 64); err == nil && unixSec > 0 {
|
||||
t := time.Unix(unixSec, 0)
|
||||
return &t
|
||||
}
|
||||
if t, err := time.Parse(time.RFC3339, raw); err == nil {
|
||||
return &t
|
||||
}
|
||||
return nil
|
||||
}
|
||||
57
backend/internal/service/antigravity_token_refresher.go
Normal file
57
backend/internal/service/antigravity_token_refresher.go
Normal file
@@ -0,0 +1,57 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AntigravityTokenRefresher 实现 TokenRefresher 接口
|
||||
type AntigravityTokenRefresher struct {
|
||||
antigravityOAuthService *AntigravityOAuthService
|
||||
}
|
||||
|
||||
func NewAntigravityTokenRefresher(antigravityOAuthService *AntigravityOAuthService) *AntigravityTokenRefresher {
|
||||
return &AntigravityTokenRefresher{
|
||||
antigravityOAuthService: antigravityOAuthService,
|
||||
}
|
||||
}
|
||||
|
||||
// CanRefresh 检查是否可以刷新此账户
|
||||
func (r *AntigravityTokenRefresher) CanRefresh(account *Account) bool {
|
||||
return account.Platform == PlatformAntigravity && account.Type == AccountTypeOAuth
|
||||
}
|
||||
|
||||
// NeedsRefresh 检查账户是否需要刷新
|
||||
func (r *AntigravityTokenRefresher) NeedsRefresh(account *Account, refreshWindow time.Duration) bool {
|
||||
if !r.CanRefresh(account) {
|
||||
return false
|
||||
}
|
||||
expiresAtStr := account.GetCredential("expires_at")
|
||||
if expiresAtStr == "" {
|
||||
return false
|
||||
}
|
||||
expiresAt, err := strconv.ParseInt(expiresAtStr, 10, 64)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
expiryTime := time.Unix(expiresAt, 0)
|
||||
return time.Until(expiryTime) < refreshWindow
|
||||
}
|
||||
|
||||
// Refresh 执行 token 刷新
|
||||
func (r *AntigravityTokenRefresher) Refresh(ctx context.Context, account *Account) (map[string]any, error) {
|
||||
tokenInfo, err := r.antigravityOAuthService.RefreshAccountToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
newCredentials := r.antigravityOAuthService.BuildAccountCredentials(tokenInfo)
|
||||
for k, v := range account.Credentials {
|
||||
if _, exists := newCredentials[k]; !exists {
|
||||
newCredentials[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
return newCredentials, nil
|
||||
}
|
||||
@@ -34,6 +34,7 @@ type ApiKeyRepository interface {
|
||||
Delete(ctx context.Context, id int64) error
|
||||
|
||||
ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams) ([]ApiKey, *pagination.PaginationResult, error)
|
||||
VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error)
|
||||
CountByUserID(ctx context.Context, userID int64) (int64, error)
|
||||
ExistsByKey(ctx context.Context, key string) (bool, error)
|
||||
ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]ApiKey, *pagination.PaginationResult, error)
|
||||
@@ -256,6 +257,18 @@ func (s *ApiKeyService) List(ctx context.Context, userID int64, params paginatio
|
||||
return keys, pagination, nil
|
||||
}
|
||||
|
||||
func (s *ApiKeyService) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
|
||||
if len(apiKeyIDs) == 0 {
|
||||
return []int64{}, nil
|
||||
}
|
||||
|
||||
validIDs, err := s.apiKeyRepo.VerifyOwnership(ctx, userID, apiKeyIDs)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("verify api key ownership: %w", err)
|
||||
}
|
||||
return validIDs, nil
|
||||
}
|
||||
|
||||
// GetByID 根据ID获取API Key
|
||||
func (s *ApiKeyService) GetByID(ctx context.Context, id int64) (*ApiKey, error) {
|
||||
apiKey, err := s.apiKeyRepo.GetByID(ctx, id)
|
||||
|
||||
@@ -20,6 +20,7 @@ var (
|
||||
ErrEmailExists = infraerrors.Conflict("EMAIL_EXISTS", "email already exists")
|
||||
ErrInvalidToken = infraerrors.Unauthorized("INVALID_TOKEN", "invalid token")
|
||||
ErrTokenExpired = infraerrors.Unauthorized("TOKEN_EXPIRED", "token has expired")
|
||||
ErrTokenRevoked = infraerrors.Unauthorized("TOKEN_REVOKED", "token has been revoked")
|
||||
ErrEmailVerifyRequired = infraerrors.BadRequest("EMAIL_VERIFY_REQUIRED", "email verification is required")
|
||||
ErrRegDisabled = infraerrors.Forbidden("REGISTRATION_DISABLED", "registration is currently disabled")
|
||||
ErrServiceUnavailable = infraerrors.ServiceUnavailable("SERVICE_UNAVAILABLE", "service temporarily unavailable")
|
||||
@@ -27,9 +28,10 @@ var (
|
||||
|
||||
// JWTClaims JWT载荷数据
|
||||
type JWTClaims struct {
|
||||
UserID int64 `json:"user_id"`
|
||||
Email string `json:"email"`
|
||||
Role string `json:"role"`
|
||||
UserID int64 `json:"user_id"`
|
||||
Email string `json:"email"`
|
||||
Role string `json:"role"`
|
||||
TokenVersion int64 `json:"token_version"` // Used to invalidate tokens on password change
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
@@ -311,9 +313,10 @@ func (s *AuthService) GenerateToken(user *User) (string, error) {
|
||||
expiresAt := now.Add(time.Duration(s.cfg.JWT.ExpireHour) * time.Hour)
|
||||
|
||||
claims := &JWTClaims{
|
||||
UserID: user.ID,
|
||||
Email: user.Email,
|
||||
Role: user.Role,
|
||||
UserID: user.ID,
|
||||
Email: user.Email,
|
||||
Role: user.Role,
|
||||
TokenVersion: user.TokenVersion,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(expiresAt),
|
||||
IssuedAt: jwt.NewNumericDate(now),
|
||||
@@ -368,6 +371,12 @@ func (s *AuthService) RefreshToken(ctx context.Context, oldTokenString string) (
|
||||
return "", ErrUserNotActive
|
||||
}
|
||||
|
||||
// Security: Check TokenVersion to prevent refreshing revoked tokens
|
||||
// This ensures tokens issued before a password change cannot be refreshed
|
||||
if claims.TokenVersion != user.TokenVersion {
|
||||
return "", ErrTokenRevoked
|
||||
}
|
||||
|
||||
// 生成新token
|
||||
return s.GenerateToken(user)
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/infrastructure/errors"
|
||||
)
|
||||
|
||||
@@ -32,14 +33,16 @@ type BillingCacheService struct {
|
||||
cache BillingCache
|
||||
userRepo UserRepository
|
||||
subRepo UserSubscriptionRepository
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
// NewBillingCacheService 创建计费缓存服务
|
||||
func NewBillingCacheService(cache BillingCache, userRepo UserRepository, subRepo UserSubscriptionRepository) *BillingCacheService {
|
||||
func NewBillingCacheService(cache BillingCache, userRepo UserRepository, subRepo UserSubscriptionRepository, cfg *config.Config) *BillingCacheService {
|
||||
return &BillingCacheService{
|
||||
cache: cache,
|
||||
userRepo: userRepo,
|
||||
subRepo: subRepo,
|
||||
cfg: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -224,6 +227,11 @@ func (s *BillingCacheService) InvalidateSubscription(ctx context.Context, userID
|
||||
// 余额模式:检查缓存余额 > 0
|
||||
// 订阅模式:检查缓存用量未超过限额(Group限额从参数传入)
|
||||
func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user *User, apiKey *ApiKey, group *Group, subscription *UserSubscription) error {
|
||||
// 简易模式:跳过所有计费检查
|
||||
if s.cfg.RunMode == config.RunModeSimple {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 判断计费模式
|
||||
isSubscriptionMode := group != nil && group.IsSubscriptionType() && subscription != nil
|
||||
|
||||
|
||||
76
backend/internal/service/deferred_service.go
Normal file
76
backend/internal/service/deferred_service.go
Normal file
@@ -0,0 +1,76 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// DeferredService provides deferred batch update functionality
|
||||
type DeferredService struct {
|
||||
accountRepo AccountRepository
|
||||
timingWheel *TimingWheelService
|
||||
interval time.Duration
|
||||
|
||||
lastUsedUpdates sync.Map
|
||||
}
|
||||
|
||||
// NewDeferredService creates a new DeferredService instance
|
||||
func NewDeferredService(accountRepo AccountRepository, timingWheel *TimingWheelService, interval time.Duration) *DeferredService {
|
||||
return &DeferredService{
|
||||
accountRepo: accountRepo,
|
||||
timingWheel: timingWheel,
|
||||
interval: interval,
|
||||
}
|
||||
}
|
||||
|
||||
// Start starts the deferred service
|
||||
func (s *DeferredService) Start() {
|
||||
s.timingWheel.ScheduleRecurring("deferred:last_used", s.interval, s.flushLastUsed)
|
||||
log.Printf("[DeferredService] Started (interval: %v)", s.interval)
|
||||
}
|
||||
|
||||
// Stop stops the deferred service
|
||||
func (s *DeferredService) Stop() {
|
||||
s.timingWheel.Cancel("deferred:last_used")
|
||||
s.flushLastUsed()
|
||||
log.Printf("[DeferredService] Service stopped")
|
||||
}
|
||||
|
||||
func (s *DeferredService) ScheduleLastUsedUpdate(accountID int64) {
|
||||
s.lastUsedUpdates.Store(accountID, time.Now())
|
||||
}
|
||||
|
||||
func (s *DeferredService) flushLastUsed() {
|
||||
updates := make(map[int64]time.Time)
|
||||
s.lastUsedUpdates.Range(func(key, value any) bool {
|
||||
id, ok := key.(int64)
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
ts, ok := value.(time.Time)
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
updates[id] = ts
|
||||
s.lastUsedUpdates.Delete(key)
|
||||
return true
|
||||
})
|
||||
|
||||
if len(updates) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.accountRepo.BatchUpdateLastUsed(ctx, updates); err != nil {
|
||||
log.Printf("[DeferredService] BatchUpdateLastUsed failed (%d accounts): %v", len(updates), err)
|
||||
for id, ts := range updates {
|
||||
s.lastUsedUpdates.Store(id, ts)
|
||||
}
|
||||
} else {
|
||||
log.Printf("[DeferredService] BatchUpdateLastUsed flushed %d accounts", len(updates))
|
||||
}
|
||||
}
|
||||
@@ -18,9 +18,10 @@ const (
|
||||
|
||||
// Platform constants
|
||||
const (
|
||||
PlatformAnthropic = "anthropic"
|
||||
PlatformOpenAI = "openai"
|
||||
PlatformGemini = "gemini"
|
||||
PlatformAnthropic = "anthropic"
|
||||
PlatformOpenAI = "openai"
|
||||
PlatformGemini = "gemini"
|
||||
PlatformAntigravity = "antigravity"
|
||||
)
|
||||
|
||||
// Account type constants
|
||||
|
||||
777
backend/internal/service/gateway_multiplatform_test.go
Normal file
777
backend/internal/service/gateway_multiplatform_test.go
Normal file
@@ -0,0 +1,777 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// testConfig 返回一个用于测试的默认配置
|
||||
func testConfig() *config.Config {
|
||||
return &config.Config{RunMode: config.RunModeStandard}
|
||||
}
|
||||
|
||||
// mockAccountRepoForPlatform 单平台测试用的 mock
|
||||
type mockAccountRepoForPlatform struct {
|
||||
accounts []Account
|
||||
accountsByID map[int64]*Account
|
||||
listPlatformFunc func(ctx context.Context, platform string) ([]Account, error)
|
||||
}
|
||||
|
||||
func (m *mockAccountRepoForPlatform) GetByID(ctx context.Context, id int64) (*Account, error) {
|
||||
if acc, ok := m.accountsByID[id]; ok {
|
||||
return acc, nil
|
||||
}
|
||||
return nil, errors.New("account not found")
|
||||
}
|
||||
|
||||
func (m *mockAccountRepoForPlatform) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
|
||||
if m.listPlatformFunc != nil {
|
||||
return m.listPlatformFunc(ctx, platform)
|
||||
}
|
||||
var result []Account
|
||||
for _, acc := range m.accounts {
|
||||
if acc.Platform == platform && acc.IsSchedulable() {
|
||||
result = append(result, acc)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (m *mockAccountRepoForPlatform) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
|
||||
return m.ListSchedulableByPlatform(ctx, platform)
|
||||
}
|
||||
|
||||
// Stub methods to implement AccountRepository interface
|
||||
func (m *mockAccountRepoForPlatform) Create(ctx context.Context, account *Account) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) Update(ctx context.Context, account *Account) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) Delete(ctx context.Context, id int64) error { return nil }
|
||||
func (m *mockAccountRepoForPlatform) List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string) ([]Account, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListActive(ctx context.Context) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListByPlatform(ctx context.Context, platform string) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) UpdateLastUsed(ctx context.Context, id int64) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) SetError(ctx context.Context, id int64, errorMsg string) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) SetSchedulable(ctx context.Context, id int64, schedulable bool) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) BindGroups(ctx context.Context, accountID int64, groupIDs []int64) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListSchedulable(ctx context.Context) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]Account, error) {
|
||||
var result []Account
|
||||
platformSet := make(map[string]bool)
|
||||
for _, p := range platforms {
|
||||
platformSet[p] = true
|
||||
}
|
||||
for _, acc := range m.accounts {
|
||||
if platformSet[acc.Platform] && acc.IsSchedulable() {
|
||||
result = append(result, acc)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error) {
|
||||
return m.ListSchedulableByPlatforms(ctx, platforms)
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) SetOverloaded(ctx context.Context, id int64, until time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ClearRateLimit(ctx context.Context, id int64) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) UpdateSessionWindow(ctx context.Context, id int64, start, end *time.Time, status string) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) BulkUpdate(ctx context.Context, ids []int64, updates AccountBulkUpdate) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Verify interface implementation
|
||||
var _ AccountRepository = (*mockAccountRepoForPlatform)(nil)
|
||||
|
||||
// mockGatewayCacheForPlatform 单平台测试用的 cache mock
|
||||
type mockGatewayCacheForPlatform struct {
|
||||
sessionBindings map[string]int64
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForPlatform) GetSessionAccountID(ctx context.Context, sessionHash string) (int64, error) {
|
||||
if id, ok := m.sessionBindings[sessionHash]; ok {
|
||||
return id, nil
|
||||
}
|
||||
return 0, errors.New("not found")
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForPlatform) SetSessionAccountID(ctx context.Context, sessionHash string, accountID int64, ttl time.Duration) error {
|
||||
if m.sessionBindings == nil {
|
||||
m.sessionBindings = make(map[string]int64)
|
||||
}
|
||||
m.sessionBindings[sessionHash] = accountID
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForPlatform) RefreshSessionTTL(ctx context.Context, sessionHash string, ttl time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func ptr[T any](v T) *T {
|
||||
return &v
|
||||
}
|
||||
|
||||
// TestGatewayService_SelectAccountForModelWithPlatform_Anthropic 测试 anthropic 单平台选择
|
||||
func TestGatewayService_SelectAccountForModelWithPlatform_Anthropic(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
{ID: 3, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true}, // 应被隔离
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(1), acc.ID, "应选择优先级最高的 anthropic 账户")
|
||||
require.Equal(t, PlatformAnthropic, acc.Platform, "应只返回 anthropic 平台账户")
|
||||
}
|
||||
|
||||
// TestGatewayService_SelectAccountForModelWithPlatform_Antigravity 测试 antigravity 单平台选择
|
||||
func TestGatewayService_SelectAccountForModelWithPlatform_Antigravity(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true}, // 应被隔离
|
||||
{ID: 2, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, PlatformAntigravity)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID)
|
||||
require.Equal(t, PlatformAntigravity, acc.Platform, "应只返回 antigravity 平台账户")
|
||||
}
|
||||
|
||||
// TestGatewayService_SelectAccountForModelWithPlatform_PriorityAndLastUsed 测试优先级和最后使用时间
|
||||
func TestGatewayService_SelectAccountForModelWithPlatform_PriorityAndLastUsed(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true, LastUsedAt: ptr(now.Add(-1 * time.Hour))},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true, LastUsedAt: ptr(now.Add(-2 * time.Hour))},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID, "同优先级应选择最久未用的账户")
|
||||
}
|
||||
|
||||
// TestGatewayService_SelectAccountForModelWithPlatform_NoAvailableAccounts 测试无可用账户
|
||||
func TestGatewayService_SelectAccountForModelWithPlatform_NoAvailableAccounts(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, acc)
|
||||
require.Contains(t, err.Error(), "no available accounts")
|
||||
}
|
||||
|
||||
// TestGatewayService_SelectAccountForModelWithPlatform_AllExcluded 测试所有账户被排除
|
||||
func TestGatewayService_SelectAccountForModelWithPlatform_AllExcluded(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
excludedIDs := map[int64]struct{}{1: {}, 2: {}}
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "", "claude-3-5-sonnet-20241022", excludedIDs, PlatformAnthropic)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, acc)
|
||||
}
|
||||
|
||||
// TestGatewayService_SelectAccountForModelWithPlatform_Schedulability 测试账户可调度性检查
|
||||
func TestGatewayService_SelectAccountForModelWithPlatform_Schedulability(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accounts []Account
|
||||
expectedID int64
|
||||
}{
|
||||
{
|
||||
name: "过载账户被跳过",
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true, OverloadUntil: ptr(now.Add(1 * time.Hour))},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
expectedID: 2,
|
||||
},
|
||||
{
|
||||
name: "限流账户被跳过",
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true, RateLimitResetAt: ptr(now.Add(1 * time.Hour))},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
expectedID: 2,
|
||||
},
|
||||
{
|
||||
name: "非active账户被跳过",
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: "error", Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
expectedID: 2,
|
||||
},
|
||||
{
|
||||
name: "schedulable=false被跳过",
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: false},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
expectedID: 2,
|
||||
},
|
||||
{
|
||||
name: "过期的过载账户可调度",
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true, OverloadUntil: ptr(now.Add(-1 * time.Hour))},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
expectedID: 1,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: tt.accounts,
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, tt.expectedID, acc.ID)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGatewayService_SelectAccountForModelWithPlatform_StickySession 测试粘性会话
|
||||
func TestGatewayService_SelectAccountForModelWithPlatform_StickySession(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("粘性会话命中-同平台", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{
|
||||
sessionBindings: map[string]int64{"session-123": 1},
|
||||
}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "session-123", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(1), acc.ID, "应返回粘性会话绑定的账户")
|
||||
})
|
||||
|
||||
t.Run("粘性会话不匹配平台-降级选择", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAntigravity, Priority: 2, Status: StatusActive, Schedulable: true}, // 粘性会话绑定但平台不匹配
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{
|
||||
sessionBindings: map[string]int64{"session-123": 1}, // 绑定 antigravity 账户
|
||||
}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
// 请求 anthropic 平台,但粘性会话绑定的是 antigravity 账户
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "session-123", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID, "粘性会话账户平台不匹配,应降级选择同平台账户")
|
||||
require.Equal(t, PlatformAnthropic, acc.Platform)
|
||||
})
|
||||
|
||||
t.Run("粘性会话账户被排除-降级选择", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{
|
||||
sessionBindings: map[string]int64{"session-123": 1},
|
||||
}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
excludedIDs := map[int64]struct{}{1: {}}
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "session-123", "claude-3-5-sonnet-20241022", excludedIDs, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID, "粘性会话账户被排除,应选择其他账户")
|
||||
})
|
||||
|
||||
t.Run("粘性会话账户不可调度-降级选择", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 2, Status: "error", Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{
|
||||
sessionBindings: map[string]int64{"session-123": 1},
|
||||
}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountForModelWithPlatform(ctx, nil, "session-123", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID, "粘性会话账户不可调度,应选择其他账户")
|
||||
})
|
||||
}
|
||||
|
||||
func TestGatewayService_isModelSupportedByAccount(t *testing.T) {
|
||||
svc := &GatewayService{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
model string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "Antigravity平台-支持claude模型",
|
||||
account: &Account{Platform: PlatformAntigravity},
|
||||
model: "claude-3-5-sonnet-20241022",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "Antigravity平台-支持gemini模型",
|
||||
account: &Account{Platform: PlatformAntigravity},
|
||||
model: "gemini-2.5-flash",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "Antigravity平台-不支持gpt模型",
|
||||
account: &Account{Platform: PlatformAntigravity},
|
||||
model: "gpt-4",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "Anthropic平台-无映射配置-支持所有模型",
|
||||
account: &Account{Platform: PlatformAnthropic},
|
||||
model: "claude-3-5-sonnet-20241022",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "Anthropic平台-有映射配置-只支持配置的模型",
|
||||
account: &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"claude-opus-4": "x"}},
|
||||
},
|
||||
model: "claude-3-5-sonnet-20241022",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "Anthropic平台-有映射配置-支持配置的模型",
|
||||
account: &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"claude-3-5-sonnet-20241022": "x"}},
|
||||
},
|
||||
model: "claude-3-5-sonnet-20241022",
|
||||
expected: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := svc.isModelSupportedByAccount(tt.account, tt.model)
|
||||
require.Equal(t, tt.expected, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGatewayService_selectAccountWithMixedScheduling 测试混合调度
|
||||
func TestGatewayService_selectAccountWithMixedScheduling(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("混合调度-包含启用mixed_scheduling的antigravity账户", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true, Extra: map[string]any{"mixed_scheduling": true}},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID, "应选择优先级最高的账户(包含启用混合调度的antigravity)")
|
||||
})
|
||||
|
||||
t.Run("混合调度-过滤未启用mixed_scheduling的antigravity账户", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true}, // 未启用 mixed_scheduling
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(1), acc.ID, "未启用mixed_scheduling的antigravity账户应被过滤")
|
||||
require.Equal(t, PlatformAnthropic, acc.Platform)
|
||||
})
|
||||
|
||||
t.Run("混合调度-粘性会话命中启用mixed_scheduling的antigravity账户", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAntigravity, Priority: 2, Status: StatusActive, Schedulable: true, Extra: map[string]any{"mixed_scheduling": true}},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{
|
||||
sessionBindings: map[string]int64{"session-123": 2},
|
||||
}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "session-123", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID, "应返回粘性会话绑定的启用mixed_scheduling的antigravity账户")
|
||||
})
|
||||
|
||||
t.Run("混合调度-粘性会话命中未启用mixed_scheduling的antigravity账户-降级选择", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformAntigravity, Priority: 2, Status: StatusActive, Schedulable: true}, // 未启用 mixed_scheduling
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{
|
||||
sessionBindings: map[string]int64{"session-123": 2},
|
||||
}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "session-123", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(1), acc.ID, "粘性会话绑定的账户未启用mixed_scheduling,应降级选择anthropic账户")
|
||||
})
|
||||
|
||||
t.Run("混合调度-仅有启用mixed_scheduling的antigravity账户", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true, Extra: map[string]any{"mixed_scheduling": true}},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(1), acc.ID)
|
||||
require.Equal(t, PlatformAntigravity, acc.Platform)
|
||||
})
|
||||
|
||||
t.Run("混合调度-无可用账户", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForPlatform{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true}, // 未启用 mixed_scheduling
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForPlatform{}
|
||||
|
||||
svc := &GatewayService{
|
||||
accountRepo: repo,
|
||||
cache: cache,
|
||||
cfg: testConfig(),
|
||||
}
|
||||
|
||||
acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, PlatformAnthropic)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, acc)
|
||||
require.Contains(t, err.Error(), "no available accounts")
|
||||
})
|
||||
}
|
||||
|
||||
// TestAccount_IsMixedSchedulingEnabled 测试混合调度开关检查
|
||||
func TestAccount_IsMixedSchedulingEnabled(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account Account
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "非antigravity平台-返回false",
|
||||
account: Account{Platform: PlatformAnthropic},
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "antigravity平台-无extra-返回false",
|
||||
account: Account{Platform: PlatformAntigravity},
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "antigravity平台-extra无mixed_scheduling-返回false",
|
||||
account: Account{Platform: PlatformAntigravity, Extra: map[string]any{}},
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "antigravity平台-mixed_scheduling=false-返回false",
|
||||
account: Account{Platform: PlatformAntigravity, Extra: map[string]any{"mixed_scheduling": false}},
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "antigravity平台-mixed_scheduling=true-返回true",
|
||||
account: Account{Platform: PlatformAntigravity, Extra: map[string]any{"mixed_scheduling": true}},
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "antigravity平台-mixed_scheduling非bool类型-返回false",
|
||||
account: Account{Platform: PlatformAntigravity, Extra: map[string]any{"mixed_scheduling": "true"}},
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := tt.account.IsMixedSchedulingEnabled()
|
||||
require.Equal(t, tt.expected, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -18,6 +18,7 @@ import (
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
|
||||
@@ -81,9 +82,19 @@ type ForwardResult struct {
|
||||
FirstTokenMs *int // 首字时间(流式请求)
|
||||
}
|
||||
|
||||
// UpstreamFailoverError indicates an upstream error that should trigger account failover.
|
||||
type UpstreamFailoverError struct {
|
||||
StatusCode int
|
||||
}
|
||||
|
||||
func (e *UpstreamFailoverError) Error() string {
|
||||
return fmt.Sprintf("upstream error: %d (failover)", e.StatusCode)
|
||||
}
|
||||
|
||||
// GatewayService handles API gateway operations
|
||||
type GatewayService struct {
|
||||
accountRepo AccountRepository
|
||||
groupRepo GroupRepository
|
||||
usageLogRepo UsageLogRepository
|
||||
userRepo UserRepository
|
||||
userSubRepo UserSubscriptionRepository
|
||||
@@ -94,11 +105,13 @@ type GatewayService struct {
|
||||
billingCacheService *BillingCacheService
|
||||
identityService *IdentityService
|
||||
httpUpstream HTTPUpstream
|
||||
deferredService *DeferredService
|
||||
}
|
||||
|
||||
// NewGatewayService creates a new GatewayService
|
||||
func NewGatewayService(
|
||||
accountRepo AccountRepository,
|
||||
groupRepo GroupRepository,
|
||||
usageLogRepo UsageLogRepository,
|
||||
userRepo UserRepository,
|
||||
userSubRepo UserSubscriptionRepository,
|
||||
@@ -109,9 +122,11 @@ func NewGatewayService(
|
||||
billingCacheService *BillingCacheService,
|
||||
identityService *IdentityService,
|
||||
httpUpstream HTTPUpstream,
|
||||
deferredService *DeferredService,
|
||||
) *GatewayService {
|
||||
return &GatewayService{
|
||||
accountRepo: accountRepo,
|
||||
groupRepo: groupRepo,
|
||||
usageLogRepo: usageLogRepo,
|
||||
userRepo: userRepo,
|
||||
userSubRepo: userSubRepo,
|
||||
@@ -122,6 +137,7 @@ func NewGatewayService(
|
||||
billingCacheService: billingCacheService,
|
||||
identityService: identityService,
|
||||
httpUpstream: httpUpstream,
|
||||
deferredService: deferredService,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -274,30 +290,77 @@ func (s *GatewayService) SelectAccount(ctx context.Context, groupID *int64, sess
|
||||
|
||||
// SelectAccountForModel 选择支持指定模型的账号(粘性会话+优先级+模型映射)
|
||||
func (s *GatewayService) SelectAccountForModel(ctx context.Context, groupID *int64, sessionHash string, requestedModel string) (*Account, error) {
|
||||
return s.SelectAccountForModelWithExclusions(ctx, groupID, sessionHash, requestedModel, nil)
|
||||
}
|
||||
|
||||
// SelectAccountForModelWithExclusions selects an account supporting the requested model while excluding specified accounts.
|
||||
func (s *GatewayService) SelectAccountForModelWithExclusions(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*Account, error) {
|
||||
// 优先检查 context 中的强制平台(/antigravity 路由)
|
||||
var platform string
|
||||
forcePlatform, hasForcePlatform := ctx.Value(ctxkey.ForcePlatform).(string)
|
||||
if hasForcePlatform && forcePlatform != "" {
|
||||
platform = forcePlatform
|
||||
} else if groupID != nil {
|
||||
// 根据分组 platform 决定查询哪种账号
|
||||
group, err := s.groupRepo.GetByID(ctx, *groupID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get group failed: %w", err)
|
||||
}
|
||||
platform = group.Platform
|
||||
} else {
|
||||
// 无分组时只使用原生 anthropic 平台
|
||||
platform = PlatformAnthropic
|
||||
}
|
||||
|
||||
// anthropic/gemini 分组支持混合调度(包含启用了 mixed_scheduling 的 antigravity 账户)
|
||||
// 注意:强制平台模式不走混合调度
|
||||
if (platform == PlatformAnthropic || platform == PlatformGemini) && !hasForcePlatform {
|
||||
return s.selectAccountWithMixedScheduling(ctx, groupID, sessionHash, requestedModel, excludedIDs, platform)
|
||||
}
|
||||
|
||||
// 强制平台模式:优先按分组查找,找不到再查全部该平台账户
|
||||
if hasForcePlatform && groupID != nil {
|
||||
account, err := s.selectAccountForModelWithPlatform(ctx, groupID, sessionHash, requestedModel, excludedIDs, platform)
|
||||
if err == nil {
|
||||
return account, nil
|
||||
}
|
||||
// 分组中找不到,回退查询全部该平台账户
|
||||
groupID = nil
|
||||
}
|
||||
|
||||
// antigravity 分组、强制平台模式或无分组使用单平台选择
|
||||
return s.selectAccountForModelWithPlatform(ctx, groupID, sessionHash, requestedModel, excludedIDs, platform)
|
||||
}
|
||||
|
||||
// selectAccountForModelWithPlatform 选择单平台账户(完全隔离)
|
||||
func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, platform string) (*Account, error) {
|
||||
// 1. 查询粘性会话
|
||||
if sessionHash != "" {
|
||||
accountID, err := s.cache.GetSessionAccountID(ctx, sessionHash)
|
||||
if err == nil && accountID > 0 {
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
// 使用IsSchedulable代替IsActive,确保限流/过载账号不会被选中
|
||||
// 同时检查模型支持
|
||||
if err == nil && account.IsSchedulable() && (requestedModel == "" || account.IsModelSupported(requestedModel)) {
|
||||
// 续期粘性会话
|
||||
if err := s.cache.RefreshSessionTTL(ctx, sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
if _, excluded := excludedIDs[accountID]; !excluded {
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
// 检查账号平台是否匹配(确保粘性会话不会跨平台)
|
||||
if err == nil && account.Platform == platform && account.IsSchedulable() && (requestedModel == "" || s.isModelSupportedByAccount(account, requestedModel)) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 获取可调度账号列表(排除限流和过载的账号,仅限 Anthropic 平台)
|
||||
// 2. 获取可调度账号列表(单平台)
|
||||
var accounts []Account
|
||||
var err error
|
||||
if groupID != nil {
|
||||
accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, PlatformAnthropic)
|
||||
if s.cfg.RunMode == config.RunModeSimple {
|
||||
// 简易模式:忽略 groupID,查询所有可用账号
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatform(ctx, platform)
|
||||
} else if groupID != nil {
|
||||
accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, platform)
|
||||
} else {
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatform(ctx, PlatformAnthropic)
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatform(ctx, platform)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query accounts failed: %w", err)
|
||||
@@ -307,19 +370,19 @@ func (s *GatewayService) SelectAccountForModel(ctx context.Context, groupID *int
|
||||
var selected *Account
|
||||
for i := range accounts {
|
||||
acc := &accounts[i]
|
||||
// 检查模型支持
|
||||
if requestedModel != "" && !acc.IsModelSupported(requestedModel) {
|
||||
if _, excluded := excludedIDs[acc.ID]; excluded {
|
||||
continue
|
||||
}
|
||||
if requestedModel != "" && !s.isModelSupportedByAccount(acc, requestedModel) {
|
||||
continue
|
||||
}
|
||||
if selected == nil {
|
||||
selected = acc
|
||||
continue
|
||||
}
|
||||
// 优先选择priority值更小的(priority值越小优先级越高)
|
||||
if acc.Priority < selected.Priority {
|
||||
selected = acc
|
||||
} else if acc.Priority == selected.Priority {
|
||||
// 优先级相同时,选最久未用的
|
||||
switch {
|
||||
case acc.LastUsedAt == nil && selected.LastUsedAt != nil:
|
||||
selected = acc
|
||||
@@ -352,6 +415,126 @@ func (s *GatewayService) SelectAccountForModel(ctx context.Context, groupID *int
|
||||
return selected, nil
|
||||
}
|
||||
|
||||
// selectAccountWithMixedScheduling 选择账户(支持混合调度)
|
||||
// 查询原生平台账户 + 启用 mixed_scheduling 的 antigravity 账户
|
||||
func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, nativePlatform string) (*Account, error) {
|
||||
platforms := []string{nativePlatform, PlatformAntigravity}
|
||||
|
||||
// 1. 查询粘性会话
|
||||
if sessionHash != "" {
|
||||
accountID, err := s.cache.GetSessionAccountID(ctx, sessionHash)
|
||||
if err == nil && accountID > 0 {
|
||||
if _, excluded := excludedIDs[accountID]; !excluded {
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
// 检查账号是否有效:原生平台直接匹配,antigravity 需要启用混合调度
|
||||
if err == nil && account.IsSchedulable() && (requestedModel == "" || s.isModelSupportedByAccount(account, requestedModel)) {
|
||||
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 获取可调度账号列表
|
||||
var accounts []Account
|
||||
var err error
|
||||
if groupID != nil {
|
||||
accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatforms(ctx, *groupID, platforms)
|
||||
} else {
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatforms(ctx, platforms)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query accounts failed: %w", err)
|
||||
}
|
||||
|
||||
// 3. 按优先级+最久未用选择(考虑模型支持和混合调度)
|
||||
var selected *Account
|
||||
for i := range accounts {
|
||||
acc := &accounts[i]
|
||||
if _, excluded := excludedIDs[acc.ID]; excluded {
|
||||
continue
|
||||
}
|
||||
// 过滤:原生平台直接通过,antigravity 需要启用混合调度
|
||||
if acc.Platform == PlatformAntigravity && !acc.IsMixedSchedulingEnabled() {
|
||||
continue
|
||||
}
|
||||
if requestedModel != "" && !s.isModelSupportedByAccount(acc, requestedModel) {
|
||||
continue
|
||||
}
|
||||
if selected == nil {
|
||||
selected = acc
|
||||
continue
|
||||
}
|
||||
if acc.Priority < selected.Priority {
|
||||
selected = acc
|
||||
} else if acc.Priority == selected.Priority {
|
||||
switch {
|
||||
case acc.LastUsedAt == nil && selected.LastUsedAt != nil:
|
||||
selected = acc
|
||||
case acc.LastUsedAt != nil && selected.LastUsedAt == nil:
|
||||
// keep selected (never used is preferred)
|
||||
case acc.LastUsedAt == nil && selected.LastUsedAt == nil:
|
||||
// keep selected (both never used)
|
||||
default:
|
||||
if acc.LastUsedAt.Before(*selected.LastUsedAt) {
|
||||
selected = acc
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if selected == nil {
|
||||
if requestedModel != "" {
|
||||
return nil, fmt.Errorf("no available accounts supporting model: %s", requestedModel)
|
||||
}
|
||||
return nil, errors.New("no available accounts")
|
||||
}
|
||||
|
||||
// 4. 建立粘性绑定
|
||||
if sessionHash != "" {
|
||||
if err := s.cache.SetSessionAccountID(ctx, sessionHash, selected.ID, stickySessionTTL); err != nil {
|
||||
log.Printf("set session account failed: session=%s account_id=%d err=%v", sessionHash, selected.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
return selected, nil
|
||||
}
|
||||
|
||||
// isModelSupportedByAccount 根据账户平台检查模型支持
|
||||
func (s *GatewayService) isModelSupportedByAccount(account *Account, requestedModel string) bool {
|
||||
if account.Platform == PlatformAntigravity {
|
||||
// Antigravity 平台使用专门的模型支持检查
|
||||
return IsAntigravityModelSupported(requestedModel)
|
||||
}
|
||||
// 其他平台使用账户的模型支持检查
|
||||
return account.IsModelSupported(requestedModel)
|
||||
}
|
||||
|
||||
// IsAntigravityModelSupported 检查 Antigravity 平台是否支持指定模型
|
||||
func IsAntigravityModelSupported(requestedModel string) bool {
|
||||
// 直接支持的模型
|
||||
if antigravitySupportedModels[requestedModel] {
|
||||
return true
|
||||
}
|
||||
// 可映射的模型
|
||||
if _, ok := antigravityModelMapping[requestedModel]; ok {
|
||||
return true
|
||||
}
|
||||
// Gemini 前缀透传
|
||||
if strings.HasPrefix(requestedModel, "gemini-") {
|
||||
return true
|
||||
}
|
||||
// Claude 模型支持(通过默认映射到 claude-sonnet-4-5)
|
||||
if strings.HasPrefix(requestedModel, "claude-") {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// GetAccessToken 获取账号凭证
|
||||
func (s *GatewayService) GetAccessToken(ctx context.Context, account *Account) (string, string, error) {
|
||||
switch account.Type {
|
||||
@@ -394,6 +577,16 @@ func (s *GatewayService) shouldRetryUpstreamError(account *Account, statusCode i
|
||||
return !account.ShouldHandleErrorCode(statusCode)
|
||||
}
|
||||
|
||||
// shouldFailoverUpstreamError determines whether an upstream error should trigger account failover.
|
||||
func (s *GatewayService) shouldFailoverUpstreamError(statusCode int) bool {
|
||||
switch statusCode {
|
||||
case 401, 403, 429, 529:
|
||||
return true
|
||||
default:
|
||||
return statusCode >= 500
|
||||
}
|
||||
}
|
||||
|
||||
// Forward 转发请求到Claude API
|
||||
func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
startTime := time.Now()
|
||||
@@ -478,9 +671,19 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
|
||||
// 处理重试耗尽的情况
|
||||
if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) {
|
||||
if s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account)
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode}
|
||||
}
|
||||
return s.handleRetryExhaustedError(ctx, resp, c, account)
|
||||
}
|
||||
|
||||
// 处理可切换账号的错误
|
||||
if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
s.handleFailoverSideEffects(ctx, resp, account)
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode}
|
||||
}
|
||||
|
||||
// 处理错误响应(不可重试的错误)
|
||||
if resp.StatusCode >= 400 {
|
||||
return s.handleErrorResponse(ctx, resp, c, account)
|
||||
@@ -492,6 +695,11 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
if req.Stream {
|
||||
streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, req.Model)
|
||||
if err != nil {
|
||||
if err.Error() == "have error in stream" {
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: 403,
|
||||
}
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
usage = streamResult.usage
|
||||
@@ -692,10 +900,7 @@ func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Res
|
||||
return nil, fmt.Errorf("upstream error: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
// handleRetryExhaustedError 处理重试耗尽后的错误
|
||||
// OAuth 403:标记账号异常
|
||||
// API Key 未配置错误码:仅返回错误,不标记账号
|
||||
func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account) (*ForwardResult, error) {
|
||||
func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account) {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
statusCode := resp.StatusCode
|
||||
|
||||
@@ -707,6 +912,18 @@ func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *ht
|
||||
// API Key 未配置错误码:不标记账号状态
|
||||
log.Printf("Account %d: upstream error %d after %d retries (not marking account)", account.ID, statusCode, maxRetries)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account) {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body)
|
||||
}
|
||||
|
||||
// handleRetryExhaustedError 处理重试耗尽后的错误
|
||||
// OAuth 403:标记账号异常
|
||||
// API Key 未配置错误码:仅返回错误,不标记账号
|
||||
func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account) (*ForwardResult, error) {
|
||||
s.handleRetryExhaustedSideEffects(ctx, resp, account)
|
||||
|
||||
// 返回统一的重试耗尽错误响应
|
||||
c.JSON(http.StatusBadGateway, gin.H{
|
||||
@@ -717,7 +934,7 @@ func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *ht
|
||||
},
|
||||
})
|
||||
|
||||
return nil, fmt.Errorf("upstream error: %d (retries exhausted)", statusCode)
|
||||
return nil, fmt.Errorf("upstream error: %d (retries exhausted)", resp.StatusCode)
|
||||
}
|
||||
|
||||
// streamingResult 流式响应结果
|
||||
@@ -757,6 +974,9 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if line == "event: error" {
|
||||
return nil, errors.New("have error in stream")
|
||||
}
|
||||
|
||||
// Extract data from SSE line (supports both "data: " and "data:" formats)
|
||||
if sseDataRe.MatchString(line) {
|
||||
@@ -1014,6 +1234,12 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
|
||||
log.Printf("Create usage log failed: %v", err)
|
||||
}
|
||||
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
||||
log.Printf("[SIMPLE MODE] Usage recorded (not billed): user=%d, tokens=%d", usageLog.UserID, usageLog.TotalTokens())
|
||||
s.deferredService.ScheduleLastUsedUpdate(account.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// 根据计费类型执行扣费
|
||||
if isSubscriptionBilling {
|
||||
// 订阅模式:更新订阅用量(使用 TotalCost 原始费用,不考虑倍率)
|
||||
@@ -1047,10 +1273,8 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
|
||||
}
|
||||
}
|
||||
|
||||
// 更新账号最后使用时间
|
||||
if err := s.accountRepo.UpdateLastUsed(ctx, account.ID); err != nil {
|
||||
log.Printf("Update last used failed: %v", err)
|
||||
}
|
||||
// Schedule batch update for account last_used_at
|
||||
s.deferredService.ScheduleLastUsedUpdate(account.ID)
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1058,6 +1282,13 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
|
||||
// ForwardCountTokens 转发 count_tokens 请求到上游 API
|
||||
// 特点:不记录使用量、仅支持非流式响应
|
||||
func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, account *Account, body []byte) error {
|
||||
// Antigravity 账户不支持 count_tokens 转发,返回估算值
|
||||
// 参考 Antigravity-Manager 和 proxycast 实现
|
||||
if account.Platform == PlatformAntigravity {
|
||||
c.JSON(http.StatusOK, gin.H{"input_tokens": 100})
|
||||
return nil
|
||||
}
|
||||
|
||||
// 应用模型映射(仅对 apikey 类型账号)
|
||||
if account.Type == AccountTypeApiKey {
|
||||
var req struct {
|
||||
|
||||
@@ -18,6 +18,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
|
||||
|
||||
@@ -33,26 +34,32 @@ const (
|
||||
)
|
||||
|
||||
type GeminiMessagesCompatService struct {
|
||||
accountRepo AccountRepository
|
||||
cache GatewayCache
|
||||
tokenProvider *GeminiTokenProvider
|
||||
rateLimitService *RateLimitService
|
||||
httpUpstream HTTPUpstream
|
||||
accountRepo AccountRepository
|
||||
groupRepo GroupRepository
|
||||
cache GatewayCache
|
||||
tokenProvider *GeminiTokenProvider
|
||||
rateLimitService *RateLimitService
|
||||
httpUpstream HTTPUpstream
|
||||
antigravityGatewayService *AntigravityGatewayService
|
||||
}
|
||||
|
||||
func NewGeminiMessagesCompatService(
|
||||
accountRepo AccountRepository,
|
||||
groupRepo GroupRepository,
|
||||
cache GatewayCache,
|
||||
tokenProvider *GeminiTokenProvider,
|
||||
rateLimitService *RateLimitService,
|
||||
httpUpstream HTTPUpstream,
|
||||
antigravityGatewayService *AntigravityGatewayService,
|
||||
) *GeminiMessagesCompatService {
|
||||
return &GeminiMessagesCompatService{
|
||||
accountRepo: accountRepo,
|
||||
cache: cache,
|
||||
tokenProvider: tokenProvider,
|
||||
rateLimitService: rateLimitService,
|
||||
httpUpstream: httpUpstream,
|
||||
accountRepo: accountRepo,
|
||||
groupRepo: groupRepo,
|
||||
cache: cache,
|
||||
tokenProvider: tokenProvider,
|
||||
rateLimitService: rateLimitService,
|
||||
httpUpstream: httpUpstream,
|
||||
antigravityGatewayService: antigravityGatewayService,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -62,24 +69,75 @@ func (s *GeminiMessagesCompatService) GetTokenProvider() *GeminiTokenProvider {
|
||||
}
|
||||
|
||||
func (s *GeminiMessagesCompatService) SelectAccountForModel(ctx context.Context, groupID *int64, sessionHash string, requestedModel string) (*Account, error) {
|
||||
return s.SelectAccountForModelWithExclusions(ctx, groupID, sessionHash, requestedModel, nil)
|
||||
}
|
||||
|
||||
func (s *GeminiMessagesCompatService) SelectAccountForModelWithExclusions(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*Account, error) {
|
||||
// 优先检查 context 中的强制平台(/antigravity 路由)
|
||||
var platform string
|
||||
forcePlatform, hasForcePlatform := ctx.Value(ctxkey.ForcePlatform).(string)
|
||||
if hasForcePlatform && forcePlatform != "" {
|
||||
platform = forcePlatform
|
||||
} else if groupID != nil {
|
||||
// 根据分组 platform 决定查询哪种账号
|
||||
group, err := s.groupRepo.GetByID(ctx, *groupID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get group failed: %w", err)
|
||||
}
|
||||
platform = group.Platform
|
||||
} else {
|
||||
// 无分组时只使用原生 gemini 平台
|
||||
platform = PlatformGemini
|
||||
}
|
||||
|
||||
// gemini 分组支持混合调度(包含启用了 mixed_scheduling 的 antigravity 账户)
|
||||
// 注意:强制平台模式不走混合调度
|
||||
useMixedScheduling := platform == PlatformGemini && !hasForcePlatform
|
||||
var queryPlatforms []string
|
||||
if useMixedScheduling {
|
||||
queryPlatforms = []string{PlatformGemini, PlatformAntigravity}
|
||||
} else {
|
||||
queryPlatforms = []string{platform}
|
||||
}
|
||||
|
||||
cacheKey := "gemini:" + sessionHash
|
||||
|
||||
if sessionHash != "" {
|
||||
accountID, err := s.cache.GetSessionAccountID(ctx, cacheKey)
|
||||
if err == nil && accountID > 0 {
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err == nil && account.IsSchedulable() && account.Platform == PlatformGemini && (requestedModel == "" || account.IsModelSupported(requestedModel)) {
|
||||
_ = s.cache.RefreshSessionTTL(ctx, cacheKey, geminiStickySessionTTL)
|
||||
return account, nil
|
||||
if _, excluded := excludedIDs[accountID]; !excluded {
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
// 检查账号是否有效:原生平台直接匹配,antigravity 需要启用混合调度
|
||||
if err == nil && account.IsSchedulable() && (requestedModel == "" || s.isModelSupportedByAccount(account, requestedModel)) {
|
||||
valid := false
|
||||
if account.Platform == platform {
|
||||
valid = true
|
||||
} else if useMixedScheduling && account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled() {
|
||||
valid = true
|
||||
}
|
||||
if valid {
|
||||
_ = s.cache.RefreshSessionTTL(ctx, cacheKey, geminiStickySessionTTL)
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 查询可调度账户(强制平台模式:优先按分组查找,找不到再查全部)
|
||||
var accounts []Account
|
||||
var err error
|
||||
if groupID != nil {
|
||||
accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, PlatformGemini)
|
||||
accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatforms(ctx, *groupID, queryPlatforms)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query accounts failed: %w", err)
|
||||
}
|
||||
// 强制平台模式下,分组中找不到账户时回退查询全部
|
||||
if len(accounts) == 0 && hasForcePlatform {
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatforms(ctx, queryPlatforms)
|
||||
}
|
||||
} else {
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatform(ctx, PlatformGemini)
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatforms(ctx, queryPlatforms)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query accounts failed: %w", err)
|
||||
@@ -88,7 +146,15 @@ func (s *GeminiMessagesCompatService) SelectAccountForModel(ctx context.Context,
|
||||
var selected *Account
|
||||
for i := range accounts {
|
||||
acc := &accounts[i]
|
||||
if requestedModel != "" && !acc.IsModelSupported(requestedModel) {
|
||||
if _, excluded := excludedIDs[acc.ID]; excluded {
|
||||
continue
|
||||
}
|
||||
// 混合调度模式下:原生平台直接通过,antigravity 需要启用 mixed_scheduling
|
||||
// 非混合调度模式(antigravity 分组):不需要过滤
|
||||
if useMixedScheduling && acc.Platform == PlatformAntigravity && !acc.IsMixedSchedulingEnabled() {
|
||||
continue
|
||||
}
|
||||
if requestedModel != "" && !s.isModelSupportedByAccount(acc, requestedModel) {
|
||||
continue
|
||||
}
|
||||
if selected == nil {
|
||||
@@ -130,6 +196,34 @@ func (s *GeminiMessagesCompatService) SelectAccountForModel(ctx context.Context,
|
||||
return selected, nil
|
||||
}
|
||||
|
||||
// isModelSupportedByAccount 根据账户平台检查模型支持
|
||||
func (s *GeminiMessagesCompatService) isModelSupportedByAccount(account *Account, requestedModel string) bool {
|
||||
if account.Platform == PlatformAntigravity {
|
||||
return IsAntigravityModelSupported(requestedModel)
|
||||
}
|
||||
return account.IsModelSupported(requestedModel)
|
||||
}
|
||||
|
||||
// GetAntigravityGatewayService 返回 AntigravityGatewayService
|
||||
func (s *GeminiMessagesCompatService) GetAntigravityGatewayService() *AntigravityGatewayService {
|
||||
return s.antigravityGatewayService
|
||||
}
|
||||
|
||||
// HasAntigravityAccounts 检查是否有可用的 antigravity 账户
|
||||
func (s *GeminiMessagesCompatService) HasAntigravityAccounts(ctx context.Context, groupID *int64) (bool, error) {
|
||||
var accounts []Account
|
||||
var err error
|
||||
if groupID != nil {
|
||||
accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, PlatformAntigravity)
|
||||
} else {
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatform(ctx, PlatformAntigravity)
|
||||
}
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return len(accounts) > 0, nil
|
||||
}
|
||||
|
||||
// SelectAccountForAIStudioEndpoints selects an account that is likely to succeed against
|
||||
// generativelanguage.googleapis.com (e.g. GET /v1beta/models).
|
||||
//
|
||||
@@ -425,6 +519,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) {
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode}
|
||||
}
|
||||
return nil, s.writeGeminiMappedError(c, resp.StatusCode, respBody)
|
||||
}
|
||||
|
||||
@@ -724,6 +821,10 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
}, nil
|
||||
}
|
||||
|
||||
if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) {
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode}
|
||||
}
|
||||
|
||||
respBody = unwrapIfNeeded(isOAuth, respBody)
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
if contentType == "" {
|
||||
@@ -795,6 +896,15 @@ func (s *GeminiMessagesCompatService) shouldRetryGeminiUpstreamError(account *Ac
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GeminiMessagesCompatService) shouldFailoverGeminiUpstreamError(statusCode int) bool {
|
||||
switch statusCode {
|
||||
case 401, 403, 429, 529:
|
||||
return true
|
||||
default:
|
||||
return statusCode >= 500
|
||||
}
|
||||
}
|
||||
|
||||
func sleepGeminiBackoff(attempt int) {
|
||||
delay := geminiRetryBaseDelay * time.Duration(1<<uint(attempt-1))
|
||||
if delay > geminiRetryMaxDelay {
|
||||
@@ -1773,7 +1883,7 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
|
||||
if statusCode != 429 {
|
||||
return
|
||||
}
|
||||
resetAt := parseGeminiRateLimitResetTime(body)
|
||||
resetAt := ParseGeminiRateLimitResetTime(body)
|
||||
if resetAt == nil {
|
||||
ra := time.Now().Add(5 * time.Minute)
|
||||
_ = s.accountRepo.SetRateLimited(ctx, account.ID, ra)
|
||||
@@ -1782,7 +1892,8 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
|
||||
_ = s.accountRepo.SetRateLimited(ctx, account.ID, time.Unix(*resetAt, 0))
|
||||
}
|
||||
|
||||
func parseGeminiRateLimitResetTime(body []byte) *int64 {
|
||||
// ParseGeminiRateLimitResetTime 解析 Gemini 格式的 429 响应,返回重置时间的 Unix 时间戳
|
||||
func ParseGeminiRateLimitResetTime(body []byte) *int64 {
|
||||
// Try to parse metadata.quotaResetDelay like "12.345s"
|
||||
var parsed map[string]any
|
||||
if err := json.Unmarshal(body, &parsed); err == nil {
|
||||
|
||||
493
backend/internal/service/gemini_multiplatform_test.go
Normal file
493
backend/internal/service/gemini_multiplatform_test.go
Normal file
@@ -0,0 +1,493 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// mockAccountRepoForGemini Gemini 测试用的 mock
|
||||
type mockAccountRepoForGemini struct {
|
||||
accounts []Account
|
||||
accountsByID map[int64]*Account
|
||||
}
|
||||
|
||||
func (m *mockAccountRepoForGemini) GetByID(ctx context.Context, id int64) (*Account, error) {
|
||||
if acc, ok := m.accountsByID[id]; ok {
|
||||
return acc, nil
|
||||
}
|
||||
return nil, errors.New("account not found")
|
||||
}
|
||||
|
||||
func (m *mockAccountRepoForGemini) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
|
||||
var result []Account
|
||||
for _, acc := range m.accounts {
|
||||
if acc.Platform == platform && acc.IsSchedulable() {
|
||||
result = append(result, acc)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (m *mockAccountRepoForGemini) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
|
||||
// 测试时不区分 groupID,直接按 platform 过滤
|
||||
return m.ListSchedulableByPlatform(ctx, platform)
|
||||
}
|
||||
|
||||
// Stub methods to implement AccountRepository interface
|
||||
func (m *mockAccountRepoForGemini) Create(ctx context.Context, account *Account) error { return nil }
|
||||
func (m *mockAccountRepoForGemini) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) {
|
||||
return nil, 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) List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string) ([]Account, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListActive(ctx context.Context) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListByPlatform(ctx context.Context, platform string) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) UpdateLastUsed(ctx context.Context, id int64) error { return nil }
|
||||
func (m *mockAccountRepoForGemini) BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) SetError(ctx context.Context, id int64, errorMsg string) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) SetSchedulable(ctx context.Context, id int64, schedulable bool) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) BindGroups(ctx context.Context, accountID int64, groupIDs []int64) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListSchedulable(ctx context.Context) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]Account, error) {
|
||||
var result []Account
|
||||
platformSet := make(map[string]bool)
|
||||
for _, p := range platforms {
|
||||
platformSet[p] = true
|
||||
}
|
||||
for _, acc := range m.accounts {
|
||||
if platformSet[acc.Platform] && acc.IsSchedulable() {
|
||||
result = append(result, acc)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error) {
|
||||
return m.ListSchedulableByPlatforms(ctx, platforms)
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) SetOverloaded(ctx context.Context, id int64, until time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ClearRateLimit(ctx context.Context, id int64) error { return nil }
|
||||
func (m *mockAccountRepoForGemini) UpdateSessionWindow(ctx context.Context, id int64, start, end *time.Time, status string) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) BulkUpdate(ctx context.Context, ids []int64, updates AccountBulkUpdate) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Verify interface implementation
|
||||
var _ AccountRepository = (*mockAccountRepoForGemini)(nil)
|
||||
|
||||
// mockGroupRepoForGemini Gemini 测试用的 group repo mock
|
||||
type mockGroupRepoForGemini struct {
|
||||
groups map[int64]*Group
|
||||
}
|
||||
|
||||
func (m *mockGroupRepoForGemini) GetByID(ctx context.Context, id int64) (*Group, error) {
|
||||
if g, ok := m.groups[id]; ok {
|
||||
return g, nil
|
||||
}
|
||||
return nil, errors.New("group not found")
|
||||
}
|
||||
|
||||
// Stub methods to implement GroupRepository interface
|
||||
func (m *mockGroupRepoForGemini) Create(ctx context.Context, group *Group) error { return nil }
|
||||
func (m *mockGroupRepoForGemini) Update(ctx context.Context, group *Group) error { return nil }
|
||||
func (m *mockGroupRepoForGemini) Delete(ctx context.Context, id int64) error { return nil }
|
||||
func (m *mockGroupRepoForGemini) DeleteCascade(ctx context.Context, id int64) ([]int64, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockGroupRepoForGemini) List(ctx context.Context, params pagination.PaginationParams) ([]Group, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockGroupRepoForGemini) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status string, isExclusive *bool) ([]Group, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockGroupRepoForGemini) ListActive(ctx context.Context) ([]Group, error) { return nil, nil }
|
||||
func (m *mockGroupRepoForGemini) ListActiveByPlatform(ctx context.Context, platform string) ([]Group, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockGroupRepoForGemini) ExistsByName(ctx context.Context, name string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
func (m *mockGroupRepoForGemini) GetAccountCount(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (m *mockGroupRepoForGemini) DeleteAccountGroupsByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
var _ GroupRepository = (*mockGroupRepoForGemini)(nil)
|
||||
|
||||
// mockGatewayCacheForGemini Gemini 测试用的 cache mock
|
||||
type mockGatewayCacheForGemini struct {
|
||||
sessionBindings map[string]int64
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForGemini) GetSessionAccountID(ctx context.Context, sessionHash string) (int64, error) {
|
||||
if id, ok := m.sessionBindings[sessionHash]; ok {
|
||||
return id, nil
|
||||
}
|
||||
return 0, errors.New("not found")
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForGemini) SetSessionAccountID(ctx context.Context, sessionHash string, accountID int64, ttl time.Duration) error {
|
||||
if m.sessionBindings == nil {
|
||||
m.sessionBindings = make(map[string]int64)
|
||||
}
|
||||
m.sessionBindings[sessionHash] = accountID
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForGemini) RefreshSessionTTL(ctx context.Context, sessionHash string, ttl time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择
|
||||
func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
repo := &mockAccountRepoForGemini{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformGemini, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
{ID: 3, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true}, // 应被隔离
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForGemini{}
|
||||
groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
|
||||
|
||||
svc := &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
groupRepo: groupRepo,
|
||||
cache: cache,
|
||||
}
|
||||
|
||||
// 无分组时使用 gemini 平台
|
||||
acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(1), acc.ID, "应选择优先级最高的 gemini 账户")
|
||||
require.Equal(t, PlatformGemini, acc.Platform, "无分组时应只返回 gemini 平台账户")
|
||||
}
|
||||
|
||||
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_AntigravityGroup 测试 antigravity 分组
|
||||
func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_AntigravityGroup(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
repo := &mockAccountRepoForGemini{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true}, // 应被隔离
|
||||
{ID: 2, Platform: PlatformAntigravity, Priority: 1, Status: StatusActive, Schedulable: true}, // 应被选择
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForGemini{}
|
||||
groupRepo := &mockGroupRepoForGemini{
|
||||
groups: map[int64]*Group{
|
||||
1: {ID: 1, Platform: PlatformAntigravity},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
groupRepo: groupRepo,
|
||||
cache: cache,
|
||||
}
|
||||
|
||||
groupID := int64(1)
|
||||
acc, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gemini-2.5-flash", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID)
|
||||
require.Equal(t, PlatformAntigravity, acc.Platform, "antigravity 分组应只返回 antigravity 账户")
|
||||
}
|
||||
|
||||
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_OAuthPreferred 测试 OAuth 优先
|
||||
func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_OAuthPreferred(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
repo := &mockAccountRepoForGemini{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformGemini, Type: AccountTypeApiKey, Priority: 1, Status: StatusActive, Schedulable: true, LastUsedAt: nil},
|
||||
{ID: 2, Platform: PlatformGemini, Type: AccountTypeOAuth, Priority: 1, Status: StatusActive, Schedulable: true, LastUsedAt: nil},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForGemini{}
|
||||
groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
|
||||
|
||||
svc := &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
groupRepo: groupRepo,
|
||||
cache: cache,
|
||||
}
|
||||
|
||||
acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID, "同优先级且都未使用时,应优先选择 OAuth 账户")
|
||||
require.Equal(t, AccountTypeOAuth, acc.Type)
|
||||
}
|
||||
|
||||
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_NoAvailableAccounts 测试无可用账户
|
||||
func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_NoAvailableAccounts(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
repo := &mockAccountRepoForGemini{
|
||||
accounts: []Account{},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForGemini{}
|
||||
groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
|
||||
|
||||
svc := &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
groupRepo: groupRepo,
|
||||
cache: cache,
|
||||
}
|
||||
|
||||
acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gemini-2.5-flash", nil)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, acc)
|
||||
require.Contains(t, err.Error(), "no available")
|
||||
}
|
||||
|
||||
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_StickySession 测试粘性会话
|
||||
func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_StickySession(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("粘性会话命中-同平台", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForGemini{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformGemini, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
// 注意:缓存键使用 "gemini:" 前缀
|
||||
cache := &mockGatewayCacheForGemini{
|
||||
sessionBindings: map[string]int64{"gemini:session-123": 1},
|
||||
}
|
||||
groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
|
||||
|
||||
svc := &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
groupRepo: groupRepo,
|
||||
cache: cache,
|
||||
}
|
||||
|
||||
acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "session-123", "gemini-2.5-flash", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(1), acc.ID, "应返回粘性会话绑定的账户")
|
||||
})
|
||||
|
||||
t.Run("粘性会话平台不匹配-降级选择", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForGemini{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformAntigravity, Priority: 2, Status: StatusActive, Schedulable: true}, // 粘性会话绑定
|
||||
{ID: 2, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
cache := &mockGatewayCacheForGemini{
|
||||
sessionBindings: map[string]int64{"gemini:session-123": 1}, // 绑定 antigravity 账户
|
||||
}
|
||||
groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
|
||||
|
||||
svc := &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
groupRepo: groupRepo,
|
||||
cache: cache,
|
||||
}
|
||||
|
||||
// 无分组时使用 gemini 平台,粘性会话绑定的 antigravity 账户平台不匹配
|
||||
acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "session-123", "gemini-2.5-flash", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
require.Equal(t, int64(2), acc.ID, "粘性会话账户平台不匹配,应降级选择 gemini 账户")
|
||||
require.Equal(t, PlatformGemini, acc.Platform)
|
||||
})
|
||||
|
||||
t.Run("粘性会话不命中无前缀缓存键", func(t *testing.T) {
|
||||
repo := &mockAccountRepoForGemini{
|
||||
accounts: []Account{
|
||||
{ID: 1, Platform: PlatformGemini, Priority: 2, Status: StatusActive, Schedulable: true},
|
||||
{ID: 2, Platform: PlatformGemini, Priority: 1, Status: StatusActive, Schedulable: true},
|
||||
},
|
||||
accountsByID: map[int64]*Account{},
|
||||
}
|
||||
for i := range repo.accounts {
|
||||
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
|
||||
}
|
||||
|
||||
// 缓存键没有 "gemini:" 前缀,不应命中
|
||||
cache := &mockGatewayCacheForGemini{
|
||||
sessionBindings: map[string]int64{"session-123": 1},
|
||||
}
|
||||
groupRepo := &mockGroupRepoForGemini{groups: map[int64]*Group{}}
|
||||
|
||||
svc := &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
groupRepo: groupRepo,
|
||||
cache: cache,
|
||||
}
|
||||
|
||||
acc, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "session-123", "gemini-2.5-flash", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, acc)
|
||||
// 粘性会话未命中,按优先级选择
|
||||
require.Equal(t, int64(2), acc.ID, "粘性会话未命中,应按优先级选择")
|
||||
})
|
||||
}
|
||||
|
||||
// TestGeminiPlatformRouting_DocumentRouteDecision 测试平台路由决策逻辑
|
||||
func TestGeminiPlatformRouting_DocumentRouteDecision(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
platform string
|
||||
expectedService string // "gemini" 表示 ForwardNative, "antigravity" 表示 ForwardGemini
|
||||
}{
|
||||
{
|
||||
name: "Gemini平台走ForwardNative",
|
||||
platform: PlatformGemini,
|
||||
expectedService: "gemini",
|
||||
},
|
||||
{
|
||||
name: "Antigravity平台走ForwardGemini",
|
||||
platform: PlatformAntigravity,
|
||||
expectedService: "antigravity",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := &Account{Platform: tt.platform}
|
||||
|
||||
// 模拟 Handler 层的路由逻辑
|
||||
var serviceName string
|
||||
if account.Platform == PlatformAntigravity {
|
||||
serviceName = "antigravity"
|
||||
} else {
|
||||
serviceName = "gemini"
|
||||
}
|
||||
|
||||
require.Equal(t, tt.expectedService, serviceName,
|
||||
"平台 %s 应该路由到 %s 服务", tt.platform, tt.expectedService)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiMessagesCompatService_isModelSupportedByAccount(t *testing.T) {
|
||||
svc := &GeminiMessagesCompatService{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
model string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "Antigravity平台-支持gemini模型",
|
||||
account: &Account{Platform: PlatformAntigravity},
|
||||
model: "gemini-2.5-flash",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "Antigravity平台-支持claude模型",
|
||||
account: &Account{Platform: PlatformAntigravity},
|
||||
model: "claude-3-5-sonnet-20241022",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "Antigravity平台-不支持gpt模型",
|
||||
account: &Account{Platform: PlatformAntigravity},
|
||||
model: "gpt-4",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "Gemini平台-无映射配置-支持所有模型",
|
||||
account: &Account{Platform: PlatformGemini},
|
||||
model: "gemini-2.5-flash",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "Gemini平台-有映射配置-只支持配置的模型",
|
||||
account: &Account{
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"gemini-1.5-pro": "x"}},
|
||||
},
|
||||
model: "gemini-2.5-flash",
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := svc.isModelSupportedByAccount(tt.account, tt.model)
|
||||
require.Equal(t, tt.expected, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -20,7 +20,7 @@ type OpenAIOAuthClient interface {
|
||||
type ClaudeOAuthClient interface {
|
||||
GetOrganizationUUID(ctx context.Context, sessionKey, proxyURL string) (string, error)
|
||||
GetAuthorizationCode(ctx context.Context, sessionKey, orgUUID, scope, codeChallenge, state, proxyURL string) (string, error)
|
||||
ExchangeCodeForToken(ctx context.Context, code, codeVerifier, state, proxyURL string) (*oauth.TokenResponse, error)
|
||||
ExchangeCodeForToken(ctx context.Context, code, codeVerifier, state, proxyURL string, isSetupToken bool) (*oauth.TokenResponse, error)
|
||||
RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*oauth.TokenResponse, error)
|
||||
}
|
||||
|
||||
@@ -142,8 +142,11 @@ func (s *OAuthService) ExchangeCode(ctx context.Context, input *ExchangeCodeInpu
|
||||
}
|
||||
}
|
||||
|
||||
// Determine if this is a setup token (scope is inference only)
|
||||
isSetupToken := session.Scope == oauth.ScopeInference
|
||||
|
||||
// Exchange code for token
|
||||
tokenInfo, err := s.exchangeCodeForToken(ctx, input.Code, session.CodeVerifier, session.State, proxyURL)
|
||||
tokenInfo, err := s.exchangeCodeForToken(ctx, input.Code, session.CodeVerifier, session.State, proxyURL, isSetupToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -172,10 +175,12 @@ func (s *OAuthService) CookieAuth(ctx context.Context, input *CookieAuthInput) (
|
||||
}
|
||||
}
|
||||
|
||||
// Determine scope
|
||||
// Determine scope and if this is a setup token
|
||||
scope := fmt.Sprintf("%s %s", oauth.ScopeProfile, oauth.ScopeInference)
|
||||
isSetupToken := false
|
||||
if input.Scope == "inference" {
|
||||
scope = oauth.ScopeInference
|
||||
isSetupToken = true
|
||||
}
|
||||
|
||||
// Step 1: Get organization info using sessionKey
|
||||
@@ -203,7 +208,7 @@ func (s *OAuthService) CookieAuth(ctx context.Context, input *CookieAuthInput) (
|
||||
}
|
||||
|
||||
// Step 4: Exchange code for token
|
||||
tokenInfo, err := s.exchangeCodeForToken(ctx, authCode, codeVerifier, state, proxyURL)
|
||||
tokenInfo, err := s.exchangeCodeForToken(ctx, authCode, codeVerifier, state, proxyURL, isSetupToken)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to exchange code: %w", err)
|
||||
}
|
||||
@@ -228,8 +233,8 @@ func (s *OAuthService) getAuthorizationCode(ctx context.Context, sessionKey, org
|
||||
}
|
||||
|
||||
// exchangeCodeForToken exchanges authorization code for tokens
|
||||
func (s *OAuthService) exchangeCodeForToken(ctx context.Context, code, codeVerifier, state, proxyURL string) (*TokenInfo, error) {
|
||||
tokenResp, err := s.oauthClient.ExchangeCodeForToken(ctx, code, codeVerifier, state, proxyURL)
|
||||
func (s *OAuthService) exchangeCodeForToken(ctx context.Context, code, codeVerifier, state, proxyURL string, isSetupToken bool) (*TokenInfo, error) {
|
||||
tokenResp, err := s.oauthClient.ExchangeCodeForToken(ctx, code, codeVerifier, state, proxyURL, isSetupToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strconv"
|
||||
@@ -83,6 +84,7 @@ type OpenAIGatewayService struct {
|
||||
rateLimitService *RateLimitService
|
||||
billingCacheService *BillingCacheService
|
||||
httpUpstream HTTPUpstream
|
||||
deferredService *DeferredService
|
||||
}
|
||||
|
||||
// NewOpenAIGatewayService creates a new OpenAIGatewayService
|
||||
@@ -97,6 +99,7 @@ func NewOpenAIGatewayService(
|
||||
rateLimitService *RateLimitService,
|
||||
billingCacheService *BillingCacheService,
|
||||
httpUpstream HTTPUpstream,
|
||||
deferredService *DeferredService,
|
||||
) *OpenAIGatewayService {
|
||||
return &OpenAIGatewayService{
|
||||
accountRepo: accountRepo,
|
||||
@@ -109,6 +112,7 @@ func NewOpenAIGatewayService(
|
||||
rateLimitService: rateLimitService,
|
||||
billingCacheService: billingCacheService,
|
||||
httpUpstream: httpUpstream,
|
||||
deferredService: deferredService,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -129,15 +133,22 @@ func (s *OpenAIGatewayService) SelectAccount(ctx context.Context, groupID *int64
|
||||
|
||||
// SelectAccountForModel selects an account supporting the requested model
|
||||
func (s *OpenAIGatewayService) SelectAccountForModel(ctx context.Context, groupID *int64, sessionHash string, requestedModel string) (*Account, error) {
|
||||
return s.SelectAccountForModelWithExclusions(ctx, groupID, sessionHash, requestedModel, nil)
|
||||
}
|
||||
|
||||
// SelectAccountForModelWithExclusions selects an account supporting the requested model while excluding specified accounts.
|
||||
func (s *OpenAIGatewayService) SelectAccountForModelWithExclusions(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*Account, error) {
|
||||
// 1. Check sticky session
|
||||
if sessionHash != "" {
|
||||
accountID, err := s.cache.GetSessionAccountID(ctx, "openai:"+sessionHash)
|
||||
if err == nil && accountID > 0 {
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err == nil && account.IsSchedulable() && account.IsOpenAI() && (requestedModel == "" || account.IsModelSupported(requestedModel)) {
|
||||
// Refresh sticky session TTL
|
||||
_ = s.cache.RefreshSessionTTL(ctx, "openai:"+sessionHash, openaiStickySessionTTL)
|
||||
return account, nil
|
||||
if _, excluded := excludedIDs[accountID]; !excluded {
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err == nil && account.IsSchedulable() && account.IsOpenAI() && (requestedModel == "" || account.IsModelSupported(requestedModel)) {
|
||||
// Refresh sticky session TTL
|
||||
_ = s.cache.RefreshSessionTTL(ctx, "openai:"+sessionHash, openaiStickySessionTTL)
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -145,7 +156,10 @@ func (s *OpenAIGatewayService) SelectAccountForModel(ctx context.Context, groupI
|
||||
// 2. Get schedulable OpenAI accounts
|
||||
var accounts []Account
|
||||
var err error
|
||||
if groupID != nil {
|
||||
// 简易模式:忽略分组限制,查询所有可用账号
|
||||
if s.cfg.RunMode == config.RunModeSimple {
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatform(ctx, PlatformOpenAI)
|
||||
} else if groupID != nil {
|
||||
accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, PlatformOpenAI)
|
||||
} else {
|
||||
accounts, err = s.accountRepo.ListSchedulableByPlatform(ctx, PlatformOpenAI)
|
||||
@@ -158,6 +172,9 @@ func (s *OpenAIGatewayService) SelectAccountForModel(ctx context.Context, groupI
|
||||
var selected *Account
|
||||
for i := range accounts {
|
||||
acc := &accounts[i]
|
||||
if _, excluded := excludedIDs[acc.ID]; excluded {
|
||||
continue
|
||||
}
|
||||
// Check model support
|
||||
if requestedModel != "" && !acc.IsModelSupported(requestedModel) {
|
||||
continue
|
||||
@@ -221,6 +238,20 @@ func (s *OpenAIGatewayService) GetAccessToken(ctx context.Context, account *Acco
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) shouldFailoverUpstreamError(statusCode int) bool {
|
||||
switch statusCode {
|
||||
case 401, 403, 429, 529:
|
||||
return true
|
||||
default:
|
||||
return statusCode >= 500
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account) {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body)
|
||||
}
|
||||
|
||||
// Forward forwards request to OpenAI API
|
||||
func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) {
|
||||
startTime := time.Now()
|
||||
@@ -288,6 +319,10 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
|
||||
// Handle error response
|
||||
if resp.StatusCode >= 400 {
|
||||
if s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
s.handleFailoverSideEffects(ctx, resp, account)
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode}
|
||||
}
|
||||
return s.handleErrorResponse(ctx, resp, c, account)
|
||||
}
|
||||
|
||||
@@ -723,6 +758,12 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
|
||||
|
||||
_ = s.usageLogRepo.Create(ctx, usageLog)
|
||||
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
||||
log.Printf("[SIMPLE MODE] Usage recorded (not billed): user=%d, tokens=%d", usageLog.UserID, usageLog.TotalTokens())
|
||||
s.deferredService.ScheduleLastUsedUpdate(account.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Deduct based on billing type
|
||||
if isSubscriptionBilling {
|
||||
if cost.TotalCost > 0 {
|
||||
@@ -744,8 +785,8 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
|
||||
}
|
||||
}
|
||||
|
||||
// Update account last used
|
||||
_ = s.accountRepo.UpdateLastUsed(ctx, account.ID)
|
||||
// Schedule batch update for account last_used_at
|
||||
s.deferredService.ScheduleLastUsedUpdate(account.ID)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -10,6 +10,13 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
)
|
||||
|
||||
// MaxExpiresAt is the maximum allowed expiration date (year 2099)
|
||||
// This prevents time.Time JSON serialization errors (RFC 3339 requires year <= 9999)
|
||||
var MaxExpiresAt = time.Date(2099, 12, 31, 23, 59, 59, 0, time.UTC)
|
||||
|
||||
// MaxValidityDays is the maximum allowed validity days for subscriptions (100 years)
|
||||
const MaxValidityDays = 36500
|
||||
|
||||
var (
|
||||
ErrSubscriptionNotFound = infraerrors.NotFound("SUBSCRIPTION_NOT_FOUND", "subscription not found")
|
||||
ErrSubscriptionExpired = infraerrors.Forbidden("SUBSCRIPTION_EXPIRED", "subscription has expired")
|
||||
@@ -111,6 +118,9 @@ func (s *SubscriptionService) AssignOrExtendSubscription(ctx context.Context, in
|
||||
if validityDays <= 0 {
|
||||
validityDays = 30
|
||||
}
|
||||
if validityDays > MaxValidityDays {
|
||||
validityDays = MaxValidityDays
|
||||
}
|
||||
|
||||
// 已有订阅,执行续期
|
||||
if existingSub != nil {
|
||||
@@ -125,6 +135,11 @@ func (s *SubscriptionService) AssignOrExtendSubscription(ctx context.Context, in
|
||||
newExpiresAt = now.AddDate(0, 0, validityDays)
|
||||
}
|
||||
|
||||
// 确保不超过最大过期时间
|
||||
if newExpiresAt.After(MaxExpiresAt) {
|
||||
newExpiresAt = MaxExpiresAt
|
||||
}
|
||||
|
||||
// 更新过期时间
|
||||
if err := s.userSubRepo.ExtendExpiry(ctx, existingSub.ID, newExpiresAt); err != nil {
|
||||
return nil, false, fmt.Errorf("extend subscription: %w", err)
|
||||
@@ -189,13 +204,21 @@ func (s *SubscriptionService) createSubscription(ctx context.Context, input *Ass
|
||||
if validityDays <= 0 {
|
||||
validityDays = 30
|
||||
}
|
||||
if validityDays > MaxValidityDays {
|
||||
validityDays = MaxValidityDays
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
expiresAt := now.AddDate(0, 0, validityDays)
|
||||
if expiresAt.After(MaxExpiresAt) {
|
||||
expiresAt = MaxExpiresAt
|
||||
}
|
||||
|
||||
sub := &UserSubscription{
|
||||
UserID: input.UserID,
|
||||
GroupID: input.GroupID,
|
||||
StartsAt: now,
|
||||
ExpiresAt: now.AddDate(0, 0, validityDays),
|
||||
ExpiresAt: expiresAt,
|
||||
Status: SubscriptionStatusActive,
|
||||
AssignedAt: now,
|
||||
Notes: input.Notes,
|
||||
@@ -291,8 +314,17 @@ func (s *SubscriptionService) ExtendSubscription(ctx context.Context, subscripti
|
||||
return nil, ErrSubscriptionNotFound
|
||||
}
|
||||
|
||||
// 限制延长天数
|
||||
if days > MaxValidityDays {
|
||||
days = MaxValidityDays
|
||||
}
|
||||
|
||||
// 计算新的过期时间
|
||||
newExpiresAt := sub.ExpiresAt.AddDate(0, 0, days)
|
||||
if newExpiresAt.After(MaxExpiresAt) {
|
||||
newExpiresAt = MaxExpiresAt
|
||||
}
|
||||
|
||||
if err := s.userSubRepo.ExtendExpiry(ctx, subscriptionID, newExpiresAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
63
backend/internal/service/timing_wheel_service.go
Normal file
63
backend/internal/service/timing_wheel_service.go
Normal file
@@ -0,0 +1,63 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/collection"
|
||||
)
|
||||
|
||||
// TimingWheelService wraps go-zero's TimingWheel for task scheduling
|
||||
type TimingWheelService struct {
|
||||
tw *collection.TimingWheel
|
||||
stopOnce sync.Once
|
||||
}
|
||||
|
||||
// NewTimingWheelService creates a new TimingWheelService instance
|
||||
func NewTimingWheelService() *TimingWheelService {
|
||||
// 1 second tick, 3600 slots = supports up to 1 hour delay
|
||||
// execute function: runs func() type tasks
|
||||
tw, err := collection.NewTimingWheel(1*time.Second, 3600, func(key, value any) {
|
||||
if fn, ok := value.(func()); ok {
|
||||
fn()
|
||||
}
|
||||
})
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return &TimingWheelService{tw: tw}
|
||||
}
|
||||
|
||||
// Start starts the timing wheel
|
||||
func (s *TimingWheelService) Start() {
|
||||
log.Println("[TimingWheel] Started (auto-start by go-zero)")
|
||||
}
|
||||
|
||||
// Stop stops the timing wheel
|
||||
func (s *TimingWheelService) Stop() {
|
||||
s.stopOnce.Do(func() {
|
||||
s.tw.Stop()
|
||||
log.Println("[TimingWheel] Stopped")
|
||||
})
|
||||
}
|
||||
|
||||
// Schedule schedules a one-time task
|
||||
func (s *TimingWheelService) Schedule(name string, delay time.Duration, fn func()) {
|
||||
_ = s.tw.SetTimer(name, fn, delay)
|
||||
}
|
||||
|
||||
// ScheduleRecurring schedules a recurring task
|
||||
func (s *TimingWheelService) ScheduleRecurring(name string, interval time.Duration, fn func()) {
|
||||
var schedule func()
|
||||
schedule = func() {
|
||||
fn()
|
||||
_ = s.tw.SetTimer(name, schedule, interval)
|
||||
}
|
||||
_ = s.tw.SetTimer(name, schedule, interval)
|
||||
}
|
||||
|
||||
// Cancel cancels a scheduled task
|
||||
func (s *TimingWheelService) Cancel(name string) {
|
||||
_ = s.tw.RemoveTimer(name)
|
||||
}
|
||||
@@ -27,6 +27,7 @@ func NewTokenRefreshService(
|
||||
oauthService *OAuthService,
|
||||
openaiOAuthService *OpenAIOAuthService,
|
||||
geminiOAuthService *GeminiOAuthService,
|
||||
antigravityOAuthService *AntigravityOAuthService,
|
||||
cfg *config.Config,
|
||||
) *TokenRefreshService {
|
||||
s := &TokenRefreshService{
|
||||
@@ -40,6 +41,7 @@ func NewTokenRefreshService(
|
||||
NewClaudeTokenRefresher(oauthService),
|
||||
NewOpenAITokenRefresher(openaiOAuthService),
|
||||
NewGeminiTokenRefresher(geminiOAuthService),
|
||||
NewAntigravityTokenRefresher(antigravityOAuthService),
|
||||
}
|
||||
|
||||
return s
|
||||
@@ -106,6 +108,9 @@ func (s *TokenRefreshService) processRefresh() {
|
||||
return
|
||||
}
|
||||
|
||||
totalAccounts := len(accounts)
|
||||
oauthAccounts := 0 // 可刷新的OAuth账号数
|
||||
needsRefresh := 0 // 需要刷新的账号数
|
||||
refreshed, failed := 0, 0
|
||||
|
||||
for i := range accounts {
|
||||
@@ -117,11 +122,15 @@ func (s *TokenRefreshService) processRefresh() {
|
||||
continue
|
||||
}
|
||||
|
||||
oauthAccounts++
|
||||
|
||||
// 检查是否需要刷新
|
||||
if !refresher.NeedsRefresh(account, refreshWindow) {
|
||||
continue
|
||||
break // 不需要刷新,跳过
|
||||
}
|
||||
|
||||
needsRefresh++
|
||||
|
||||
// 执行刷新
|
||||
if err := s.refreshWithRetry(ctx, account, refresher); err != nil {
|
||||
log.Printf("[TokenRefresh] Account %d (%s) failed: %v", account.ID, account.Name, err)
|
||||
@@ -136,9 +145,9 @@ func (s *TokenRefreshService) processRefresh() {
|
||||
}
|
||||
}
|
||||
|
||||
if refreshed > 0 || failed > 0 {
|
||||
log.Printf("[TokenRefresh] Cycle complete: %d refreshed, %d failed", refreshed, failed)
|
||||
}
|
||||
// 始终打印周期日志,便于跟踪服务运行状态
|
||||
log.Printf("[TokenRefresh] Cycle complete: total=%d, oauth=%d, needs_refresh=%d, refreshed=%d, failed=%d",
|
||||
totalAccounts, oauthAccounts, needsRefresh, refreshed, failed)
|
||||
}
|
||||
|
||||
// listActiveAccounts 获取所有active状态的账号
|
||||
|
||||
@@ -43,18 +43,17 @@ func (r *ClaudeTokenRefresher) CanRefresh(account *Account) bool {
|
||||
// NeedsRefresh 检查token是否需要刷新
|
||||
// 基于 expires_at 字段判断是否在刷新窗口内
|
||||
func (r *ClaudeTokenRefresher) NeedsRefresh(account *Account, refreshWindow time.Duration) bool {
|
||||
expiresAtStr := account.GetCredential("expires_at")
|
||||
if expiresAtStr == "" {
|
||||
s := account.GetCredential("expires_at")
|
||||
if s == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
expiresAt, err := strconv.ParseInt(expiresAtStr, 10, 64)
|
||||
expiresAt, err := strconv.ParseInt(s, 10, 64)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
expiryTime := time.Unix(expiresAt, 0)
|
||||
return time.Until(expiryTime) < refreshWindow
|
||||
return time.Until(time.Unix(expiresAt, 0)) < refreshWindow
|
||||
}
|
||||
|
||||
// Refresh 执行token刷新
|
||||
|
||||
214
backend/internal/service/token_refresher_test.go
Normal file
214
backend/internal/service/token_refresher_test.go
Normal file
@@ -0,0 +1,214 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestClaudeTokenRefresher_NeedsRefresh(t *testing.T) {
|
||||
refresher := &ClaudeTokenRefresher{}
|
||||
refreshWindow := 30 * time.Minute
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
credentials map[string]any
|
||||
wantRefresh bool
|
||||
}{
|
||||
{
|
||||
name: "expires_at as string - expired",
|
||||
credentials: map[string]any{
|
||||
"expires_at": "1000", // 1970-01-01 00:16:40 UTC, 已过期
|
||||
},
|
||||
wantRefresh: true,
|
||||
},
|
||||
{
|
||||
name: "expires_at as float64 - expired",
|
||||
credentials: map[string]any{
|
||||
"expires_at": float64(1000), // 数字类型,已过期
|
||||
},
|
||||
wantRefresh: true,
|
||||
},
|
||||
{
|
||||
name: "expires_at as string - far future",
|
||||
credentials: map[string]any{
|
||||
"expires_at": "9999999999", // 远未来
|
||||
},
|
||||
wantRefresh: false,
|
||||
},
|
||||
{
|
||||
name: "expires_at as float64 - far future",
|
||||
credentials: map[string]any{
|
||||
"expires_at": float64(9999999999), // 远未来,数字类型
|
||||
},
|
||||
wantRefresh: false,
|
||||
},
|
||||
{
|
||||
name: "expires_at missing",
|
||||
credentials: map[string]any{},
|
||||
wantRefresh: false,
|
||||
},
|
||||
{
|
||||
name: "expires_at is nil",
|
||||
credentials: map[string]any{
|
||||
"expires_at": nil,
|
||||
},
|
||||
wantRefresh: false,
|
||||
},
|
||||
{
|
||||
name: "expires_at is invalid string",
|
||||
credentials: map[string]any{
|
||||
"expires_at": "invalid",
|
||||
},
|
||||
wantRefresh: false,
|
||||
},
|
||||
{
|
||||
name: "credentials is nil",
|
||||
credentials: nil,
|
||||
wantRefresh: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: tt.credentials,
|
||||
}
|
||||
|
||||
got := refresher.NeedsRefresh(account, refreshWindow)
|
||||
require.Equal(t, tt.wantRefresh, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaudeTokenRefresher_NeedsRefresh_WithinWindow(t *testing.T) {
|
||||
refresher := &ClaudeTokenRefresher{}
|
||||
refreshWindow := 30 * time.Minute
|
||||
|
||||
// 设置一个在刷新窗口内的时间(当前时间 + 15分钟)
|
||||
expiresAt := time.Now().Add(15 * time.Minute).Unix()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
credentials map[string]any
|
||||
}{
|
||||
{
|
||||
name: "string type - within refresh window",
|
||||
credentials: map[string]any{
|
||||
"expires_at": strconv.FormatInt(expiresAt, 10),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "float64 type - within refresh window",
|
||||
credentials: map[string]any{
|
||||
"expires_at": float64(expiresAt),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: tt.credentials,
|
||||
}
|
||||
|
||||
got := refresher.NeedsRefresh(account, refreshWindow)
|
||||
require.True(t, got, "should need refresh when within window")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaudeTokenRefresher_NeedsRefresh_OutsideWindow(t *testing.T) {
|
||||
refresher := &ClaudeTokenRefresher{}
|
||||
refreshWindow := 30 * time.Minute
|
||||
|
||||
// 设置一个在刷新窗口外的时间(当前时间 + 1小时)
|
||||
expiresAt := time.Now().Add(1 * time.Hour).Unix()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
credentials map[string]any
|
||||
}{
|
||||
{
|
||||
name: "string type - outside refresh window",
|
||||
credentials: map[string]any{
|
||||
"expires_at": strconv.FormatInt(expiresAt, 10),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "float64 type - outside refresh window",
|
||||
credentials: map[string]any{
|
||||
"expires_at": float64(expiresAt),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: tt.credentials,
|
||||
}
|
||||
|
||||
got := refresher.NeedsRefresh(account, refreshWindow)
|
||||
require.False(t, got, "should not need refresh when outside window")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaudeTokenRefresher_CanRefresh(t *testing.T) {
|
||||
refresher := &ClaudeTokenRefresher{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
platform string
|
||||
accType string
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
name: "anthropic oauth - can refresh",
|
||||
platform: PlatformAnthropic,
|
||||
accType: AccountTypeOAuth,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "anthropic api-key - cannot refresh",
|
||||
platform: PlatformAnthropic,
|
||||
accType: AccountTypeApiKey,
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "openai oauth - cannot refresh",
|
||||
platform: PlatformOpenAI,
|
||||
accType: AccountTypeOAuth,
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "gemini oauth - cannot refresh",
|
||||
platform: PlatformGemini,
|
||||
accType: AccountTypeOAuth,
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: tt.platform,
|
||||
Type: tt.accType,
|
||||
}
|
||||
|
||||
got := refresher.CanRefresh(account)
|
||||
require.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -148,22 +148,40 @@ func (s *UsageService) ListByAccount(ctx context.Context, accountID int64, param
|
||||
|
||||
// GetStatsByUser 获取用户的使用统计
|
||||
func (s *UsageService) GetStatsByUser(ctx context.Context, userID int64, startTime, endTime time.Time) (*UsageStats, error) {
|
||||
logs, _, err := s.usageRepo.ListByUserAndTimeRange(ctx, userID, startTime, endTime)
|
||||
stats, err := s.usageRepo.GetUserStatsAggregated(ctx, userID, startTime, endTime)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list usage logs: %w", err)
|
||||
return nil, fmt.Errorf("get user stats: %w", err)
|
||||
}
|
||||
|
||||
return s.calculateStats(logs), nil
|
||||
return &UsageStats{
|
||||
TotalRequests: stats.TotalRequests,
|
||||
TotalInputTokens: stats.TotalInputTokens,
|
||||
TotalOutputTokens: stats.TotalOutputTokens,
|
||||
TotalCacheTokens: stats.TotalCacheTokens,
|
||||
TotalTokens: stats.TotalTokens,
|
||||
TotalCost: stats.TotalCost,
|
||||
TotalActualCost: stats.TotalActualCost,
|
||||
AverageDurationMs: stats.AverageDurationMs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetStatsByApiKey 获取API Key的使用统计
|
||||
func (s *UsageService) GetStatsByApiKey(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) (*UsageStats, error) {
|
||||
logs, _, err := s.usageRepo.ListByApiKeyAndTimeRange(ctx, apiKeyID, startTime, endTime)
|
||||
stats, err := s.usageRepo.GetApiKeyStatsAggregated(ctx, apiKeyID, startTime, endTime)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list usage logs: %w", err)
|
||||
return nil, fmt.Errorf("get api key stats: %w", err)
|
||||
}
|
||||
|
||||
return s.calculateStats(logs), nil
|
||||
return &UsageStats{
|
||||
TotalRequests: stats.TotalRequests,
|
||||
TotalInputTokens: stats.TotalInputTokens,
|
||||
TotalOutputTokens: stats.TotalOutputTokens,
|
||||
TotalCacheTokens: stats.TotalCacheTokens,
|
||||
TotalTokens: stats.TotalTokens,
|
||||
TotalCost: stats.TotalCost,
|
||||
TotalActualCost: stats.TotalActualCost,
|
||||
AverageDurationMs: stats.AverageDurationMs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetStatsByAccount 获取账号的使用统计
|
||||
|
||||
@@ -18,6 +18,7 @@ type User struct {
|
||||
Concurrency int
|
||||
Status string
|
||||
AllowedGroups []int64
|
||||
TokenVersion int64 // Incremented on password change to invalidate existing tokens
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
|
||||
|
||||
@@ -116,6 +116,7 @@ func (s *UserService) UpdateProfile(ctx context.Context, userID int64, req Updat
|
||||
}
|
||||
|
||||
// ChangePassword 修改密码
|
||||
// Security: Increments TokenVersion to invalidate all existing JWT tokens
|
||||
func (s *UserService) ChangePassword(ctx context.Context, userID int64, req ChangePasswordRequest) error {
|
||||
user, err := s.userRepo.GetByID(ctx, userID)
|
||||
if err != nil {
|
||||
@@ -131,6 +132,10 @@ func (s *UserService) ChangePassword(ctx context.Context, userID int64, req Chan
|
||||
return fmt.Errorf("set password: %w", err)
|
||||
}
|
||||
|
||||
// Increment TokenVersion to invalidate all existing tokens
|
||||
// This ensures that any tokens issued before the password change become invalid
|
||||
user.TokenVersion++
|
||||
|
||||
if err := s.userRepo.Update(ctx, user); err != nil {
|
||||
return fmt.Errorf("update user: %w", err)
|
||||
}
|
||||
@@ -164,6 +169,14 @@ func (s *UserService) UpdateBalance(ctx context.Context, userID int64, amount fl
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateConcurrency 更新用户并发数(管理员功能)
|
||||
func (s *UserService) UpdateConcurrency(ctx context.Context, userID int64, concurrency int) error {
|
||||
if err := s.userRepo.UpdateConcurrency(ctx, userID, concurrency); err != nil {
|
||||
return fmt.Errorf("update concurrency: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateStatus 更新用户状态(管理员功能)
|
||||
func (s *UserService) UpdateStatus(ctx context.Context, userID int64, status string) error {
|
||||
user, err := s.userRepo.GetByID(ctx, userID)
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/google/wire"
|
||||
)
|
||||
@@ -15,7 +17,7 @@ type BuildInfo struct {
|
||||
func ProvidePricingService(cfg *config.Config, remoteClient PricingRemoteClient) (*PricingService, error) {
|
||||
svc := NewPricingService(cfg, remoteClient)
|
||||
if err := svc.Initialize(); err != nil {
|
||||
// 价格服务初始化失败不应阻止启动,使用回退价格
|
||||
// Pricing service initialization failure should not block startup, use fallback prices
|
||||
println("[Service] Warning: Pricing service initialization failed:", err.Error())
|
||||
}
|
||||
return svc, nil
|
||||
@@ -37,9 +39,36 @@ func ProvideTokenRefreshService(
|
||||
oauthService *OAuthService,
|
||||
openaiOAuthService *OpenAIOAuthService,
|
||||
geminiOAuthService *GeminiOAuthService,
|
||||
antigravityOAuthService *AntigravityOAuthService,
|
||||
cfg *config.Config,
|
||||
) *TokenRefreshService {
|
||||
svc := NewTokenRefreshService(accountRepo, oauthService, openaiOAuthService, geminiOAuthService, cfg)
|
||||
svc := NewTokenRefreshService(accountRepo, oauthService, openaiOAuthService, geminiOAuthService, antigravityOAuthService, cfg)
|
||||
svc.Start()
|
||||
return svc
|
||||
}
|
||||
|
||||
// ProvideTimingWheelService creates and starts TimingWheelService
|
||||
func ProvideTimingWheelService() *TimingWheelService {
|
||||
svc := NewTimingWheelService()
|
||||
svc.Start()
|
||||
return svc
|
||||
}
|
||||
|
||||
// ProvideAntigravityQuotaRefresher creates and starts AntigravityQuotaRefresher
|
||||
func ProvideAntigravityQuotaRefresher(
|
||||
accountRepo AccountRepository,
|
||||
proxyRepo ProxyRepository,
|
||||
oauthSvc *AntigravityOAuthService,
|
||||
cfg *config.Config,
|
||||
) *AntigravityQuotaRefresher {
|
||||
svc := NewAntigravityQuotaRefresher(accountRepo, proxyRepo, oauthSvc, cfg)
|
||||
svc.Start()
|
||||
return svc
|
||||
}
|
||||
|
||||
// ProvideDeferredService creates and starts DeferredService
|
||||
func ProvideDeferredService(accountRepo AccountRepository, timingWheel *TimingWheelService) *DeferredService {
|
||||
svc := NewDeferredService(accountRepo, timingWheel, 10*time.Second)
|
||||
svc.Start()
|
||||
return svc
|
||||
}
|
||||
@@ -65,8 +94,11 @@ var ProviderSet = wire.NewSet(
|
||||
NewOAuthService,
|
||||
NewOpenAIOAuthService,
|
||||
NewGeminiOAuthService,
|
||||
NewAntigravityOAuthService,
|
||||
NewGeminiTokenProvider,
|
||||
NewGeminiMessagesCompatService,
|
||||
NewAntigravityTokenProvider,
|
||||
NewAntigravityGatewayService,
|
||||
NewRateLimitService,
|
||||
NewAccountUsageService,
|
||||
NewAccountTestService,
|
||||
@@ -80,4 +112,7 @@ var ProviderSet = wire.NewSet(
|
||||
NewCRSSyncService,
|
||||
ProvideUpdateService,
|
||||
ProvideTokenRefreshService,
|
||||
ProvideTimingWheelService,
|
||||
ProvideDeferredService,
|
||||
ProvideAntigravityQuotaRefresher,
|
||||
)
|
||||
|
||||
@@ -271,7 +271,9 @@ func initializeDatabase(cfg *SetupConfig) error {
|
||||
}
|
||||
}()
|
||||
|
||||
return repository.AutoMigrate(db)
|
||||
// setup 阶段使用 standard 模式(不创建默认分组)
|
||||
// 默认分组将在正常启动时根据实际 runMode 决定是否创建
|
||||
return repository.AutoMigrate(db, "standard")
|
||||
}
|
||||
|
||||
func createAdminUser(cfg *SetupConfig) error {
|
||||
|
||||
@@ -28,8 +28,10 @@ func ServeEmbeddedFrontend() gin.HandlerFunc {
|
||||
if strings.HasPrefix(path, "/api/") ||
|
||||
strings.HasPrefix(path, "/v1/") ||
|
||||
strings.HasPrefix(path, "/v1beta/") ||
|
||||
strings.HasPrefix(path, "/antigravity/") ||
|
||||
strings.HasPrefix(path, "/setup/") ||
|
||||
path == "/health" {
|
||||
path == "/health" ||
|
||||
path == "/responses" {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
@@ -20,6 +20,10 @@ SERVER_PORT=8080
|
||||
# Server mode: release or debug
|
||||
SERVER_MODE=release
|
||||
|
||||
# 运行模式: standard (默认) 或 simple (内部自用)
|
||||
# standard: 完整 SaaS 功能,包含计费/余额校验;simple: 隐藏 SaaS 功能并跳过计费/余额校验
|
||||
RUN_MODE=standard
|
||||
|
||||
# Timezone
|
||||
TZ=Asia/Shanghai
|
||||
|
||||
|
||||
184
deploy/Caddyfile
Normal file
184
deploy/Caddyfile
Normal file
@@ -0,0 +1,184 @@
|
||||
# =============================================================================
|
||||
# Sub2API Caddy Reverse Proxy Configuration (宿主机部署)
|
||||
# =============================================================================
|
||||
# 使用方法:
|
||||
# 1. 安装 Caddy: https://caddyserver.com/docs/install
|
||||
# 2. 修改下方 example.com 为你的域名
|
||||
# 3. 确保域名 DNS 已指向服务器
|
||||
# 4. 复制配置: sudo cp Caddyfile /etc/caddy/Caddyfile
|
||||
# 5. 重载配置: sudo systemctl reload caddy
|
||||
#
|
||||
# Caddy 会自动申请和续期 Let's Encrypt SSL 证书
|
||||
# =============================================================================
|
||||
|
||||
# 全局配置
|
||||
{
|
||||
# Let's Encrypt 邮箱通知
|
||||
email admin@example.com
|
||||
|
||||
# 服务器配置
|
||||
servers {
|
||||
# 启用 HTTP/2 和 HTTP/3
|
||||
protocols h1 h2 h3
|
||||
|
||||
# 超时配置
|
||||
timeouts {
|
||||
read_body 30s
|
||||
read_header 10s
|
||||
write 60s
|
||||
idle 120s
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# 修改为你的域名
|
||||
example.com {
|
||||
# =========================================================================
|
||||
# TLS 安全配置
|
||||
# =========================================================================
|
||||
tls {
|
||||
# 仅使用 TLS 1.2 和 1.3
|
||||
protocols tls1.2 tls1.3
|
||||
|
||||
# 优先使用的加密套件
|
||||
ciphers TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384 TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384 TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256 TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256 TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256 TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256
|
||||
}
|
||||
|
||||
# =========================================================================
|
||||
# 反向代理配置
|
||||
# =========================================================================
|
||||
reverse_proxy localhost:8080 {
|
||||
# 健康检查
|
||||
health_uri /health
|
||||
health_interval 30s
|
||||
health_timeout 10s
|
||||
health_status 200
|
||||
|
||||
# 负载均衡策略(单节点可忽略,多节点时有用)
|
||||
lb_policy round_robin
|
||||
lb_try_duration 5s
|
||||
lb_try_interval 250ms
|
||||
|
||||
# 传递真实客户端信息
|
||||
# 兼容 Cloudflare 和直连:后端应优先读取 CF-Connecting-IP,其次 X-Real-IP
|
||||
header_up X-Real-IP {remote_host}
|
||||
header_up X-Forwarded-For {remote_host}
|
||||
header_up X-Forwarded-Proto {scheme}
|
||||
header_up X-Forwarded-Host {host}
|
||||
# 保留 Cloudflare 原始头(如果存在)
|
||||
# 后端获取 IP 的优先级建议: CF-Connecting-IP → X-Real-IP → X-Forwarded-For
|
||||
header_up CF-Connecting-IP {http.request.header.CF-Connecting-IP}
|
||||
|
||||
# 连接池优化
|
||||
transport http {
|
||||
keepalive 120s
|
||||
keepalive_idle_conns 256
|
||||
read_buffer 16KB
|
||||
write_buffer 16KB
|
||||
compression off
|
||||
}
|
||||
|
||||
# 故障转移
|
||||
fail_duration 30s
|
||||
max_fails 3
|
||||
unhealthy_status 500 502 503 504
|
||||
}
|
||||
|
||||
# =========================================================================
|
||||
# 压缩配置
|
||||
# =========================================================================
|
||||
encode {
|
||||
zstd
|
||||
gzip 6
|
||||
minimum_length 256
|
||||
match {
|
||||
header Content-Type text/*
|
||||
header Content-Type application/json*
|
||||
header Content-Type application/javascript*
|
||||
header Content-Type application/xml*
|
||||
header Content-Type application/rss+xml*
|
||||
header Content-Type image/svg+xml*
|
||||
}
|
||||
}
|
||||
|
||||
# =========================================================================
|
||||
# 速率限制 (需要 caddy-ratelimit 插件)
|
||||
# 如未安装插件,请注释掉此段
|
||||
# =========================================================================
|
||||
# rate_limit {
|
||||
# zone api {
|
||||
# key {remote_host}
|
||||
# events 100
|
||||
# window 1m
|
||||
# }
|
||||
# }
|
||||
|
||||
# =========================================================================
|
||||
# 安全响应头
|
||||
# =========================================================================
|
||||
header {
|
||||
# 防止点击劫持
|
||||
X-Frame-Options "SAMEORIGIN"
|
||||
|
||||
# XSS 保护
|
||||
X-XSS-Protection "1; mode=block"
|
||||
|
||||
# 防止 MIME 类型嗅探
|
||||
X-Content-Type-Options "nosniff"
|
||||
|
||||
# 引用策略
|
||||
Referrer-Policy "strict-origin-when-cross-origin"
|
||||
|
||||
# HSTS - 强制 HTTPS (max-age=1年)
|
||||
Strict-Transport-Security "max-age=31536000; includeSubDomains; preload"
|
||||
|
||||
# 内容安全策略 (根据需要调整)
|
||||
# Content-Security-Policy "default-src 'self'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'; img-src 'self' data: https:; font-src 'self' data:; connect-src 'self' https:;"
|
||||
|
||||
# 权限策略
|
||||
Permissions-Policy "accelerometer=(), camera=(), geolocation=(), gyroscope=(), magnetometer=(), microphone=(), payment=(), usb=()"
|
||||
|
||||
# 跨域资源策略
|
||||
Cross-Origin-Opener-Policy "same-origin"
|
||||
Cross-Origin-Embedder-Policy "require-corp"
|
||||
Cross-Origin-Resource-Policy "same-origin"
|
||||
|
||||
# 移除敏感头
|
||||
-Server
|
||||
-X-Powered-By
|
||||
}
|
||||
|
||||
# =========================================================================
|
||||
# 请求大小限制 (防止大文件攻击)
|
||||
# =========================================================================
|
||||
request_body {
|
||||
max_size 100MB
|
||||
}
|
||||
|
||||
# =========================================================================
|
||||
# 日志配置
|
||||
# =========================================================================
|
||||
log {
|
||||
output file /var/log/caddy/sub2api.log {
|
||||
roll_size 50mb
|
||||
roll_keep 10
|
||||
roll_keep_for 720h
|
||||
}
|
||||
format json
|
||||
level INFO
|
||||
}
|
||||
|
||||
# =========================================================================
|
||||
# 错误处理
|
||||
# =========================================================================
|
||||
handle_errors {
|
||||
respond "{err.status_code} {err.status_text}"
|
||||
}
|
||||
}
|
||||
|
||||
# =============================================================================
|
||||
# HTTP 重定向到 HTTPS (Caddy 默认自动处理,此处显式声明)
|
||||
# =============================================================================
|
||||
; http://example.com {
|
||||
; redir https://{host}{uri} permanent
|
||||
; }
|
||||
@@ -281,6 +281,30 @@ To change after installation:
|
||||
sudo systemctl restart sub2api
|
||||
```
|
||||
|
||||
#### Gemini OAuth Configuration
|
||||
|
||||
If you need to use AI Studio OAuth for Gemini accounts, add the OAuth client credentials to the systemd service file:
|
||||
|
||||
1. Edit the service file:
|
||||
```bash
|
||||
sudo nano /etc/systemd/system/sub2api.service
|
||||
```
|
||||
|
||||
2. Add your OAuth credentials in the `[Service]` section (after the existing `Environment=` lines):
|
||||
```ini
|
||||
Environment=GEMINI_OAUTH_CLIENT_ID=your-client-id.apps.googleusercontent.com
|
||||
Environment=GEMINI_OAUTH_CLIENT_SECRET=GOCSPX-your-client-secret
|
||||
```
|
||||
|
||||
3. Reload and restart:
|
||||
```bash
|
||||
sudo systemctl daemon-reload
|
||||
sudo systemctl restart sub2api
|
||||
```
|
||||
|
||||
> **Note:** Code Assist OAuth does not require any configuration - it uses the built-in Gemini CLI client.
|
||||
> See the [Gemini OAuth Configuration](#gemini-oauth-configuration) section above for detailed setup instructions.
|
||||
|
||||
#### Application Configuration
|
||||
|
||||
The main config file is at `/etc/sub2api/config.yaml` (created by Setup Wizard).
|
||||
|
||||
@@ -13,6 +13,14 @@ server:
|
||||
# Mode: "debug" for development, "release" for production
|
||||
mode: "release"
|
||||
|
||||
# =============================================================================
|
||||
# Run Mode Configuration
|
||||
# =============================================================================
|
||||
# Run mode: "standard" (default) or "simple" (for internal use)
|
||||
# - standard: Full SaaS features with billing/balance checks
|
||||
# - simple: Hides SaaS features and skips billing/balance checks
|
||||
run_mode: "standard"
|
||||
|
||||
# =============================================================================
|
||||
# Database Configuration (PostgreSQL)
|
||||
# =============================================================================
|
||||
|
||||
@@ -36,6 +36,7 @@ services:
|
||||
- SERVER_HOST=0.0.0.0
|
||||
- SERVER_PORT=8080
|
||||
- SERVER_MODE=${SERVER_MODE:-release}
|
||||
- RUN_MODE=${RUN_MODE:-standard}
|
||||
|
||||
# =======================================================================
|
||||
# Database Configuration (PostgreSQL)
|
||||
@@ -121,8 +122,8 @@ services:
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
start_period: 10s
|
||||
ports:
|
||||
- 5433:5432
|
||||
# 注意:不暴露端口到宿主机,应用通过内部网络连接
|
||||
# 如需调试,可临时添加:ports: ["127.0.0.1:5433:5432"]
|
||||
|
||||
# ===========================================================================
|
||||
# Redis Cache
|
||||
|
||||
@@ -404,7 +404,9 @@ configure_server() {
|
||||
|
||||
# Check if running as root
|
||||
check_root() {
|
||||
if [ "$EUID" -ne 0 ]; then
|
||||
# Use 'id -u' instead of $EUID for better compatibility
|
||||
# $EUID may not work reliably when script is piped to bash
|
||||
if [ "$(id -u)" -ne 0 ]; then
|
||||
print_error "$(msg 'run_as_root')"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
1229
frontend/package-lock.json
generated
1229
frontend/package-lock.json
generated
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user