mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
Merge branch 'release/custom-0.1.104'
This commit is contained in:
@@ -271,3 +271,36 @@ jobs:
|
||||
parse_mode: "Markdown",
|
||||
disable_web_page_preview: true
|
||||
}')"
|
||||
|
||||
sync-version-file:
|
||||
needs: [release]
|
||||
if: ${{ needs.release.result == 'success' }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout default branch
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.repository.default_branch }}
|
||||
|
||||
- name: Sync VERSION file to released tag
|
||||
run: |
|
||||
if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
|
||||
VERSION=${{ github.event.inputs.tag }}
|
||||
VERSION=${VERSION#v}
|
||||
else
|
||||
VERSION=${GITHUB_REF#refs/tags/v}
|
||||
fi
|
||||
|
||||
CURRENT_VERSION=$(tr -d '\r\n' < backend/cmd/server/VERSION || true)
|
||||
if [ "$CURRENT_VERSION" = "$VERSION" ]; then
|
||||
echo "VERSION file already matches $VERSION"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "$VERSION" > backend/cmd/server/VERSION
|
||||
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
|
||||
git add backend/cmd/server/VERSION
|
||||
git commit -m "chore: sync VERSION to ${VERSION} [skip ci]"
|
||||
git push origin HEAD:${{ github.event.repository.default_branch }}
|
||||
|
||||
@@ -128,6 +128,7 @@ deploy/docker-compose.override.yml
|
||||
.gocache/
|
||||
vite.config.js
|
||||
docs/*
|
||||
!docs/ACCOUNT_SCHEDULING_FLOW.md
|
||||
.serena/
|
||||
|
||||
# ===================
|
||||
|
||||
@@ -47,6 +47,8 @@ dockers:
|
||||
- "ghcr.io/{{ .Env.GITHUB_REPO_OWNER_LOWER }}/sub2api:latest"
|
||||
dockerfile: Dockerfile.goreleaser
|
||||
use: buildx
|
||||
extra_files:
|
||||
- deploy/docker-entrypoint.sh
|
||||
build_flag_templates:
|
||||
- "--platform=linux/amd64"
|
||||
- "--label=org.opencontainers.image.version={{ .Version }}"
|
||||
|
||||
@@ -63,6 +63,8 @@ dockers:
|
||||
- "{{ .Env.DOCKERHUB_USERNAME }}/sub2api:{{ .Version }}-amd64"
|
||||
dockerfile: Dockerfile.goreleaser
|
||||
use: buildx
|
||||
extra_files:
|
||||
- deploy/docker-entrypoint.sh
|
||||
build_flag_templates:
|
||||
- "--platform=linux/amd64"
|
||||
- "--label=org.opencontainers.image.version={{ .Version }}"
|
||||
@@ -76,6 +78,8 @@ dockers:
|
||||
- "{{ .Env.DOCKERHUB_USERNAME }}/sub2api:{{ .Version }}-arm64"
|
||||
dockerfile: Dockerfile.goreleaser
|
||||
use: buildx
|
||||
extra_files:
|
||||
- deploy/docker-entrypoint.sh
|
||||
build_flag_templates:
|
||||
- "--platform=linux/arm64"
|
||||
- "--label=org.opencontainers.image.version={{ .Version }}"
|
||||
@@ -89,6 +93,8 @@ dockers:
|
||||
- "ghcr.io/{{ .Env.GITHUB_REPO_OWNER_LOWER }}/sub2api:{{ .Version }}-amd64"
|
||||
dockerfile: Dockerfile.goreleaser
|
||||
use: buildx
|
||||
extra_files:
|
||||
- deploy/docker-entrypoint.sh
|
||||
build_flag_templates:
|
||||
- "--platform=linux/amd64"
|
||||
- "--label=org.opencontainers.image.version={{ .Version }}"
|
||||
@@ -102,6 +108,8 @@ dockers:
|
||||
- "ghcr.io/{{ .Env.GITHUB_REPO_OWNER_LOWER }}/sub2api:{{ .Version }}-arm64"
|
||||
dockerfile: Dockerfile.goreleaser
|
||||
use: buildx
|
||||
extra_files:
|
||||
- deploy/docker-entrypoint.sh
|
||||
build_flag_templates:
|
||||
- "--platform=linux/arm64"
|
||||
- "--label=org.opencontainers.image.version={{ .Version }}"
|
||||
|
||||
+7
-4
@@ -92,6 +92,7 @@ LABEL org.opencontainers.image.source="https://github.com/Wei-Shaw/sub2api"
|
||||
RUN apk add --no-cache \
|
||||
ca-certificates \
|
||||
tzdata \
|
||||
su-exec \
|
||||
libpq \
|
||||
zstd-libs \
|
||||
lz4-libs \
|
||||
@@ -120,8 +121,9 @@ COPY --from=backend-builder --chown=sub2api:sub2api /app/backend/resources /app/
|
||||
# Create data directory
|
||||
RUN mkdir -p /app/data && chown sub2api:sub2api /app/data
|
||||
|
||||
# Switch to non-root user
|
||||
USER sub2api
|
||||
# Copy entrypoint script (fixes volume permissions then drops to sub2api)
|
||||
COPY deploy/docker-entrypoint.sh /app/docker-entrypoint.sh
|
||||
RUN chmod +x /app/docker-entrypoint.sh
|
||||
|
||||
# Expose port (can be overridden by SERVER_PORT env var)
|
||||
EXPOSE 8080
|
||||
@@ -130,5 +132,6 @@ EXPOSE 8080
|
||||
HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=3 \
|
||||
CMD wget -q -T 5 -O /dev/null http://localhost:${SERVER_PORT:-8080}/health || exit 1
|
||||
|
||||
# Run the application
|
||||
ENTRYPOINT ["/app/sub2api"]
|
||||
# Run the application (entrypoint fixes /app/data ownership then execs as sub2api)
|
||||
ENTRYPOINT ["/app/docker-entrypoint.sh"]
|
||||
CMD ["/app/sub2api"]
|
||||
|
||||
@@ -21,6 +21,7 @@ RUN apk add --no-cache \
|
||||
ca-certificates \
|
||||
tzdata \
|
||||
curl \
|
||||
su-exec \
|
||||
libpq \
|
||||
zstd-libs \
|
||||
lz4-libs \
|
||||
@@ -47,11 +48,15 @@ COPY sub2api /app/sub2api
|
||||
# Create data directory
|
||||
RUN mkdir -p /app/data && chown -R sub2api:sub2api /app
|
||||
|
||||
USER sub2api
|
||||
# Copy entrypoint script (fixes volume permissions then drops to sub2api)
|
||||
COPY deploy/docker-entrypoint.sh /app/docker-entrypoint.sh
|
||||
RUN chmod +x /app/docker-entrypoint.sh
|
||||
|
||||
EXPOSE 8080
|
||||
|
||||
HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=3 \
|
||||
CMD curl -f http://localhost:${SERVER_PORT:-8080}/health || exit 1
|
||||
|
||||
ENTRYPOINT ["/app/sub2api"]
|
||||
# Run the application (entrypoint fixes /app/data ownership then execs as sub2api)
|
||||
ENTRYPOINT ["/app/docker-entrypoint.sh"]
|
||||
CMD ["/app/sub2api"]
|
||||
|
||||
@@ -178,10 +178,10 @@ mkdir -p sub2api-deploy && cd sub2api-deploy
|
||||
curl -sSL https://raw.githubusercontent.com/Wei-Shaw/sub2api/main/deploy/docker-deploy.sh | bash
|
||||
|
||||
# Start services
|
||||
docker-compose up -d
|
||||
docker compose up -d
|
||||
|
||||
# View logs
|
||||
docker-compose logs -f sub2api
|
||||
docker compose logs -f sub2api
|
||||
```
|
||||
|
||||
**What the script does:**
|
||||
@@ -245,16 +245,16 @@ mkdir -p data postgres_data redis_data
|
||||
|
||||
# 5. Start all services
|
||||
# Option A: Local directory version (recommended - easy migration)
|
||||
docker-compose -f docker-compose.local.yml up -d
|
||||
docker compose -f docker-compose.local.yml up -d
|
||||
|
||||
# Option B: Named volumes version (simple setup)
|
||||
docker-compose up -d
|
||||
docker compose up -d
|
||||
|
||||
# 6. Check status
|
||||
docker-compose -f docker-compose.local.yml ps
|
||||
docker compose -f docker-compose.local.yml ps
|
||||
|
||||
# 7. View logs
|
||||
docker-compose -f docker-compose.local.yml logs -f sub2api
|
||||
docker compose -f docker-compose.local.yml logs -f sub2api
|
||||
```
|
||||
|
||||
#### Deployment Versions
|
||||
@@ -272,15 +272,15 @@ Open `http://YOUR_SERVER_IP:8080` in your browser.
|
||||
|
||||
If admin password was auto-generated, find it in logs:
|
||||
```bash
|
||||
docker-compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
```
|
||||
|
||||
#### Upgrade
|
||||
|
||||
```bash
|
||||
# Pull latest image and recreate container
|
||||
docker-compose -f docker-compose.local.yml pull
|
||||
docker-compose -f docker-compose.local.yml up -d
|
||||
docker compose -f docker-compose.local.yml pull
|
||||
docker compose -f docker-compose.local.yml up -d
|
||||
```
|
||||
|
||||
#### Easy Migration (Local Directory Version)
|
||||
@@ -289,7 +289,7 @@ When using `docker-compose.local.yml`, migrate to a new server easily:
|
||||
|
||||
```bash
|
||||
# On source server
|
||||
docker-compose -f docker-compose.local.yml down
|
||||
docker compose -f docker-compose.local.yml down
|
||||
cd ..
|
||||
tar czf sub2api-complete.tar.gz sub2api-deploy/
|
||||
|
||||
@@ -299,23 +299,23 @@ scp sub2api-complete.tar.gz user@new-server:/path/
|
||||
# On new server
|
||||
tar xzf sub2api-complete.tar.gz
|
||||
cd sub2api-deploy/
|
||||
docker-compose -f docker-compose.local.yml up -d
|
||||
docker compose -f docker-compose.local.yml up -d
|
||||
```
|
||||
|
||||
#### Useful Commands
|
||||
|
||||
```bash
|
||||
# Stop all services
|
||||
docker-compose -f docker-compose.local.yml down
|
||||
docker compose -f docker-compose.local.yml down
|
||||
|
||||
# Restart
|
||||
docker-compose -f docker-compose.local.yml restart
|
||||
docker compose -f docker-compose.local.yml restart
|
||||
|
||||
# View all logs
|
||||
docker-compose -f docker-compose.local.yml logs -f
|
||||
docker compose -f docker-compose.local.yml logs -f
|
||||
|
||||
# Remove all data (caution!)
|
||||
docker-compose -f docker-compose.local.yml down
|
||||
docker compose -f docker-compose.local.yml down
|
||||
rm -rf data/ postgres_data/ redis_data/
|
||||
```
|
||||
|
||||
|
||||
+15
-15
@@ -177,10 +177,10 @@ mkdir -p sub2api-deploy && cd sub2api-deploy
|
||||
curl -sSL https://raw.githubusercontent.com/Wei-Shaw/sub2api/main/deploy/docker-deploy.sh | bash
|
||||
|
||||
# 启动服务
|
||||
docker-compose up -d
|
||||
docker compose up -d
|
||||
|
||||
# 查看日志
|
||||
docker-compose logs -f sub2api
|
||||
docker compose logs -f sub2api
|
||||
```
|
||||
|
||||
**脚本功能:**
|
||||
@@ -244,16 +244,16 @@ mkdir -p data postgres_data redis_data
|
||||
|
||||
# 5. 启动所有服务
|
||||
# 选项 A:本地目录版(推荐 - 易于迁移)
|
||||
docker-compose -f docker-compose.local.yml up -d
|
||||
docker compose -f docker-compose.local.yml up -d
|
||||
|
||||
# 选项 B:命名卷版(简单设置)
|
||||
docker-compose up -d
|
||||
docker compose up -d
|
||||
|
||||
# 6. 查看状态
|
||||
docker-compose -f docker-compose.local.yml ps
|
||||
docker compose -f docker-compose.local.yml ps
|
||||
|
||||
# 7. 查看日志
|
||||
docker-compose -f docker-compose.local.yml logs -f sub2api
|
||||
docker compose -f docker-compose.local.yml logs -f sub2api
|
||||
```
|
||||
|
||||
#### 部署版本对比
|
||||
@@ -283,15 +283,15 @@ docker-compose -f docker-compose.local.yml logs -f sub2api
|
||||
|
||||
如果管理员密码是自动生成的,在日志中查找:
|
||||
```bash
|
||||
docker-compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
```
|
||||
|
||||
#### 升级
|
||||
|
||||
```bash
|
||||
# 拉取最新镜像并重建容器
|
||||
docker-compose -f docker-compose.local.yml pull
|
||||
docker-compose -f docker-compose.local.yml up -d
|
||||
docker compose -f docker-compose.local.yml pull
|
||||
docker compose -f docker-compose.local.yml up -d
|
||||
```
|
||||
|
||||
#### 轻松迁移(本地目录版)
|
||||
@@ -300,7 +300,7 @@ docker-compose -f docker-compose.local.yml up -d
|
||||
|
||||
```bash
|
||||
# 源服务器
|
||||
docker-compose -f docker-compose.local.yml down
|
||||
docker compose -f docker-compose.local.yml down
|
||||
cd ..
|
||||
tar czf sub2api-complete.tar.gz sub2api-deploy/
|
||||
|
||||
@@ -310,23 +310,23 @@ scp sub2api-complete.tar.gz user@new-server:/path/
|
||||
# 新服务器
|
||||
tar xzf sub2api-complete.tar.gz
|
||||
cd sub2api-deploy/
|
||||
docker-compose -f docker-compose.local.yml up -d
|
||||
docker compose -f docker-compose.local.yml up -d
|
||||
```
|
||||
|
||||
#### 常用命令
|
||||
|
||||
```bash
|
||||
# 停止所有服务
|
||||
docker-compose -f docker-compose.local.yml down
|
||||
docker compose -f docker-compose.local.yml down
|
||||
|
||||
# 重启
|
||||
docker-compose -f docker-compose.local.yml restart
|
||||
docker compose -f docker-compose.local.yml restart
|
||||
|
||||
# 查看所有日志
|
||||
docker-compose -f docker-compose.local.yml logs -f
|
||||
docker compose -f docker-compose.local.yml logs -f
|
||||
|
||||
# 删除所有数据(谨慎!)
|
||||
docker-compose -f docker-compose.local.yml down
|
||||
docker compose -f docker-compose.local.yml down
|
||||
rm -rf data/ postgres_data/ redis_data/
|
||||
```
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
0.1.101.1
|
||||
0.1.104.1
|
||||
|
||||
@@ -110,7 +110,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
concurrencyCache := repository.ProvideConcurrencyCache(redisClient, configConfig)
|
||||
concurrencyService := service.ProvideConcurrencyService(concurrencyCache, accountRepository, configConfig)
|
||||
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService)
|
||||
groupHandler := admin.NewGroupHandler(adminService)
|
||||
claudeOAuthClient := repository.NewClaudeOAuthClient()
|
||||
oAuthService := service.NewOAuthService(proxyRepository, claudeOAuthClient)
|
||||
openAIOAuthClient := repository.NewOpenAIOAuthClient()
|
||||
@@ -137,12 +136,14 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
gatewayCache := repository.NewGatewayCache(redisClient)
|
||||
schedulerOutboxRepository := repository.NewSchedulerOutboxRepository(db)
|
||||
schedulerSnapshotService := service.ProvideSchedulerSnapshotService(schedulerCache, schedulerOutboxRepository, accountRepository, groupRepository, configConfig)
|
||||
antigravityTokenProvider := service.ProvideAntigravityTokenProvider(accountRepository, geminiTokenCache, antigravityOAuthService, oauthRefreshAPI)
|
||||
antigravityTokenProvider := service.ProvideAntigravityTokenProvider(accountRepository, geminiTokenCache, antigravityOAuthService, oauthRefreshAPI, tempUnschedCache)
|
||||
antigravityGatewayService := service.NewAntigravityGatewayService(accountRepository, gatewayCache, schedulerSnapshotService, antigravityTokenProvider, rateLimitService, httpUpstream, settingService)
|
||||
accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, antigravityGatewayService, httpUpstream, configConfig)
|
||||
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
|
||||
sessionLimitCache := repository.ProvideSessionLimitCache(redisClient, configConfig)
|
||||
rpmCache := repository.NewRPMCache(redisClient)
|
||||
groupCapacityService := service.NewGroupCapacityService(accountRepository, groupRepository, concurrencyService, sessionLimitCache, rpmCache)
|
||||
groupHandler := admin.NewGroupHandler(adminService, dashboardService, groupCapacityService)
|
||||
accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, gatewayCache)
|
||||
adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService)
|
||||
dataManagementService := service.NewDataManagementService()
|
||||
|
||||
@@ -717,6 +717,7 @@ var (
|
||||
{Name: "id", Type: field.TypeInt64, Increment: true},
|
||||
{Name: "request_id", Type: field.TypeString, Size: 64},
|
||||
{Name: "model", Type: field.TypeString, Size: 100},
|
||||
{Name: "upstream_model", Type: field.TypeString, Nullable: true, Size: 100},
|
||||
{Name: "input_tokens", Type: field.TypeInt, Default: 0},
|
||||
{Name: "output_tokens", Type: field.TypeInt, Default: 0},
|
||||
{Name: "cache_creation_tokens", Type: field.TypeInt, Default: 0},
|
||||
@@ -756,31 +757,31 @@ var (
|
||||
ForeignKeys: []*schema.ForeignKey{
|
||||
{
|
||||
Symbol: "usage_logs_api_keys_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[28]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[29]},
|
||||
RefColumns: []*schema.Column{APIKeysColumns[0]},
|
||||
OnDelete: schema.NoAction,
|
||||
},
|
||||
{
|
||||
Symbol: "usage_logs_accounts_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[29]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[30]},
|
||||
RefColumns: []*schema.Column{AccountsColumns[0]},
|
||||
OnDelete: schema.NoAction,
|
||||
},
|
||||
{
|
||||
Symbol: "usage_logs_groups_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[30]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[31]},
|
||||
RefColumns: []*schema.Column{GroupsColumns[0]},
|
||||
OnDelete: schema.SetNull,
|
||||
},
|
||||
{
|
||||
Symbol: "usage_logs_users_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[31]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[32]},
|
||||
RefColumns: []*schema.Column{UsersColumns[0]},
|
||||
OnDelete: schema.NoAction,
|
||||
},
|
||||
{
|
||||
Symbol: "usage_logs_user_subscriptions_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[32]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[33]},
|
||||
RefColumns: []*schema.Column{UserSubscriptionsColumns[0]},
|
||||
OnDelete: schema.SetNull,
|
||||
},
|
||||
@@ -789,32 +790,32 @@ var (
|
||||
{
|
||||
Name: "usagelog_user_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[31]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[32]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_api_key_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[28]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[29]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_account_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[29]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[30]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_group_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[30]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[31]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_subscription_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[32]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[33]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[27]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[28]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_model",
|
||||
@@ -829,17 +830,17 @@ var (
|
||||
{
|
||||
Name: "usagelog_user_id_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[31], UsageLogsColumns[27]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[32], UsageLogsColumns[28]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_api_key_id_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[28], UsageLogsColumns[27]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[29], UsageLogsColumns[28]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_group_id_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[30], UsageLogsColumns[27]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[31], UsageLogsColumns[28]},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
+74
-1
@@ -18293,6 +18293,7 @@ type UsageLogMutation struct {
|
||||
id *int64
|
||||
request_id *string
|
||||
model *string
|
||||
upstream_model *string
|
||||
input_tokens *int
|
||||
addinput_tokens *int
|
||||
output_tokens *int
|
||||
@@ -18630,6 +18631,55 @@ func (m *UsageLogMutation) ResetModel() {
|
||||
m.model = nil
|
||||
}
|
||||
|
||||
// SetUpstreamModel sets the "upstream_model" field.
|
||||
func (m *UsageLogMutation) SetUpstreamModel(s string) {
|
||||
m.upstream_model = &s
|
||||
}
|
||||
|
||||
// UpstreamModel returns the value of the "upstream_model" field in the mutation.
|
||||
func (m *UsageLogMutation) UpstreamModel() (r string, exists bool) {
|
||||
v := m.upstream_model
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldUpstreamModel returns the old "upstream_model" field's value of the UsageLog entity.
|
||||
// If the UsageLog object wasn't provided to the builder, the object is fetched from the database.
|
||||
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
|
||||
func (m *UsageLogMutation) OldUpstreamModel(ctx context.Context) (v *string, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldUpstreamModel is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldUpstreamModel requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldUpstreamModel: %w", err)
|
||||
}
|
||||
return oldValue.UpstreamModel, nil
|
||||
}
|
||||
|
||||
// ClearUpstreamModel clears the value of the "upstream_model" field.
|
||||
func (m *UsageLogMutation) ClearUpstreamModel() {
|
||||
m.upstream_model = nil
|
||||
m.clearedFields[usagelog.FieldUpstreamModel] = struct{}{}
|
||||
}
|
||||
|
||||
// UpstreamModelCleared returns if the "upstream_model" field was cleared in this mutation.
|
||||
func (m *UsageLogMutation) UpstreamModelCleared() bool {
|
||||
_, ok := m.clearedFields[usagelog.FieldUpstreamModel]
|
||||
return ok
|
||||
}
|
||||
|
||||
// ResetUpstreamModel resets all changes to the "upstream_model" field.
|
||||
func (m *UsageLogMutation) ResetUpstreamModel() {
|
||||
m.upstream_model = nil
|
||||
delete(m.clearedFields, usagelog.FieldUpstreamModel)
|
||||
}
|
||||
|
||||
// SetGroupID sets the "group_id" field.
|
||||
func (m *UsageLogMutation) SetGroupID(i int64) {
|
||||
m.group = &i
|
||||
@@ -20251,7 +20301,7 @@ func (m *UsageLogMutation) Type() string {
|
||||
// order to get all numeric fields that were incremented/decremented, call
|
||||
// AddedFields().
|
||||
func (m *UsageLogMutation) Fields() []string {
|
||||
fields := make([]string, 0, 32)
|
||||
fields := make([]string, 0, 33)
|
||||
if m.user != nil {
|
||||
fields = append(fields, usagelog.FieldUserID)
|
||||
}
|
||||
@@ -20267,6 +20317,9 @@ func (m *UsageLogMutation) Fields() []string {
|
||||
if m.model != nil {
|
||||
fields = append(fields, usagelog.FieldModel)
|
||||
}
|
||||
if m.upstream_model != nil {
|
||||
fields = append(fields, usagelog.FieldUpstreamModel)
|
||||
}
|
||||
if m.group != nil {
|
||||
fields = append(fields, usagelog.FieldGroupID)
|
||||
}
|
||||
@@ -20366,6 +20419,8 @@ func (m *UsageLogMutation) Field(name string) (ent.Value, bool) {
|
||||
return m.RequestID()
|
||||
case usagelog.FieldModel:
|
||||
return m.Model()
|
||||
case usagelog.FieldUpstreamModel:
|
||||
return m.UpstreamModel()
|
||||
case usagelog.FieldGroupID:
|
||||
return m.GroupID()
|
||||
case usagelog.FieldSubscriptionID:
|
||||
@@ -20439,6 +20494,8 @@ func (m *UsageLogMutation) OldField(ctx context.Context, name string) (ent.Value
|
||||
return m.OldRequestID(ctx)
|
||||
case usagelog.FieldModel:
|
||||
return m.OldModel(ctx)
|
||||
case usagelog.FieldUpstreamModel:
|
||||
return m.OldUpstreamModel(ctx)
|
||||
case usagelog.FieldGroupID:
|
||||
return m.OldGroupID(ctx)
|
||||
case usagelog.FieldSubscriptionID:
|
||||
@@ -20537,6 +20594,13 @@ func (m *UsageLogMutation) SetField(name string, value ent.Value) error {
|
||||
}
|
||||
m.SetModel(v)
|
||||
return nil
|
||||
case usagelog.FieldUpstreamModel:
|
||||
v, ok := value.(string)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetUpstreamModel(v)
|
||||
return nil
|
||||
case usagelog.FieldGroupID:
|
||||
v, ok := value.(int64)
|
||||
if !ok {
|
||||
@@ -20975,6 +21039,9 @@ func (m *UsageLogMutation) AddField(name string, value ent.Value) error {
|
||||
// mutation.
|
||||
func (m *UsageLogMutation) ClearedFields() []string {
|
||||
var fields []string
|
||||
if m.FieldCleared(usagelog.FieldUpstreamModel) {
|
||||
fields = append(fields, usagelog.FieldUpstreamModel)
|
||||
}
|
||||
if m.FieldCleared(usagelog.FieldGroupID) {
|
||||
fields = append(fields, usagelog.FieldGroupID)
|
||||
}
|
||||
@@ -21016,6 +21083,9 @@ func (m *UsageLogMutation) FieldCleared(name string) bool {
|
||||
// error if the field is not defined in the schema.
|
||||
func (m *UsageLogMutation) ClearField(name string) error {
|
||||
switch name {
|
||||
case usagelog.FieldUpstreamModel:
|
||||
m.ClearUpstreamModel()
|
||||
return nil
|
||||
case usagelog.FieldGroupID:
|
||||
m.ClearGroupID()
|
||||
return nil
|
||||
@@ -21066,6 +21136,9 @@ func (m *UsageLogMutation) ResetField(name string) error {
|
||||
case usagelog.FieldModel:
|
||||
m.ResetModel()
|
||||
return nil
|
||||
case usagelog.FieldUpstreamModel:
|
||||
m.ResetUpstreamModel()
|
||||
return nil
|
||||
case usagelog.FieldGroupID:
|
||||
m.ResetGroupID()
|
||||
return nil
|
||||
|
||||
@@ -825,92 +825,96 @@ func init() {
|
||||
return nil
|
||||
}
|
||||
}()
|
||||
// usagelogDescUpstreamModel is the schema descriptor for upstream_model field.
|
||||
usagelogDescUpstreamModel := usagelogFields[5].Descriptor()
|
||||
// usagelog.UpstreamModelValidator is a validator for the "upstream_model" field. It is called by the builders before save.
|
||||
usagelog.UpstreamModelValidator = usagelogDescUpstreamModel.Validators[0].(func(string) error)
|
||||
// usagelogDescInputTokens is the schema descriptor for input_tokens field.
|
||||
usagelogDescInputTokens := usagelogFields[7].Descriptor()
|
||||
usagelogDescInputTokens := usagelogFields[8].Descriptor()
|
||||
// usagelog.DefaultInputTokens holds the default value on creation for the input_tokens field.
|
||||
usagelog.DefaultInputTokens = usagelogDescInputTokens.Default.(int)
|
||||
// usagelogDescOutputTokens is the schema descriptor for output_tokens field.
|
||||
usagelogDescOutputTokens := usagelogFields[8].Descriptor()
|
||||
usagelogDescOutputTokens := usagelogFields[9].Descriptor()
|
||||
// usagelog.DefaultOutputTokens holds the default value on creation for the output_tokens field.
|
||||
usagelog.DefaultOutputTokens = usagelogDescOutputTokens.Default.(int)
|
||||
// usagelogDescCacheCreationTokens is the schema descriptor for cache_creation_tokens field.
|
||||
usagelogDescCacheCreationTokens := usagelogFields[9].Descriptor()
|
||||
usagelogDescCacheCreationTokens := usagelogFields[10].Descriptor()
|
||||
// usagelog.DefaultCacheCreationTokens holds the default value on creation for the cache_creation_tokens field.
|
||||
usagelog.DefaultCacheCreationTokens = usagelogDescCacheCreationTokens.Default.(int)
|
||||
// usagelogDescCacheReadTokens is the schema descriptor for cache_read_tokens field.
|
||||
usagelogDescCacheReadTokens := usagelogFields[10].Descriptor()
|
||||
usagelogDescCacheReadTokens := usagelogFields[11].Descriptor()
|
||||
// usagelog.DefaultCacheReadTokens holds the default value on creation for the cache_read_tokens field.
|
||||
usagelog.DefaultCacheReadTokens = usagelogDescCacheReadTokens.Default.(int)
|
||||
// usagelogDescCacheCreation5mTokens is the schema descriptor for cache_creation_5m_tokens field.
|
||||
usagelogDescCacheCreation5mTokens := usagelogFields[11].Descriptor()
|
||||
usagelogDescCacheCreation5mTokens := usagelogFields[12].Descriptor()
|
||||
// usagelog.DefaultCacheCreation5mTokens holds the default value on creation for the cache_creation_5m_tokens field.
|
||||
usagelog.DefaultCacheCreation5mTokens = usagelogDescCacheCreation5mTokens.Default.(int)
|
||||
// usagelogDescCacheCreation1hTokens is the schema descriptor for cache_creation_1h_tokens field.
|
||||
usagelogDescCacheCreation1hTokens := usagelogFields[12].Descriptor()
|
||||
usagelogDescCacheCreation1hTokens := usagelogFields[13].Descriptor()
|
||||
// usagelog.DefaultCacheCreation1hTokens holds the default value on creation for the cache_creation_1h_tokens field.
|
||||
usagelog.DefaultCacheCreation1hTokens = usagelogDescCacheCreation1hTokens.Default.(int)
|
||||
// usagelogDescInputCost is the schema descriptor for input_cost field.
|
||||
usagelogDescInputCost := usagelogFields[13].Descriptor()
|
||||
usagelogDescInputCost := usagelogFields[14].Descriptor()
|
||||
// usagelog.DefaultInputCost holds the default value on creation for the input_cost field.
|
||||
usagelog.DefaultInputCost = usagelogDescInputCost.Default.(float64)
|
||||
// usagelogDescOutputCost is the schema descriptor for output_cost field.
|
||||
usagelogDescOutputCost := usagelogFields[14].Descriptor()
|
||||
usagelogDescOutputCost := usagelogFields[15].Descriptor()
|
||||
// usagelog.DefaultOutputCost holds the default value on creation for the output_cost field.
|
||||
usagelog.DefaultOutputCost = usagelogDescOutputCost.Default.(float64)
|
||||
// usagelogDescCacheCreationCost is the schema descriptor for cache_creation_cost field.
|
||||
usagelogDescCacheCreationCost := usagelogFields[15].Descriptor()
|
||||
usagelogDescCacheCreationCost := usagelogFields[16].Descriptor()
|
||||
// usagelog.DefaultCacheCreationCost holds the default value on creation for the cache_creation_cost field.
|
||||
usagelog.DefaultCacheCreationCost = usagelogDescCacheCreationCost.Default.(float64)
|
||||
// usagelogDescCacheReadCost is the schema descriptor for cache_read_cost field.
|
||||
usagelogDescCacheReadCost := usagelogFields[16].Descriptor()
|
||||
usagelogDescCacheReadCost := usagelogFields[17].Descriptor()
|
||||
// usagelog.DefaultCacheReadCost holds the default value on creation for the cache_read_cost field.
|
||||
usagelog.DefaultCacheReadCost = usagelogDescCacheReadCost.Default.(float64)
|
||||
// usagelogDescTotalCost is the schema descriptor for total_cost field.
|
||||
usagelogDescTotalCost := usagelogFields[17].Descriptor()
|
||||
usagelogDescTotalCost := usagelogFields[18].Descriptor()
|
||||
// usagelog.DefaultTotalCost holds the default value on creation for the total_cost field.
|
||||
usagelog.DefaultTotalCost = usagelogDescTotalCost.Default.(float64)
|
||||
// usagelogDescActualCost is the schema descriptor for actual_cost field.
|
||||
usagelogDescActualCost := usagelogFields[18].Descriptor()
|
||||
usagelogDescActualCost := usagelogFields[19].Descriptor()
|
||||
// usagelog.DefaultActualCost holds the default value on creation for the actual_cost field.
|
||||
usagelog.DefaultActualCost = usagelogDescActualCost.Default.(float64)
|
||||
// usagelogDescRateMultiplier is the schema descriptor for rate_multiplier field.
|
||||
usagelogDescRateMultiplier := usagelogFields[19].Descriptor()
|
||||
usagelogDescRateMultiplier := usagelogFields[20].Descriptor()
|
||||
// usagelog.DefaultRateMultiplier holds the default value on creation for the rate_multiplier field.
|
||||
usagelog.DefaultRateMultiplier = usagelogDescRateMultiplier.Default.(float64)
|
||||
// usagelogDescBillingType is the schema descriptor for billing_type field.
|
||||
usagelogDescBillingType := usagelogFields[21].Descriptor()
|
||||
usagelogDescBillingType := usagelogFields[22].Descriptor()
|
||||
// usagelog.DefaultBillingType holds the default value on creation for the billing_type field.
|
||||
usagelog.DefaultBillingType = usagelogDescBillingType.Default.(int8)
|
||||
// usagelogDescStream is the schema descriptor for stream field.
|
||||
usagelogDescStream := usagelogFields[22].Descriptor()
|
||||
usagelogDescStream := usagelogFields[23].Descriptor()
|
||||
// usagelog.DefaultStream holds the default value on creation for the stream field.
|
||||
usagelog.DefaultStream = usagelogDescStream.Default.(bool)
|
||||
// usagelogDescUserAgent is the schema descriptor for user_agent field.
|
||||
usagelogDescUserAgent := usagelogFields[25].Descriptor()
|
||||
usagelogDescUserAgent := usagelogFields[26].Descriptor()
|
||||
// usagelog.UserAgentValidator is a validator for the "user_agent" field. It is called by the builders before save.
|
||||
usagelog.UserAgentValidator = usagelogDescUserAgent.Validators[0].(func(string) error)
|
||||
// usagelogDescIPAddress is the schema descriptor for ip_address field.
|
||||
usagelogDescIPAddress := usagelogFields[26].Descriptor()
|
||||
usagelogDescIPAddress := usagelogFields[27].Descriptor()
|
||||
// usagelog.IPAddressValidator is a validator for the "ip_address" field. It is called by the builders before save.
|
||||
usagelog.IPAddressValidator = usagelogDescIPAddress.Validators[0].(func(string) error)
|
||||
// usagelogDescImageCount is the schema descriptor for image_count field.
|
||||
usagelogDescImageCount := usagelogFields[27].Descriptor()
|
||||
usagelogDescImageCount := usagelogFields[28].Descriptor()
|
||||
// usagelog.DefaultImageCount holds the default value on creation for the image_count field.
|
||||
usagelog.DefaultImageCount = usagelogDescImageCount.Default.(int)
|
||||
// usagelogDescImageSize is the schema descriptor for image_size field.
|
||||
usagelogDescImageSize := usagelogFields[28].Descriptor()
|
||||
usagelogDescImageSize := usagelogFields[29].Descriptor()
|
||||
// usagelog.ImageSizeValidator is a validator for the "image_size" field. It is called by the builders before save.
|
||||
usagelog.ImageSizeValidator = usagelogDescImageSize.Validators[0].(func(string) error)
|
||||
// usagelogDescMediaType is the schema descriptor for media_type field.
|
||||
usagelogDescMediaType := usagelogFields[29].Descriptor()
|
||||
usagelogDescMediaType := usagelogFields[30].Descriptor()
|
||||
// usagelog.MediaTypeValidator is a validator for the "media_type" field. It is called by the builders before save.
|
||||
usagelog.MediaTypeValidator = usagelogDescMediaType.Validators[0].(func(string) error)
|
||||
// usagelogDescCacheTTLOverridden is the schema descriptor for cache_ttl_overridden field.
|
||||
usagelogDescCacheTTLOverridden := usagelogFields[30].Descriptor()
|
||||
usagelogDescCacheTTLOverridden := usagelogFields[31].Descriptor()
|
||||
// usagelog.DefaultCacheTTLOverridden holds the default value on creation for the cache_ttl_overridden field.
|
||||
usagelog.DefaultCacheTTLOverridden = usagelogDescCacheTTLOverridden.Default.(bool)
|
||||
// usagelogDescCreatedAt is the schema descriptor for created_at field.
|
||||
usagelogDescCreatedAt := usagelogFields[31].Descriptor()
|
||||
usagelogDescCreatedAt := usagelogFields[32].Descriptor()
|
||||
// usagelog.DefaultCreatedAt holds the default value on creation for the created_at field.
|
||||
usagelog.DefaultCreatedAt = usagelogDescCreatedAt.Default.(func() time.Time)
|
||||
userMixin := schema.User{}.Mixin()
|
||||
|
||||
@@ -41,6 +41,12 @@ func (UsageLog) Fields() []ent.Field {
|
||||
field.String("model").
|
||||
MaxLen(100).
|
||||
NotEmpty(),
|
||||
// UpstreamModel stores the actual upstream model name when model mapping
|
||||
// is applied. NULL means no mapping — the requested model was used as-is.
|
||||
field.String("upstream_model").
|
||||
MaxLen(100).
|
||||
Optional().
|
||||
Nillable(),
|
||||
field.Int64("group_id").
|
||||
Optional().
|
||||
Nillable(),
|
||||
|
||||
+15
-1
@@ -32,6 +32,8 @@ type UsageLog struct {
|
||||
RequestID string `json:"request_id,omitempty"`
|
||||
// Model holds the value of the "model" field.
|
||||
Model string `json:"model,omitempty"`
|
||||
// UpstreamModel holds the value of the "upstream_model" field.
|
||||
UpstreamModel *string `json:"upstream_model,omitempty"`
|
||||
// GroupID holds the value of the "group_id" field.
|
||||
GroupID *int64 `json:"group_id,omitempty"`
|
||||
// SubscriptionID holds the value of the "subscription_id" field.
|
||||
@@ -175,7 +177,7 @@ func (*UsageLog) scanValues(columns []string) ([]any, error) {
|
||||
values[i] = new(sql.NullFloat64)
|
||||
case usagelog.FieldID, usagelog.FieldUserID, usagelog.FieldAPIKeyID, usagelog.FieldAccountID, usagelog.FieldGroupID, usagelog.FieldSubscriptionID, usagelog.FieldInputTokens, usagelog.FieldOutputTokens, usagelog.FieldCacheCreationTokens, usagelog.FieldCacheReadTokens, usagelog.FieldCacheCreation5mTokens, usagelog.FieldCacheCreation1hTokens, usagelog.FieldBillingType, usagelog.FieldDurationMs, usagelog.FieldFirstTokenMs, usagelog.FieldImageCount:
|
||||
values[i] = new(sql.NullInt64)
|
||||
case usagelog.FieldRequestID, usagelog.FieldModel, usagelog.FieldUserAgent, usagelog.FieldIPAddress, usagelog.FieldImageSize, usagelog.FieldMediaType:
|
||||
case usagelog.FieldRequestID, usagelog.FieldModel, usagelog.FieldUpstreamModel, usagelog.FieldUserAgent, usagelog.FieldIPAddress, usagelog.FieldImageSize, usagelog.FieldMediaType:
|
||||
values[i] = new(sql.NullString)
|
||||
case usagelog.FieldCreatedAt:
|
||||
values[i] = new(sql.NullTime)
|
||||
@@ -230,6 +232,13 @@ func (_m *UsageLog) assignValues(columns []string, values []any) error {
|
||||
} else if value.Valid {
|
||||
_m.Model = value.String
|
||||
}
|
||||
case usagelog.FieldUpstreamModel:
|
||||
if value, ok := values[i].(*sql.NullString); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field upstream_model", values[i])
|
||||
} else if value.Valid {
|
||||
_m.UpstreamModel = new(string)
|
||||
*_m.UpstreamModel = value.String
|
||||
}
|
||||
case usagelog.FieldGroupID:
|
||||
if value, ok := values[i].(*sql.NullInt64); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field group_id", values[i])
|
||||
@@ -477,6 +486,11 @@ func (_m *UsageLog) String() string {
|
||||
builder.WriteString("model=")
|
||||
builder.WriteString(_m.Model)
|
||||
builder.WriteString(", ")
|
||||
if v := _m.UpstreamModel; v != nil {
|
||||
builder.WriteString("upstream_model=")
|
||||
builder.WriteString(*v)
|
||||
}
|
||||
builder.WriteString(", ")
|
||||
if v := _m.GroupID; v != nil {
|
||||
builder.WriteString("group_id=")
|
||||
builder.WriteString(fmt.Sprintf("%v", *v))
|
||||
|
||||
@@ -24,6 +24,8 @@ const (
|
||||
FieldRequestID = "request_id"
|
||||
// FieldModel holds the string denoting the model field in the database.
|
||||
FieldModel = "model"
|
||||
// FieldUpstreamModel holds the string denoting the upstream_model field in the database.
|
||||
FieldUpstreamModel = "upstream_model"
|
||||
// FieldGroupID holds the string denoting the group_id field in the database.
|
||||
FieldGroupID = "group_id"
|
||||
// FieldSubscriptionID holds the string denoting the subscription_id field in the database.
|
||||
@@ -135,6 +137,7 @@ var Columns = []string{
|
||||
FieldAccountID,
|
||||
FieldRequestID,
|
||||
FieldModel,
|
||||
FieldUpstreamModel,
|
||||
FieldGroupID,
|
||||
FieldSubscriptionID,
|
||||
FieldInputTokens,
|
||||
@@ -179,6 +182,8 @@ var (
|
||||
RequestIDValidator func(string) error
|
||||
// ModelValidator is a validator for the "model" field. It is called by the builders before save.
|
||||
ModelValidator func(string) error
|
||||
// UpstreamModelValidator is a validator for the "upstream_model" field. It is called by the builders before save.
|
||||
UpstreamModelValidator func(string) error
|
||||
// DefaultInputTokens holds the default value on creation for the "input_tokens" field.
|
||||
DefaultInputTokens int
|
||||
// DefaultOutputTokens holds the default value on creation for the "output_tokens" field.
|
||||
@@ -258,6 +263,11 @@ func ByModel(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldModel, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByUpstreamModel orders the results by the upstream_model field.
|
||||
func ByUpstreamModel(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldUpstreamModel, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByGroupID orders the results by the group_id field.
|
||||
func ByGroupID(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldGroupID, opts...).ToFunc()
|
||||
|
||||
@@ -80,6 +80,11 @@ func Model(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModel applies equality check predicate on the "upstream_model" field. It's identical to UpstreamModelEQ.
|
||||
func UpstreamModel(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// GroupID applies equality check predicate on the "group_id" field. It's identical to GroupIDEQ.
|
||||
func GroupID(v int64) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldGroupID, v))
|
||||
@@ -405,6 +410,81 @@ func ModelContainsFold(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldContainsFold(FieldModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelEQ applies the EQ predicate on the "upstream_model" field.
|
||||
func UpstreamModelEQ(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelNEQ applies the NEQ predicate on the "upstream_model" field.
|
||||
func UpstreamModelNEQ(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldNEQ(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelIn applies the In predicate on the "upstream_model" field.
|
||||
func UpstreamModelIn(vs ...string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldIn(FieldUpstreamModel, vs...))
|
||||
}
|
||||
|
||||
// UpstreamModelNotIn applies the NotIn predicate on the "upstream_model" field.
|
||||
func UpstreamModelNotIn(vs ...string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldNotIn(FieldUpstreamModel, vs...))
|
||||
}
|
||||
|
||||
// UpstreamModelGT applies the GT predicate on the "upstream_model" field.
|
||||
func UpstreamModelGT(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldGT(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelGTE applies the GTE predicate on the "upstream_model" field.
|
||||
func UpstreamModelGTE(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldGTE(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelLT applies the LT predicate on the "upstream_model" field.
|
||||
func UpstreamModelLT(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldLT(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelLTE applies the LTE predicate on the "upstream_model" field.
|
||||
func UpstreamModelLTE(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldLTE(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelContains applies the Contains predicate on the "upstream_model" field.
|
||||
func UpstreamModelContains(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldContains(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelHasPrefix applies the HasPrefix predicate on the "upstream_model" field.
|
||||
func UpstreamModelHasPrefix(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldHasPrefix(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelHasSuffix applies the HasSuffix predicate on the "upstream_model" field.
|
||||
func UpstreamModelHasSuffix(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldHasSuffix(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelIsNil applies the IsNil predicate on the "upstream_model" field.
|
||||
func UpstreamModelIsNil() predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldIsNull(FieldUpstreamModel))
|
||||
}
|
||||
|
||||
// UpstreamModelNotNil applies the NotNil predicate on the "upstream_model" field.
|
||||
func UpstreamModelNotNil() predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldNotNull(FieldUpstreamModel))
|
||||
}
|
||||
|
||||
// UpstreamModelEqualFold applies the EqualFold predicate on the "upstream_model" field.
|
||||
func UpstreamModelEqualFold(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEqualFold(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelContainsFold applies the ContainsFold predicate on the "upstream_model" field.
|
||||
func UpstreamModelContainsFold(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldContainsFold(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// GroupIDEQ applies the EQ predicate on the "group_id" field.
|
||||
func GroupIDEQ(v int64) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldGroupID, v))
|
||||
|
||||
@@ -57,6 +57,20 @@ func (_c *UsageLogCreate) SetModel(v string) *UsageLogCreate {
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetUpstreamModel sets the "upstream_model" field.
|
||||
func (_c *UsageLogCreate) SetUpstreamModel(v string) *UsageLogCreate {
|
||||
_c.mutation.SetUpstreamModel(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetNillableUpstreamModel sets the "upstream_model" field if the given value is not nil.
|
||||
func (_c *UsageLogCreate) SetNillableUpstreamModel(v *string) *UsageLogCreate {
|
||||
if v != nil {
|
||||
_c.SetUpstreamModel(*v)
|
||||
}
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetGroupID sets the "group_id" field.
|
||||
func (_c *UsageLogCreate) SetGroupID(v int64) *UsageLogCreate {
|
||||
_c.mutation.SetGroupID(v)
|
||||
@@ -596,6 +610,11 @@ func (_c *UsageLogCreate) check() error {
|
||||
return &ValidationError{Name: "model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _c.mutation.UpstreamModel(); ok {
|
||||
if err := usagelog.UpstreamModelValidator(v); err != nil {
|
||||
return &ValidationError{Name: "upstream_model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.upstream_model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if _, ok := _c.mutation.InputTokens(); !ok {
|
||||
return &ValidationError{Name: "input_tokens", err: errors.New(`ent: missing required field "UsageLog.input_tokens"`)}
|
||||
}
|
||||
@@ -714,6 +733,10 @@ func (_c *UsageLogCreate) createSpec() (*UsageLog, *sqlgraph.CreateSpec) {
|
||||
_spec.SetField(usagelog.FieldModel, field.TypeString, value)
|
||||
_node.Model = value
|
||||
}
|
||||
if value, ok := _c.mutation.UpstreamModel(); ok {
|
||||
_spec.SetField(usagelog.FieldUpstreamModel, field.TypeString, value)
|
||||
_node.UpstreamModel = &value
|
||||
}
|
||||
if value, ok := _c.mutation.InputTokens(); ok {
|
||||
_spec.SetField(usagelog.FieldInputTokens, field.TypeInt, value)
|
||||
_node.InputTokens = value
|
||||
@@ -1011,6 +1034,24 @@ func (u *UsageLogUpsert) UpdateModel() *UsageLogUpsert {
|
||||
return u
|
||||
}
|
||||
|
||||
// SetUpstreamModel sets the "upstream_model" field.
|
||||
func (u *UsageLogUpsert) SetUpstreamModel(v string) *UsageLogUpsert {
|
||||
u.Set(usagelog.FieldUpstreamModel, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateUpstreamModel sets the "upstream_model" field to the value that was provided on create.
|
||||
func (u *UsageLogUpsert) UpdateUpstreamModel() *UsageLogUpsert {
|
||||
u.SetExcluded(usagelog.FieldUpstreamModel)
|
||||
return u
|
||||
}
|
||||
|
||||
// ClearUpstreamModel clears the value of the "upstream_model" field.
|
||||
func (u *UsageLogUpsert) ClearUpstreamModel() *UsageLogUpsert {
|
||||
u.SetNull(usagelog.FieldUpstreamModel)
|
||||
return u
|
||||
}
|
||||
|
||||
// SetGroupID sets the "group_id" field.
|
||||
func (u *UsageLogUpsert) SetGroupID(v int64) *UsageLogUpsert {
|
||||
u.Set(usagelog.FieldGroupID, v)
|
||||
@@ -1600,6 +1641,27 @@ func (u *UsageLogUpsertOne) UpdateModel() *UsageLogUpsertOne {
|
||||
})
|
||||
}
|
||||
|
||||
// SetUpstreamModel sets the "upstream_model" field.
|
||||
func (u *UsageLogUpsertOne) SetUpstreamModel(v string) *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.SetUpstreamModel(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateUpstreamModel sets the "upstream_model" field to the value that was provided on create.
|
||||
func (u *UsageLogUpsertOne) UpdateUpstreamModel() *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.UpdateUpstreamModel()
|
||||
})
|
||||
}
|
||||
|
||||
// ClearUpstreamModel clears the value of the "upstream_model" field.
|
||||
func (u *UsageLogUpsertOne) ClearUpstreamModel() *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.ClearUpstreamModel()
|
||||
})
|
||||
}
|
||||
|
||||
// SetGroupID sets the "group_id" field.
|
||||
func (u *UsageLogUpsertOne) SetGroupID(v int64) *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
@@ -2434,6 +2496,27 @@ func (u *UsageLogUpsertBulk) UpdateModel() *UsageLogUpsertBulk {
|
||||
})
|
||||
}
|
||||
|
||||
// SetUpstreamModel sets the "upstream_model" field.
|
||||
func (u *UsageLogUpsertBulk) SetUpstreamModel(v string) *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.SetUpstreamModel(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateUpstreamModel sets the "upstream_model" field to the value that was provided on create.
|
||||
func (u *UsageLogUpsertBulk) UpdateUpstreamModel() *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.UpdateUpstreamModel()
|
||||
})
|
||||
}
|
||||
|
||||
// ClearUpstreamModel clears the value of the "upstream_model" field.
|
||||
func (u *UsageLogUpsertBulk) ClearUpstreamModel() *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.ClearUpstreamModel()
|
||||
})
|
||||
}
|
||||
|
||||
// SetGroupID sets the "group_id" field.
|
||||
func (u *UsageLogUpsertBulk) SetGroupID(v int64) *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
|
||||
@@ -102,6 +102,26 @@ func (_u *UsageLogUpdate) SetNillableModel(v *string) *UsageLogUpdate {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetUpstreamModel sets the "upstream_model" field.
|
||||
func (_u *UsageLogUpdate) SetUpstreamModel(v string) *UsageLogUpdate {
|
||||
_u.mutation.SetUpstreamModel(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableUpstreamModel sets the "upstream_model" field if the given value is not nil.
|
||||
func (_u *UsageLogUpdate) SetNillableUpstreamModel(v *string) *UsageLogUpdate {
|
||||
if v != nil {
|
||||
_u.SetUpstreamModel(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// ClearUpstreamModel clears the value of the "upstream_model" field.
|
||||
func (_u *UsageLogUpdate) ClearUpstreamModel() *UsageLogUpdate {
|
||||
_u.mutation.ClearUpstreamModel()
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetGroupID sets the "group_id" field.
|
||||
func (_u *UsageLogUpdate) SetGroupID(v int64) *UsageLogUpdate {
|
||||
_u.mutation.SetGroupID(v)
|
||||
@@ -745,6 +765,11 @@ func (_u *UsageLogUpdate) check() error {
|
||||
return &ValidationError{Name: "model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _u.mutation.UpstreamModel(); ok {
|
||||
if err := usagelog.UpstreamModelValidator(v); err != nil {
|
||||
return &ValidationError{Name: "upstream_model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.upstream_model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _u.mutation.UserAgent(); ok {
|
||||
if err := usagelog.UserAgentValidator(v); err != nil {
|
||||
return &ValidationError{Name: "user_agent", err: fmt.Errorf(`ent: validator failed for field "UsageLog.user_agent": %w`, err)}
|
||||
@@ -795,6 +820,12 @@ func (_u *UsageLogUpdate) sqlSave(ctx context.Context) (_node int, err error) {
|
||||
if value, ok := _u.mutation.Model(); ok {
|
||||
_spec.SetField(usagelog.FieldModel, field.TypeString, value)
|
||||
}
|
||||
if value, ok := _u.mutation.UpstreamModel(); ok {
|
||||
_spec.SetField(usagelog.FieldUpstreamModel, field.TypeString, value)
|
||||
}
|
||||
if _u.mutation.UpstreamModelCleared() {
|
||||
_spec.ClearField(usagelog.FieldUpstreamModel, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.InputTokens(); ok {
|
||||
_spec.SetField(usagelog.FieldInputTokens, field.TypeInt, value)
|
||||
}
|
||||
@@ -1177,6 +1208,26 @@ func (_u *UsageLogUpdateOne) SetNillableModel(v *string) *UsageLogUpdateOne {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetUpstreamModel sets the "upstream_model" field.
|
||||
func (_u *UsageLogUpdateOne) SetUpstreamModel(v string) *UsageLogUpdateOne {
|
||||
_u.mutation.SetUpstreamModel(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableUpstreamModel sets the "upstream_model" field if the given value is not nil.
|
||||
func (_u *UsageLogUpdateOne) SetNillableUpstreamModel(v *string) *UsageLogUpdateOne {
|
||||
if v != nil {
|
||||
_u.SetUpstreamModel(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// ClearUpstreamModel clears the value of the "upstream_model" field.
|
||||
func (_u *UsageLogUpdateOne) ClearUpstreamModel() *UsageLogUpdateOne {
|
||||
_u.mutation.ClearUpstreamModel()
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetGroupID sets the "group_id" field.
|
||||
func (_u *UsageLogUpdateOne) SetGroupID(v int64) *UsageLogUpdateOne {
|
||||
_u.mutation.SetGroupID(v)
|
||||
@@ -1833,6 +1884,11 @@ func (_u *UsageLogUpdateOne) check() error {
|
||||
return &ValidationError{Name: "model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _u.mutation.UpstreamModel(); ok {
|
||||
if err := usagelog.UpstreamModelValidator(v); err != nil {
|
||||
return &ValidationError{Name: "upstream_model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.upstream_model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _u.mutation.UserAgent(); ok {
|
||||
if err := usagelog.UserAgentValidator(v); err != nil {
|
||||
return &ValidationError{Name: "user_agent", err: fmt.Errorf(`ent: validator failed for field "UsageLog.user_agent": %w`, err)}
|
||||
@@ -1900,6 +1956,12 @@ func (_u *UsageLogUpdateOne) sqlSave(ctx context.Context) (_node *UsageLog, err
|
||||
if value, ok := _u.mutation.Model(); ok {
|
||||
_spec.SetField(usagelog.FieldModel, field.TypeString, value)
|
||||
}
|
||||
if value, ok := _u.mutation.UpstreamModel(); ok {
|
||||
_spec.SetField(usagelog.FieldUpstreamModel, field.TypeString, value)
|
||||
}
|
||||
if _u.mutation.UpstreamModelCleared() {
|
||||
_spec.ClearField(usagelog.FieldUpstreamModel, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.InputTokens(); ok {
|
||||
_spec.SetField(usagelog.FieldInputTokens, field.TypeInt, value)
|
||||
}
|
||||
|
||||
+4
-9
@@ -26,8 +26,6 @@ github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwTo
|
||||
github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
|
||||
github.com/apparentlymart/go-textseg/v15 v15.0.0 h1:uYvfpb3DyLSCGWnctWKGj857c6ew1u1fNQOlOtuGxQY=
|
||||
github.com/apparentlymart/go-textseg/v15 v15.0.0/go.mod h1:K8XmNZdhEBkdlyDdvbmmsvpAG721bKi0joRfFdHIWJ4=
|
||||
github.com/aws/aws-sdk-go-v2 v1.41.2 h1:LuT2rzqNQsauaGkPK/7813XxcZ3o3yePY0Iy891T2ls=
|
||||
github.com/aws/aws-sdk-go-v2 v1.41.2/go.mod h1:IvvlAZQXvTXznUPfRVfryiG1fbzE2NGK6m9u39YQ+S4=
|
||||
github.com/aws/aws-sdk-go-v2 v1.41.3 h1:4kQ/fa22KjDt13QCy1+bYADvdgcxpfH18f0zP542kZA=
|
||||
github.com/aws/aws-sdk-go-v2 v1.41.3/go.mod h1:mwsPRE8ceUUpiTgF7QmQIJ7lgsKUPQOUl3o72QBrE1o=
|
||||
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.5 h1:zWFmPmgw4sveAYi1mRqG+E/g0461cJ5M4bJ8/nc6d3Q=
|
||||
@@ -64,8 +62,6 @@ github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.15 h1:edCcNp9eGIUDUCrzoCu1jWA
|
||||
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.15/go.mod h1:lyRQKED9xWfgkYC/wmmYfv7iVIM68Z5OQ88ZdcV1QbU=
|
||||
github.com/aws/aws-sdk-go-v2/service/sts v1.41.7 h1:NITQpgo9A5NrDZ57uOWj+abvXSb83BbyggcUBVksN7c=
|
||||
github.com/aws/aws-sdk-go-v2/service/sts v1.41.7/go.mod h1:sks5UWBhEuWYDPdwlnRFn1w7xWdH29Jcpe+/PJQefEs=
|
||||
github.com/aws/smithy-go v1.24.1 h1:VbyeNfmYkWoxMVpGUAbQumkODcYmfMRfZ8yQiH30SK0=
|
||||
github.com/aws/smithy-go v1.24.1/go.mod h1:LEj2LM3rBRQJxPZTB4KuzZkaZYnZPnvgIhb4pu07mx0=
|
||||
github.com/aws/smithy-go v1.24.2 h1:FzA3bu/nt/vDvmnkg+R8Xl46gmzEDam6mZ1hzmwXFng=
|
||||
github.com/aws/smithy-go v1.24.2/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc=
|
||||
github.com/bdandy/go-errors v1.2.2 h1:WdFv/oukjTJCLa79UfkGmwX7ZxONAihKu4V0mLIs11Q=
|
||||
@@ -217,7 +213,6 @@ github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17k
|
||||
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
|
||||
github.com/google/s2a-go v0.1.7 h1:60BLSyTrOV4/haCDW4zb1guZItoSq8foHCXrAnjBo/o=
|
||||
github.com/google/s2a-go v0.1.7/go.mod h1:50CgR4k1jNlWBu4UfS4AcfhVe1r6pdZPygJ3R8F0Qdw=
|
||||
github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk=
|
||||
github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
@@ -375,8 +370,6 @@ github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSS
|
||||
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
|
||||
github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
|
||||
github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
|
||||
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/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
@@ -385,6 +378,8 @@ github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o
|
||||
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=
|
||||
github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
|
||||
github.com/tam7t/hpkp v0.0.0-20160821193359-2b70b4024ed5 h1:YqAladjX7xpA6BM04leXMWAEjS0mTZ5kUU9KRBriQJc=
|
||||
@@ -500,11 +495,11 @@ golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBc
|
||||
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=
|
||||
golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
|
||||
golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
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/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.40.0 h1:36e4zGLqU4yhjlmxEaagx2KuYbJq3EwY8K943ZsHcvg=
|
||||
golang.org/x/term v0.40.0/go.mod h1:w2P8uVp06p2iyKKuvXIm7N/y0UCRt3UfJTfZ7oOpglM=
|
||||
|
||||
@@ -82,8 +82,8 @@ var DefaultAntigravityModelMapping = map[string]string{
|
||||
"claude-opus-4-5-20251101": "claude-opus-4-6-thinking", // 迁移旧模型
|
||||
"claude-sonnet-4-5-20250929": "claude-sonnet-4-5",
|
||||
// Claude Haiku → Sonnet(无 Haiku 支持)
|
||||
"claude-haiku-4-5": "claude-sonnet-4-5",
|
||||
"claude-haiku-4-5-20251001": "claude-sonnet-4-5",
|
||||
"claude-haiku-4-5": "claude-sonnet-4-6",
|
||||
"claude-haiku-4-5-20251001": "claude-sonnet-4-6",
|
||||
// Gemini 2.5 白名单
|
||||
"gemini-2.5-flash": "gemini-2.5-flash",
|
||||
"gemini-2.5-flash-image": "gemini-2.5-flash-image",
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// 账号亲和调度 API 与契约说明见 docs/ACCOUNT_SCHEDULING_FLOW.md 。
|
||||
|
||||
// GetAffinityClients returns the list of affinity clients for an account with last active timestamps.
|
||||
// GET /api/v1/admin/accounts/:id/affinity-clients
|
||||
func (h *AccountHandler) GetAffinityClients(c *gin.Context) {
|
||||
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
|
||||
account, err := h.adminService.GetAccount(c.Request.Context(), accountID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
if !account.IsAffinityEnabled() {
|
||||
response.Success(c, []service.AffinityClient{})
|
||||
return
|
||||
}
|
||||
|
||||
if h.gatewayCache == nil || len(account.GroupIDs) == 0 {
|
||||
response.Success(c, []service.AffinityClient{})
|
||||
return
|
||||
}
|
||||
|
||||
clients, err := h.gatewayCache.GetAccountAffinityClientsWithScores(
|
||||
c.Request.Context(), accountID, account.GroupIDs, service.ClientAffinityTTL,
|
||||
)
|
||||
if err != nil {
|
||||
response.Success(c, []service.AffinityClient{})
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, clients)
|
||||
}
|
||||
|
||||
// clearAccountAffinity 清除指定账号在所有分组的亲和记录。
|
||||
func (h *AccountHandler) clearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) {
|
||||
if h.gatewayCache == nil || len(groupIDs) == 0 {
|
||||
return
|
||||
}
|
||||
if err := h.gatewayCache.ClearAccountAffinity(ctx, accountID, groupIDs); err != nil {
|
||||
// 清理失败不影响主流程,记录日志即可
|
||||
slog.Warn("clear account affinity failed",
|
||||
"account_id", accountID,
|
||||
"error", err,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// countUniqueUsersFromAffinityMembers 从 "{userID}/{clientID}" 格式的成员列表中计算唯一用户数。
|
||||
func countUniqueUsersFromAffinityMembers(members []string) int64 {
|
||||
users := make(map[string]struct{}, len(members))
|
||||
for _, m := range members {
|
||||
if idx := strings.Index(m, "/"); idx > 0 {
|
||||
users[m[:idx]] = struct{}{}
|
||||
}
|
||||
}
|
||||
return int64(len(users))
|
||||
}
|
||||
|
||||
// AffinityDetailClient 亲和详情中的单个客户端信息
|
||||
type AffinityDetailClient struct {
|
||||
ClientID string `json:"client_id"`
|
||||
LastActive time.Time `json:"last_active"`
|
||||
IsPinned bool `json:"is_pinned"`
|
||||
}
|
||||
|
||||
// AffinityDetailUser 亲和详情中按用户分组的信息
|
||||
type AffinityDetailUser struct {
|
||||
UserID int64 `json:"user_id"`
|
||||
UserEmail string `json:"user_email"`
|
||||
ClientCount int `json:"client_count"`
|
||||
IsPinned bool `json:"is_pinned"`
|
||||
Clients []AffinityDetailClient `json:"clients"`
|
||||
}
|
||||
|
||||
// AffinityDetailsResponse 亲和详情响应
|
||||
type AffinityDetailsResponse struct {
|
||||
Users []AffinityDetailUser `json:"users"`
|
||||
TotalUsers int `json:"total_users"`
|
||||
TotalClients int `json:"total_clients"`
|
||||
PinnedUsers []int64 `json:"pinned_users"`
|
||||
}
|
||||
|
||||
type affinityStateSnapshot struct {
|
||||
enabled bool
|
||||
groupIDs []int64
|
||||
}
|
||||
|
||||
func (h *AccountHandler) captureAffinityStates(ctx context.Context, accountIDs []int64) map[int64]affinityStateSnapshot {
|
||||
states := make(map[int64]affinityStateSnapshot, len(accountIDs))
|
||||
if h.gatewayCache == nil || len(accountIDs) == 0 {
|
||||
return states
|
||||
}
|
||||
for _, accountID := range accountIDs {
|
||||
account, err := h.adminService.GetAccount(ctx, accountID)
|
||||
if err != nil || account == nil {
|
||||
continue
|
||||
}
|
||||
states[accountID] = affinityStateSnapshot{
|
||||
enabled: account.IsAffinityEnabled(),
|
||||
groupIDs: append([]int64(nil), account.GroupIDs...),
|
||||
}
|
||||
}
|
||||
return states
|
||||
}
|
||||
|
||||
func (h *AccountHandler) clearAffinityCacheIfDisabled(ctx context.Context, accountID int64, oldState affinityStateSnapshot) {
|
||||
if h.gatewayCache == nil || !oldState.enabled {
|
||||
return
|
||||
}
|
||||
account, err := h.adminService.GetAccount(ctx, accountID)
|
||||
if err != nil || account == nil || account.IsAffinityEnabled() {
|
||||
return
|
||||
}
|
||||
groupIDs := oldState.groupIDs
|
||||
if len(account.GroupIDs) > 0 {
|
||||
groupIDs = mergeGroupIDs(groupIDs, account.GroupIDs)
|
||||
}
|
||||
h.clearAccountAffinity(ctx, accountID, groupIDs)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) clearAffinityCacheForBulkIfDisabled(ctx context.Context, accountIDs []int64, oldStates map[int64]affinityStateSnapshot) {
|
||||
for _, accountID := range accountIDs {
|
||||
oldState, ok := oldStates[accountID]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
h.clearAffinityCacheIfDisabled(ctx, accountID, oldState)
|
||||
}
|
||||
}
|
||||
|
||||
// GetAffinityDetails returns the affinity details grouped by user for an account.
|
||||
// GET /api/v1/admin/accounts/:id/affinity-details
|
||||
func (h *AccountHandler) GetAffinityDetails(c *gin.Context) {
|
||||
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
|
||||
emptyResp := AffinityDetailsResponse{
|
||||
Users: []AffinityDetailUser{},
|
||||
TotalUsers: 0,
|
||||
TotalClients: 0,
|
||||
PinnedUsers: []int64{},
|
||||
}
|
||||
|
||||
account, err := h.adminService.GetAccount(c.Request.Context(), accountID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
if !account.IsAffinityEnabled() {
|
||||
response.Success(c, emptyResp)
|
||||
return
|
||||
}
|
||||
|
||||
pinnedUsers := account.GetPinnedUsers()
|
||||
if pinnedUsers == nil {
|
||||
pinnedUsers = []int64{}
|
||||
}
|
||||
|
||||
if h.gatewayCache == nil || len(account.GroupIDs) == 0 {
|
||||
emptyResp.PinnedUsers = pinnedUsers
|
||||
response.Success(c, emptyResp)
|
||||
return
|
||||
}
|
||||
|
||||
clients, err := h.gatewayCache.GetAccountAffinityClientsWithScores(
|
||||
c.Request.Context(), accountID, account.GroupIDs, service.ClientAffinityTTL,
|
||||
)
|
||||
if err != nil {
|
||||
emptyResp.PinnedUsers = pinnedUsers
|
||||
response.Success(c, emptyResp)
|
||||
return
|
||||
}
|
||||
|
||||
pinnedSet := make(map[int64]struct{}, len(pinnedUsers))
|
||||
for _, uid := range pinnedUsers {
|
||||
pinnedSet[uid] = struct{}{}
|
||||
}
|
||||
|
||||
// 按 UserID 分组
|
||||
userMap := make(map[int64]*AffinityDetailUser)
|
||||
var userOrder []int64
|
||||
for _, cl := range clients {
|
||||
u, ok := userMap[cl.UserID]
|
||||
if !ok {
|
||||
_, pinned := pinnedSet[cl.UserID]
|
||||
u = &AffinityDetailUser{
|
||||
UserID: cl.UserID,
|
||||
UserEmail: "",
|
||||
ClientCount: 0,
|
||||
IsPinned: pinned,
|
||||
Clients: []AffinityDetailClient{},
|
||||
}
|
||||
userMap[cl.UserID] = u
|
||||
userOrder = append(userOrder, cl.UserID)
|
||||
}
|
||||
u.Clients = append(u.Clients, AffinityDetailClient{
|
||||
ClientID: cl.ClientID,
|
||||
LastActive: cl.LastActive,
|
||||
IsPinned: false, // 客户端级别暂无 pinned 概念
|
||||
})
|
||||
}
|
||||
|
||||
// 关联用户邮箱(查询失败时保持空字符串,不影响主流程)
|
||||
for _, uid := range userOrder {
|
||||
user, uErr := h.adminService.GetUser(c.Request.Context(), uid)
|
||||
if uErr != nil || user == nil {
|
||||
continue
|
||||
}
|
||||
if u, ok := userMap[uid]; ok {
|
||||
if user.Email != "" {
|
||||
u.UserEmail = user.Email
|
||||
} else {
|
||||
u.UserEmail = user.Username
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
users := make([]AffinityDetailUser, 0, len(userOrder))
|
||||
for _, uid := range userOrder {
|
||||
u := userMap[uid]
|
||||
u.ClientCount = len(u.Clients)
|
||||
users = append(users, *u)
|
||||
}
|
||||
|
||||
response.Success(c, AffinityDetailsResponse{
|
||||
Users: users,
|
||||
TotalUsers: len(users),
|
||||
TotalClients: len(clients),
|
||||
PinnedUsers: pinnedUsers,
|
||||
})
|
||||
}
|
||||
@@ -9,7 +9,6 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -169,6 +168,8 @@ type AccountWithConcurrency struct {
|
||||
CurrentRPM *int `json:"current_rpm,omitempty"` // 当前分钟 RPM 计数
|
||||
}
|
||||
|
||||
const accountListGroupUngroupedQueryValue = "ungrouped"
|
||||
|
||||
func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, account *service.Account) AccountWithConcurrency {
|
||||
item := AccountWithConcurrency{
|
||||
Account: dto.AccountFromService(account),
|
||||
@@ -211,7 +212,7 @@ func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, ac
|
||||
}
|
||||
|
||||
// 亲和客户端数据(启用亲和的账号始终返回 count,即使为 0)
|
||||
if account.IsClientAffinityEnabled() {
|
||||
if account.IsAffinityEnabled() {
|
||||
if h.gatewayCache != nil && len(account.GroupIDs) > 0 {
|
||||
accountGroups := map[int64][]int64{account.ID: account.GroupIDs}
|
||||
if clients, err := h.gatewayCache.GetAccountAffinityClientsBatch(ctx, accountGroups, service.ClientAffinityTTL); err == nil {
|
||||
@@ -219,17 +220,22 @@ func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, ac
|
||||
count := int64(len(cl))
|
||||
item.AffinityClientCount = &count
|
||||
item.AffinityClients = cl
|
||||
userCount := countUniqueUsersFromAffinityMembers(cl)
|
||||
item.AffinityUserCount = &userCount
|
||||
} else {
|
||||
zero := int64(0)
|
||||
item.AffinityClientCount = &zero
|
||||
item.AffinityUserCount = &zero
|
||||
}
|
||||
} else {
|
||||
zero := int64(0)
|
||||
item.AffinityClientCount = &zero
|
||||
item.AffinityUserCount = &zero
|
||||
}
|
||||
} else {
|
||||
zero := int64(0)
|
||||
item.AffinityClientCount = &zero
|
||||
item.AffinityUserCount = &zero
|
||||
}
|
||||
}
|
||||
|
||||
@@ -253,7 +259,20 @@ func (h *AccountHandler) List(c *gin.Context) {
|
||||
|
||||
var groupID int64
|
||||
if groupIDStr := c.Query("group"); groupIDStr != "" {
|
||||
groupID, _ = strconv.ParseInt(groupIDStr, 10, 64)
|
||||
if groupIDStr == accountListGroupUngroupedQueryValue {
|
||||
groupID = service.AccountListGroupUngrouped
|
||||
} else {
|
||||
parsedGroupID, parseErr := strconv.ParseInt(groupIDStr, 10, 64)
|
||||
if parseErr != nil {
|
||||
response.ErrorFrom(c, infraerrors.BadRequest("INVALID_GROUP_FILTER", "invalid group filter"))
|
||||
return
|
||||
}
|
||||
if parsedGroupID < 0 {
|
||||
response.ErrorFrom(c, infraerrors.BadRequest("INVALID_GROUP_FILTER", "invalid group filter"))
|
||||
return
|
||||
}
|
||||
groupID = parsedGroupID
|
||||
}
|
||||
}
|
||||
|
||||
accounts, total, err := h.adminService.ListAccounts(c.Request.Context(), page, pageSize, platform, accountType, status, search, groupID)
|
||||
@@ -351,7 +370,7 @@ func (h *AccountHandler) List(c *gin.Context) {
|
||||
accountGroups := make(map[int64][]int64)
|
||||
for i := range accounts {
|
||||
acc := &accounts[i]
|
||||
if acc.IsClientAffinityEnabled() && len(acc.GroupIDs) > 0 {
|
||||
if acc.IsAffinityEnabled() && len(acc.GroupIDs) > 0 {
|
||||
accountGroups[acc.ID] = acc.GroupIDs
|
||||
}
|
||||
}
|
||||
@@ -391,14 +410,18 @@ func (h *AccountHandler) List(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 注入亲和客户端数据到 DTO(启用亲和的账号始终返回 count,即使为 0)
|
||||
if acc.IsClientAffinityEnabled() {
|
||||
if acc.IsAffinityEnabled() {
|
||||
if clients, ok := affinityClients[acc.ID]; ok && len(clients) > 0 {
|
||||
count := int64(len(clients))
|
||||
item.AffinityClientCount = &count
|
||||
item.AffinityClients = clients
|
||||
// 从成员列表中解析唯一用户数
|
||||
userCount := countUniqueUsersFromAffinityMembers(clients)
|
||||
item.AffinityUserCount = &userCount
|
||||
} else {
|
||||
zero := int64(0)
|
||||
item.AffinityClientCount = &zero
|
||||
item.AffinityUserCount = &zero
|
||||
}
|
||||
}
|
||||
|
||||
@@ -622,15 +645,7 @@ func (h *AccountHandler) Update(c *gin.Context) {
|
||||
// base_rpm 输入校验:负值归零,超过 10000 截断
|
||||
sanitizeExtraBaseRPM(req.Extra)
|
||||
|
||||
// 记录更新前的亲和状态,用于检测亲和关闭时清理 Redis 记录
|
||||
oldAffinityEnabled := false
|
||||
var oldGroupIDs []int64
|
||||
if len(req.Extra) > 0 && h.gatewayCache != nil {
|
||||
if oldAccount, err := h.adminService.GetAccount(c.Request.Context(), accountID); err == nil {
|
||||
oldAffinityEnabled = oldAccount.IsClientAffinityEnabled()
|
||||
oldGroupIDs = oldAccount.GroupIDs
|
||||
}
|
||||
}
|
||||
oldStates := h.captureAffinityStates(c.Request.Context(), []int64{accountID})
|
||||
|
||||
// 确定是否跳过混合渠道检查
|
||||
skipCheck := req.ConfirmMixedChannelRisk != nil && *req.ConfirmMixedChannelRisk
|
||||
@@ -668,14 +683,7 @@ func (h *AccountHandler) Update(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 亲和关闭时清理 Redis 中的亲和记录
|
||||
if oldAffinityEnabled && !account.IsClientAffinityEnabled() {
|
||||
groupIDs := oldGroupIDs
|
||||
if len(account.GroupIDs) > 0 {
|
||||
groupIDs = mergeGroupIDs(oldGroupIDs, account.GroupIDs)
|
||||
}
|
||||
h.clearAccountAffinity(c.Request.Context(), accountID, groupIDs)
|
||||
}
|
||||
h.clearAffinityCacheForBulkIfDisabled(c.Request.Context(), []int64{accountID}, oldStates)
|
||||
|
||||
response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), account))
|
||||
}
|
||||
@@ -1393,6 +1401,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
oldStates := h.captureAffinityStates(c.Request.Context(), req.AccountIDs)
|
||||
|
||||
result, err := h.adminService.BulkUpdateAccounts(c.Request.Context(), &service.BulkUpdateAccountsInput{
|
||||
AccountIDs: req.AccountIDs,
|
||||
Name: req.Name,
|
||||
@@ -1427,6 +1437,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
h.clearAffinityCacheForBulkIfDisabled(c.Request.Context(), req.AccountIDs, oldStates)
|
||||
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
@@ -1575,7 +1587,7 @@ func (h *OAuthHandler) SetupTokenCookieAuth(c *gin.Context) {
|
||||
}
|
||||
|
||||
// GetUsage handles getting account usage information
|
||||
// GET /api/v1/admin/accounts/:id/usage
|
||||
// GET /api/v1/admin/accounts/:id/usage?source=passive|active
|
||||
func (h *AccountHandler) GetUsage(c *gin.Context) {
|
||||
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
@@ -1583,7 +1595,14 @@ func (h *AccountHandler) GetUsage(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
usage, err := h.accountUsageService.GetUsage(c.Request.Context(), accountID)
|
||||
source := c.DefaultQuery("source", "active")
|
||||
|
||||
var usage *service.UsageInfo
|
||||
if source == "passive" {
|
||||
usage, err = h.accountUsageService.GetPassiveUsage(c.Request.Context(), accountID)
|
||||
} else {
|
||||
usage, err = h.accountUsageService.GetUsage(c.Request.Context(), accountID)
|
||||
}
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
@@ -1639,56 +1658,6 @@ func (h *AccountHandler) ResetQuota(c *gin.Context) {
|
||||
response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), account))
|
||||
}
|
||||
|
||||
// GetAffinityClients returns the list of affinity clients for an account with last active timestamps.
|
||||
// GET /api/v1/admin/accounts/:id/affinity-clients
|
||||
func (h *AccountHandler) GetAffinityClients(c *gin.Context) {
|
||||
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
|
||||
account, err := h.adminService.GetAccount(c.Request.Context(), accountID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
if !account.IsClientAffinityEnabled() {
|
||||
response.Success(c, []service.AffinityClient{})
|
||||
return
|
||||
}
|
||||
|
||||
if h.gatewayCache == nil || len(account.GroupIDs) == 0 {
|
||||
response.Success(c, []service.AffinityClient{})
|
||||
return
|
||||
}
|
||||
|
||||
clients, err := h.gatewayCache.GetAccountAffinityClientsWithScores(
|
||||
c.Request.Context(), accountID, account.GroupIDs, service.ClientAffinityTTL,
|
||||
)
|
||||
if err != nil {
|
||||
response.Success(c, []service.AffinityClient{})
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, clients)
|
||||
}
|
||||
|
||||
// clearAccountAffinity 清除指定账号在所有分组的亲和记录。
|
||||
func (h *AccountHandler) clearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) {
|
||||
if h.gatewayCache == nil || len(groupIDs) == 0 {
|
||||
return
|
||||
}
|
||||
if err := h.gatewayCache.ClearAccountAffinity(ctx, accountID, groupIDs); err != nil {
|
||||
// 清理失败不影响主流程,记录日志即可
|
||||
slog.Warn("clear account affinity failed",
|
||||
"account_id", accountID,
|
||||
"error", err,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// mergeGroupIDs 合并两个 groupID 切片并去重。
|
||||
func mergeGroupIDs(a, b []int64) []int64 {
|
||||
seen := make(map[int64]struct{}, len(a)+len(b))
|
||||
|
||||
@@ -0,0 +1,298 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type affinityCacheAdminService struct {
|
||||
*stubAdminService
|
||||
accounts map[int64]*service.Account
|
||||
bulkResultIDs []int64
|
||||
bulkFailedCount int
|
||||
}
|
||||
|
||||
func (s *affinityCacheAdminService) GetAccount(_ context.Context, id int64) (*service.Account, error) {
|
||||
if acc, ok := s.accounts[id]; ok {
|
||||
cloned := *acc
|
||||
return &cloned, nil
|
||||
}
|
||||
return s.stubAdminService.GetAccount(context.Background(), id)
|
||||
}
|
||||
|
||||
func (s *affinityCacheAdminService) UpdateAccount(_ context.Context, id int64, input *service.UpdateAccountInput) (*service.Account, error) {
|
||||
acc := s.accounts[id]
|
||||
if acc == nil {
|
||||
return s.stubAdminService.UpdateAccount(context.Background(), id, input)
|
||||
}
|
||||
if input.Extra != nil {
|
||||
acc.Extra = input.Extra
|
||||
}
|
||||
if input.GroupIDs != nil {
|
||||
acc.GroupIDs = *input.GroupIDs
|
||||
}
|
||||
cloned := *acc
|
||||
return &cloned, nil
|
||||
}
|
||||
|
||||
func (s *affinityCacheAdminService) BulkUpdateAccounts(_ context.Context, input *service.BulkUpdateAccountsInput) (*service.BulkUpdateAccountsResult, error) {
|
||||
for _, accountID := range input.AccountIDs {
|
||||
acc := s.accounts[accountID]
|
||||
if acc == nil {
|
||||
continue
|
||||
}
|
||||
if input.Extra != nil {
|
||||
acc.Extra = input.Extra
|
||||
}
|
||||
if input.GroupIDs != nil {
|
||||
acc.GroupIDs = *input.GroupIDs
|
||||
}
|
||||
}
|
||||
successIDs := append([]int64(nil), input.AccountIDs...)
|
||||
failed := 0
|
||||
if len(s.bulkResultIDs) > 0 {
|
||||
successIDs = append([]int64(nil), s.bulkResultIDs...)
|
||||
failed = s.bulkFailedCount
|
||||
}
|
||||
return &service.BulkUpdateAccountsResult{
|
||||
Success: len(successIDs),
|
||||
Failed: failed,
|
||||
SuccessIDs: successIDs,
|
||||
Results: []service.BulkUpdateAccountResult{},
|
||||
}, nil
|
||||
}
|
||||
|
||||
type affinityCacheClearRecorder struct {
|
||||
clearCalls []struct {
|
||||
accountID int64
|
||||
groupIDs []int64
|
||||
}
|
||||
}
|
||||
|
||||
func (r *affinityCacheClearRecorder) GetSessionAccountID(_ context.Context, _ int64, _ string) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) SetSessionAccountID(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) RefreshSessionTTL(_ context.Context, _ int64, _ string, _ time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error {
|
||||
return nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
|
||||
return map[int64]int64{}, nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) {
|
||||
return map[int64][]string{}, nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]service.AffinityClient, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) ClearAccountAffinity(_ context.Context, accountID int64, groupIDs []int64) error {
|
||||
r.clearCalls = append(r.clearCalls, struct {
|
||||
accountID int64
|
||||
groupIDs []int64
|
||||
}{
|
||||
accountID: accountID,
|
||||
groupIDs: append([]int64(nil), groupIDs...),
|
||||
})
|
||||
return nil
|
||||
}
|
||||
func (r *affinityCacheClearRecorder) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (users, clients, perUser int64, err error) {
|
||||
return 0, 0, 0, nil
|
||||
}
|
||||
|
||||
func setupAffinityCacheRouter(adminSvc service.AdminService, gatewayCache service.GatewayCache) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, gatewayCache)
|
||||
router.PUT("/api/v1/admin/accounts/:id", handler.Update)
|
||||
router.POST("/api/v1/admin/accounts/bulk-update", handler.BulkUpdate)
|
||||
return router
|
||||
}
|
||||
|
||||
func TestAccountHandlerUpdate_DisablingAffinityClearsCache(t *testing.T) {
|
||||
adminSvc := &affinityCacheAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
accounts: map[int64]*service.Account{
|
||||
3: {
|
||||
ID: 3,
|
||||
Platform: service.PlatformAnthropic,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{1, 2},
|
||||
Extra: map[string]any{
|
||||
"client_affinity_enabled": true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
cache := &affinityCacheClearRecorder{}
|
||||
router := setupAffinityCacheRouter(adminSvc, cache)
|
||||
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"extra": map[string]any{},
|
||||
})
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPut, "/api/v1/admin/accounts/3", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Len(t, cache.clearCalls, 1)
|
||||
require.Equal(t, int64(3), cache.clearCalls[0].accountID)
|
||||
require.Equal(t, []int64{1, 2}, cache.clearCalls[0].groupIDs)
|
||||
}
|
||||
|
||||
func TestAccountHandlerBulkUpdate_DisablingAffinityClearsCache(t *testing.T) {
|
||||
adminSvc := &affinityCacheAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
accounts: map[int64]*service.Account{
|
||||
1: {
|
||||
ID: 1,
|
||||
Platform: service.PlatformAnthropic,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{11},
|
||||
Extra: map[string]any{
|
||||
"client_affinity_enabled": true,
|
||||
},
|
||||
},
|
||||
2: {
|
||||
ID: 2,
|
||||
Platform: service.PlatformAnthropic,
|
||||
Type: service.AccountTypeSetupToken,
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{22},
|
||||
Extra: map[string]any{
|
||||
"client_affinity_enabled": true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
cache := &affinityCacheClearRecorder{}
|
||||
router := setupAffinityCacheRouter(adminSvc, cache)
|
||||
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"account_ids": []int64{1, 2},
|
||||
"extra": map[string]any{
|
||||
"client_affinity_enabled": false,
|
||||
},
|
||||
})
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Len(t, cache.clearCalls, 2)
|
||||
require.Equal(t, int64(1), cache.clearCalls[0].accountID)
|
||||
require.Equal(t, []int64{11}, cache.clearCalls[0].groupIDs)
|
||||
require.Equal(t, int64(2), cache.clearCalls[1].accountID)
|
||||
require.Equal(t, []int64{22}, cache.clearCalls[1].groupIDs)
|
||||
}
|
||||
|
||||
func TestAccountHandlerBulkUpdate_DisablingAffinityWithNewFlagClearsCache(t *testing.T) {
|
||||
adminSvc := &affinityCacheAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
accounts: map[int64]*service.Account{
|
||||
1: {
|
||||
ID: 1,
|
||||
Platform: service.PlatformAnthropic,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{33},
|
||||
Extra: map[string]any{
|
||||
"affinity_enabled": true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
cache := &affinityCacheClearRecorder{}
|
||||
router := setupAffinityCacheRouter(adminSvc, cache)
|
||||
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"account_ids": []int64{1},
|
||||
"extra": map[string]any{
|
||||
"affinity_enabled": false,
|
||||
},
|
||||
})
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Len(t, cache.clearCalls, 1)
|
||||
require.Equal(t, int64(1), cache.clearCalls[0].accountID)
|
||||
require.Equal(t, []int64{33}, cache.clearCalls[0].groupIDs)
|
||||
}
|
||||
|
||||
func TestAccountHandlerBulkUpdate_PartialFailureStillClearsDisabledAffinityCache(t *testing.T) {
|
||||
adminSvc := &affinityCacheAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
accounts: map[int64]*service.Account{
|
||||
1: {
|
||||
ID: 1,
|
||||
Platform: service.PlatformAnthropic,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{44},
|
||||
Extra: map[string]any{
|
||||
"affinity_enabled": true,
|
||||
},
|
||||
},
|
||||
2: {
|
||||
ID: 2,
|
||||
Platform: service.PlatformAnthropic,
|
||||
Type: service.AccountTypeSetupToken,
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{55},
|
||||
Extra: map[string]any{
|
||||
"client_affinity_enabled": true,
|
||||
},
|
||||
},
|
||||
},
|
||||
bulkResultIDs: []int64{1},
|
||||
bulkFailedCount: 1,
|
||||
}
|
||||
cache := &affinityCacheClearRecorder{}
|
||||
router := setupAffinityCacheRouter(adminSvc, cache)
|
||||
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"account_ids": []int64{1, 2},
|
||||
"group_ids": []int64{101},
|
||||
"extra": map[string]any{
|
||||
"affinity_enabled": false,
|
||||
"client_affinity_enabled": false,
|
||||
},
|
||||
})
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Len(t, cache.clearCalls, 2)
|
||||
require.Equal(t, int64(1), cache.clearCalls[0].accountID)
|
||||
require.Equal(t, []int64{44, 101}, cache.clearCalls[0].groupIDs)
|
||||
require.Equal(t, int64(2), cache.clearCalls[1].accountID)
|
||||
require.Equal(t, []int64{55, 101}, cache.clearCalls[1].groupIDs)
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type affinityDetailsAdminService struct {
|
||||
*stubAdminService
|
||||
account service.Account
|
||||
usersByID map[int64]service.User
|
||||
}
|
||||
|
||||
func (s *affinityDetailsAdminService) GetAccount(_ context.Context, id int64) (*service.Account, error) {
|
||||
if s.account.ID == id {
|
||||
acc := s.account
|
||||
return &acc, nil
|
||||
}
|
||||
return s.stubAdminService.GetAccount(context.Background(), id)
|
||||
}
|
||||
|
||||
func (s *affinityDetailsAdminService) GetUser(_ context.Context, id int64) (*service.User, error) {
|
||||
if u, ok := s.usersByID[id]; ok {
|
||||
user := u
|
||||
return &user, nil
|
||||
}
|
||||
return s.stubAdminService.GetUser(context.Background(), id)
|
||||
}
|
||||
|
||||
type affinityDetailsGatewayCacheStub struct {
|
||||
clients []service.AffinityClient
|
||||
}
|
||||
|
||||
func (s *affinityDetailsGatewayCacheStub) GetSessionAccountID(_ context.Context, _ int64, _ string) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) SetSessionAccountID(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) RefreshSessionTTL(_ context.Context, _ int64, _ string, _ time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error {
|
||||
return nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
|
||||
return map[int64]int64{}, nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) {
|
||||
return map[int64][]string{}, nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]service.AffinityClient, error) {
|
||||
return s.clients, nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error {
|
||||
return nil
|
||||
}
|
||||
func (s *affinityDetailsGatewayCacheStub) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (users, clients, perUser int64, err error) {
|
||||
return 0, 0, 0, nil
|
||||
}
|
||||
|
||||
func setupAffinityDetailsRouter(adminSvc service.AdminService, gatewayCache service.GatewayCache) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, gatewayCache)
|
||||
router.GET("/api/v1/admin/accounts/:id/affinity-details", handler.GetAffinityDetails)
|
||||
return router
|
||||
}
|
||||
|
||||
func TestAccountHandlerGetAffinityDetails_ContractFields(t *testing.T) {
|
||||
adminSvc := &affinityDetailsAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
account: service.Account{
|
||||
ID: 88,
|
||||
Platform: service.PlatformAnthropic,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{1},
|
||||
Extra: map[string]any{
|
||||
"client_affinity_enabled": true,
|
||||
"pinned_users": []any{float64(42), float64(99)},
|
||||
},
|
||||
},
|
||||
usersByID: map[int64]service.User{
|
||||
42: {ID: 42, Email: "user42@example.com"},
|
||||
7: {ID: 7, Username: "user7"},
|
||||
},
|
||||
}
|
||||
cache := &affinityDetailsGatewayCacheStub{
|
||||
clients: []service.AffinityClient{
|
||||
{UserID: 42, ClientID: "client-a", LastActive: time.Now().Add(-time.Minute)},
|
||||
{UserID: 7, ClientID: "client-b", LastActive: time.Now().Add(-2 * time.Minute)},
|
||||
{UserID: 42, ClientID: "client-c", LastActive: time.Now().Add(-3 * time.Minute)},
|
||||
},
|
||||
}
|
||||
router := setupAffinityDetailsRouter(adminSvc, cache)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/88/affinity-details", nil)
|
||||
router.ServeHTTP(rec, req)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var resp struct {
|
||||
Code int `json:"code"`
|
||||
Data struct {
|
||||
Users []struct {
|
||||
UserID int64 `json:"user_id"`
|
||||
UserEmail string `json:"user_email"`
|
||||
ClientCount int `json:"client_count"`
|
||||
IsPinned bool `json:"is_pinned"`
|
||||
Clients []struct {
|
||||
ClientID string `json:"client_id"`
|
||||
} `json:"clients"`
|
||||
} `json:"users"`
|
||||
TotalUsers int `json:"total_users"`
|
||||
TotalClients int `json:"total_clients"`
|
||||
PinnedUsers []int64 `json:"pinned_users"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, 0, resp.Code)
|
||||
require.Equal(t, 2, resp.Data.TotalUsers)
|
||||
require.Equal(t, 3, resp.Data.TotalClients)
|
||||
require.Equal(t, []int64{42, 99}, resp.Data.PinnedUsers)
|
||||
require.Len(t, resp.Data.Users, 2)
|
||||
|
||||
usersByID := make(map[int64]struct {
|
||||
UserEmail string
|
||||
ClientCount int
|
||||
IsPinned bool
|
||||
ClientLen int
|
||||
}, len(resp.Data.Users))
|
||||
for _, u := range resp.Data.Users {
|
||||
usersByID[u.UserID] = struct {
|
||||
UserEmail string
|
||||
ClientCount int
|
||||
IsPinned bool
|
||||
ClientLen int
|
||||
}{
|
||||
UserEmail: u.UserEmail,
|
||||
ClientCount: u.ClientCount,
|
||||
IsPinned: u.IsPinned,
|
||||
ClientLen: len(u.Clients),
|
||||
}
|
||||
}
|
||||
require.Equal(t, "user42@example.com", usersByID[42].UserEmail)
|
||||
require.Equal(t, 2, usersByID[42].ClientCount)
|
||||
require.True(t, usersByID[42].IsPinned)
|
||||
require.Equal(t, 2, usersByID[42].ClientLen)
|
||||
|
||||
require.Equal(t, "user7", usersByID[7].UserEmail)
|
||||
require.Equal(t, 1, usersByID[7].ClientCount)
|
||||
require.False(t, usersByID[7].IsPinned)
|
||||
require.Equal(t, 1, usersByID[7].ClientLen)
|
||||
}
|
||||
|
||||
func TestAccountHandlerGetAffinityDetails_DisabledReturnsEmptyContract(t *testing.T) {
|
||||
adminSvc := &affinityDetailsAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
account: service.Account{
|
||||
ID: 89,
|
||||
Platform: service.PlatformAnthropic,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{1},
|
||||
},
|
||||
usersByID: map[int64]service.User{},
|
||||
}
|
||||
router := setupAffinityDetailsRouter(adminSvc, &affinityDetailsGatewayCacheStub{})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/89/affinity-details", nil)
|
||||
router.ServeHTTP(rec, req)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var resp struct {
|
||||
Code int `json:"code"`
|
||||
Data struct {
|
||||
Users []any `json:"users"`
|
||||
TotalUsers int `json:"total_users"`
|
||||
TotalClients int `json:"total_clients"`
|
||||
PinnedUsers []int64 `json:"pinned_users"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, 0, resp.Code)
|
||||
require.Empty(t, resp.Data.Users)
|
||||
require.Equal(t, 0, resp.Data.TotalUsers)
|
||||
require.Equal(t, 0, resp.Data.TotalClients)
|
||||
require.Empty(t, resp.Data.PinnedUsers)
|
||||
}
|
||||
@@ -17,7 +17,7 @@ func setupAdminRouter() (*gin.Engine, *stubAdminService) {
|
||||
adminSvc := newStubAdminService()
|
||||
|
||||
userHandler := NewUserHandler(adminSvc, nil)
|
||||
groupHandler := NewGroupHandler(adminSvc)
|
||||
groupHandler := NewGroupHandler(adminSvc, nil, nil)
|
||||
proxyHandler := NewProxyHandler(adminSvc)
|
||||
redeemHandler := NewRedeemHandler(adminSvc, nil)
|
||||
|
||||
|
||||
@@ -445,5 +445,9 @@ func (s *stubAdminService) EnsureOpenAIPrivacy(ctx context.Context, account *ser
|
||||
return ""
|
||||
}
|
||||
|
||||
func (s *stubAdminService) ReplaceUserGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (*service.ReplaceUserGroupResult, error) {
|
||||
return &service.ReplaceUserGroupResult{MigratedKeys: 0}, nil
|
||||
}
|
||||
|
||||
// Ensure stub implements interface.
|
||||
var _ service.AdminService = (*stubAdminService)(nil)
|
||||
|
||||
@@ -98,12 +98,12 @@ func (h *BackupHandler) CreateBackup(c *gin.Context) {
|
||||
expireDays = *req.ExpireDays
|
||||
}
|
||||
|
||||
record, err := h.backupService.CreateBackup(c.Request.Context(), "manual", expireDays)
|
||||
record, err := h.backupService.StartBackup(c.Request.Context(), "manual", expireDays)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, record)
|
||||
response.Accepted(c, record)
|
||||
}
|
||||
|
||||
func (h *BackupHandler) ListBackups(c *gin.Context) {
|
||||
@@ -196,9 +196,10 @@ func (h *BackupHandler) RestoreBackup(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.backupService.RestoreBackup(c.Request.Context(), backupID); err != nil {
|
||||
record, err := h.backupService.StartRestore(c.Request.Context(), backupID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, gin.H{"restored": true})
|
||||
response.Accepted(c, record)
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -272,6 +273,7 @@ func (h *DashboardHandler) GetModelStats(c *gin.Context) {
|
||||
|
||||
// Parse optional filter params
|
||||
var userID, apiKeyID, accountID, groupID int64
|
||||
modelSource := usagestats.ModelSourceRequested
|
||||
var requestType *int16
|
||||
var stream *bool
|
||||
var billingType *int8
|
||||
@@ -296,6 +298,13 @@ func (h *DashboardHandler) GetModelStats(c *gin.Context) {
|
||||
groupID = id
|
||||
}
|
||||
}
|
||||
if rawModelSource := strings.TrimSpace(c.Query("model_source")); rawModelSource != "" {
|
||||
if !usagestats.IsValidModelSource(rawModelSource) {
|
||||
response.BadRequest(c, "Invalid model_source, use requested/upstream/mapping")
|
||||
return
|
||||
}
|
||||
modelSource = rawModelSource
|
||||
}
|
||||
if requestTypeStr := strings.TrimSpace(c.Query("request_type")); requestTypeStr != "" {
|
||||
parsed, err := service.ParseUsageRequestType(requestTypeStr)
|
||||
if err != nil {
|
||||
@@ -322,7 +331,7 @@ func (h *DashboardHandler) GetModelStats(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
stats, hit, err := h.getModelStatsCached(c.Request.Context(), startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType)
|
||||
stats, hit, err := h.getModelStatsCached(c.Request.Context(), startTime, endTime, userID, apiKeyID, accountID, groupID, modelSource, requestType, stream, billingType)
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get model statistics")
|
||||
return
|
||||
@@ -604,3 +613,47 @@ func (h *DashboardHandler) GetBatchAPIKeysUsage(c *gin.Context) {
|
||||
c.Header("X-Snapshot-Cache", "miss")
|
||||
response.Success(c, payload)
|
||||
}
|
||||
|
||||
// GetUserBreakdown handles getting per-user usage breakdown within a dimension.
|
||||
// GET /api/v1/admin/dashboard/user-breakdown
|
||||
// Query params: start_date, end_date, group_id, model, endpoint, endpoint_type, limit
|
||||
func (h *DashboardHandler) GetUserBreakdown(c *gin.Context) {
|
||||
startTime, endTime := parseTimeRange(c)
|
||||
|
||||
dim := usagestats.UserBreakdownDimension{}
|
||||
if v := c.Query("group_id"); v != "" {
|
||||
if id, err := strconv.ParseInt(v, 10, 64); err == nil {
|
||||
dim.GroupID = id
|
||||
}
|
||||
}
|
||||
dim.Model = c.Query("model")
|
||||
rawModelSource := strings.TrimSpace(c.DefaultQuery("model_source", usagestats.ModelSourceRequested))
|
||||
if !usagestats.IsValidModelSource(rawModelSource) {
|
||||
response.BadRequest(c, "Invalid model_source, use requested/upstream/mapping")
|
||||
return
|
||||
}
|
||||
dim.ModelType = rawModelSource
|
||||
dim.Endpoint = c.Query("endpoint")
|
||||
dim.EndpointType = c.DefaultQuery("endpoint_type", "inbound")
|
||||
|
||||
limit := 50
|
||||
if v := c.Query("limit"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 200 {
|
||||
limit = n
|
||||
}
|
||||
}
|
||||
|
||||
stats, err := h.dashboardService.GetUserBreakdownStats(
|
||||
c.Request.Context(), startTime, endTime, dim, limit,
|
||||
)
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get user breakdown stats")
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, gin.H{
|
||||
"users": stats,
|
||||
"start_date": startTime.Format("2006-01-02"),
|
||||
"end_date": endTime.Add(-24 * time.Hour).Format("2006-01-02"),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -149,6 +149,28 @@ func TestDashboardModelStatsInvalidStream(t *testing.T) {
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
}
|
||||
|
||||
func TestDashboardModelStatsInvalidModelSource(t *testing.T) {
|
||||
repo := &dashboardUsageRepoCapture{}
|
||||
router := newDashboardRequestTypeTestRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/admin/dashboard/models?model_source=invalid", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
}
|
||||
|
||||
func TestDashboardModelStatsValidModelSource(t *testing.T) {
|
||||
repo := &dashboardUsageRepoCapture{}
|
||||
router := newDashboardRequestTypeTestRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/admin/dashboard/models?model_source=upstream", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
|
||||
func TestDashboardUsersRankingLimitAndCache(t *testing.T) {
|
||||
dashboardUsersRankingCache = newSnapshotCache(5 * time.Minute)
|
||||
repo := &dashboardUsageRepoCapture{
|
||||
|
||||
@@ -0,0 +1,229 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// --- mock repo ---
|
||||
|
||||
type userBreakdownRepoCapture struct {
|
||||
service.UsageLogRepository
|
||||
capturedDim usagestats.UserBreakdownDimension
|
||||
capturedLimit int
|
||||
result []usagestats.UserBreakdownItem
|
||||
}
|
||||
|
||||
func (r *userBreakdownRepoCapture) GetUserBreakdownStats(
|
||||
_ context.Context, _, _ time.Time,
|
||||
dim usagestats.UserBreakdownDimension, limit int,
|
||||
) ([]usagestats.UserBreakdownItem, error) {
|
||||
r.capturedDim = dim
|
||||
r.capturedLimit = limit
|
||||
if r.result != nil {
|
||||
return r.result, nil
|
||||
}
|
||||
return []usagestats.UserBreakdownItem{}, nil
|
||||
}
|
||||
|
||||
func newUserBreakdownRouter(repo *userBreakdownRepoCapture) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := service.NewDashboardService(repo, nil, nil, nil)
|
||||
h := NewDashboardHandler(svc, nil)
|
||||
router := gin.New()
|
||||
router.GET("/admin/dashboard/user-breakdown", h.GetUserBreakdown)
|
||||
return router
|
||||
}
|
||||
|
||||
// --- tests ---
|
||||
|
||||
func TestGetUserBreakdown_GroupIDFilter(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&group_id=42", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, int64(42), repo.capturedDim.GroupID)
|
||||
require.Empty(t, repo.capturedDim.Model)
|
||||
require.Empty(t, repo.capturedDim.Endpoint)
|
||||
require.Equal(t, 50, repo.capturedLimit) // default limit
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_ModelFilter(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&model=claude-opus-4-6", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, "claude-opus-4-6", repo.capturedDim.Model)
|
||||
require.Equal(t, usagestats.ModelSourceRequested, repo.capturedDim.ModelType)
|
||||
require.Equal(t, int64(0), repo.capturedDim.GroupID)
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_ModelSourceFilter(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&model=claude-opus-4-6&model_source=upstream", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, usagestats.ModelSourceUpstream, repo.capturedDim.ModelType)
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_InvalidModelSource(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&model_source=foobar", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, w.Code)
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_EndpointFilter(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&endpoint=/v1/messages&endpoint_type=upstream", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, "/v1/messages", repo.capturedDim.Endpoint)
|
||||
require.Equal(t, "upstream", repo.capturedDim.EndpointType)
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_DefaultEndpointType(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&endpoint=/chat", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, "inbound", repo.capturedDim.EndpointType)
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_CustomLimit(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&model=test&limit=100", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, 100, repo.capturedLimit)
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_LimitClamped(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
// limit > 200 should fall back to default 50
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&model=test&limit=999", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, 50, repo.capturedLimit)
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_ResponseFormat(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{
|
||||
result: []usagestats.UserBreakdownItem{
|
||||
{UserID: 1, Email: "alice@test.com", Requests: 100, TotalTokens: 50000, Cost: 1.5, ActualCost: 1.2},
|
||||
{UserID: 2, Email: "bob@test.com", Requests: 50, TotalTokens: 25000, Cost: 0.8, ActualCost: 0.6},
|
||||
},
|
||||
}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&group_id=1", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
|
||||
var resp struct {
|
||||
Code int `json:"code"`
|
||||
Data struct {
|
||||
Users []usagestats.UserBreakdownItem `json:"users"`
|
||||
StartDate string `json:"start_date"`
|
||||
EndDate string `json:"end_date"`
|
||||
} `json:"data"`
|
||||
}
|
||||
err := json.Unmarshal(w.Body.Bytes(), &resp)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, resp.Code)
|
||||
require.Len(t, resp.Data.Users, 2)
|
||||
require.Equal(t, int64(1), resp.Data.Users[0].UserID)
|
||||
require.Equal(t, "alice@test.com", resp.Data.Users[0].Email)
|
||||
require.Equal(t, int64(100), resp.Data.Users[0].Requests)
|
||||
require.InDelta(t, 1.2, resp.Data.Users[0].ActualCost, 0.001)
|
||||
require.Equal(t, "2026-03-01", resp.Data.StartDate)
|
||||
require.Equal(t, "2026-03-16", resp.Data.EndDate)
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_EmptyResult(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&group_id=999", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
|
||||
var resp struct {
|
||||
Data struct {
|
||||
Users []usagestats.UserBreakdownItem `json:"users"`
|
||||
} `json:"data"`
|
||||
}
|
||||
err := json.Unmarshal(w.Body.Bytes(), &resp)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, resp.Data.Users)
|
||||
}
|
||||
|
||||
func TestGetUserBreakdown_NoFilters(t *testing.T) {
|
||||
repo := &userBreakdownRepoCapture{}
|
||||
router := newUserBreakdownRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, int64(0), repo.capturedDim.GroupID)
|
||||
require.Empty(t, repo.capturedDim.Model)
|
||||
require.Empty(t, repo.capturedDim.Endpoint)
|
||||
}
|
||||
@@ -38,6 +38,7 @@ type dashboardModelGroupCacheKey struct {
|
||||
APIKeyID int64 `json:"api_key_id"`
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
ModelSource string `json:"model_source,omitempty"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
Stream *bool `json:"stream"`
|
||||
BillingType *int8 `json:"billing_type"`
|
||||
@@ -111,6 +112,7 @@ func (h *DashboardHandler) getModelStatsCached(
|
||||
ctx context.Context,
|
||||
startTime, endTime time.Time,
|
||||
userID, apiKeyID, accountID, groupID int64,
|
||||
modelSource string,
|
||||
requestType *int16,
|
||||
stream *bool,
|
||||
billingType *int8,
|
||||
@@ -122,12 +124,13 @@ func (h *DashboardHandler) getModelStatsCached(
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
ModelSource: usagestats.NormalizeModelSource(modelSource),
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
})
|
||||
entry, hit, err := dashboardModelStatsCache.GetOrLoad(key, func() (any, error) {
|
||||
return h.dashboardService.GetModelStatsWithFilters(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType)
|
||||
return h.dashboardService.GetModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType, modelSource)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, hit, err
|
||||
|
||||
@@ -200,6 +200,7 @@ func (h *DashboardHandler) buildSnapshotV2Response(
|
||||
filters.APIKeyID,
|
||||
filters.AccountID,
|
||||
filters.GroupID,
|
||||
usagestats.ModelSourceRequested,
|
||||
filters.RequestType,
|
||||
filters.Stream,
|
||||
filters.BillingType,
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -16,7 +17,9 @@ import (
|
||||
|
||||
// GroupHandler handles admin group management
|
||||
type GroupHandler struct {
|
||||
adminService service.AdminService
|
||||
adminService service.AdminService
|
||||
dashboardService *service.DashboardService
|
||||
groupCapacityService *service.GroupCapacityService
|
||||
}
|
||||
|
||||
type optionalLimitField struct {
|
||||
@@ -69,9 +72,11 @@ func (f optionalLimitField) ToServiceInput() *float64 {
|
||||
}
|
||||
|
||||
// NewGroupHandler creates a new admin group handler
|
||||
func NewGroupHandler(adminService service.AdminService) *GroupHandler {
|
||||
func NewGroupHandler(adminService service.AdminService, dashboardService *service.DashboardService, groupCapacityService *service.GroupCapacityService) *GroupHandler {
|
||||
return &GroupHandler{
|
||||
adminService: adminService,
|
||||
adminService: adminService,
|
||||
dashboardService: dashboardService,
|
||||
groupCapacityService: groupCapacityService,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -367,6 +372,33 @@ func (h *GroupHandler) GetStats(c *gin.Context) {
|
||||
_ = groupID // TODO: implement actual stats
|
||||
}
|
||||
|
||||
// GetUsageSummary returns today's and cumulative cost for all groups.
|
||||
// GET /api/v1/admin/groups/usage-summary?timezone=Asia/Shanghai
|
||||
func (h *GroupHandler) GetUsageSummary(c *gin.Context) {
|
||||
userTZ := c.Query("timezone")
|
||||
now := timezone.NowInUserLocation(userTZ)
|
||||
todayStart := timezone.StartOfDayInUserLocation(now, userTZ)
|
||||
|
||||
results, err := h.dashboardService.GetGroupUsageSummary(c.Request.Context(), todayStart)
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get group usage summary")
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, results)
|
||||
}
|
||||
|
||||
// GetCapacitySummary returns aggregated capacity (concurrency/sessions/RPM) for all active groups.
|
||||
// GET /api/v1/admin/groups/capacity-summary
|
||||
func (h *GroupHandler) GetCapacitySummary(c *gin.Context) {
|
||||
results, err := h.groupCapacityService.GetAllGroupCapacity(c.Request.Context())
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get group capacity summary")
|
||||
return
|
||||
}
|
||||
response.Success(c, results)
|
||||
}
|
||||
|
||||
// GetGroupAPIKeys handles getting API keys in a group
|
||||
// GET /api/v1/admin/groups/:id/api-keys
|
||||
func (h *GroupHandler) GetGroupAPIKeys(c *gin.Context) {
|
||||
|
||||
@@ -129,6 +129,7 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
OpsQueryModeDefault: settings.OpsQueryModeDefault,
|
||||
OpsMetricsIntervalSeconds: settings.OpsMetricsIntervalSeconds,
|
||||
MinClaudeCodeVersion: settings.MinClaudeCodeVersion,
|
||||
MaxClaudeCodeVersion: settings.MaxClaudeCodeVersion,
|
||||
AllowUngroupedKeyScheduling: settings.AllowUngroupedKeyScheduling,
|
||||
BackendModeEnabled: settings.BackendModeEnabled,
|
||||
})
|
||||
@@ -203,6 +204,7 @@ type UpdateSettingsRequest struct {
|
||||
OpsMetricsIntervalSeconds *int `json:"ops_metrics_interval_seconds"`
|
||||
|
||||
MinClaudeCodeVersion string `json:"min_claude_code_version"`
|
||||
MaxClaudeCodeVersion string `json:"max_claude_code_version"`
|
||||
|
||||
// 分组隔离
|
||||
AllowUngroupedKeyScheduling bool `json:"allow_ungrouped_key_scheduling"`
|
||||
@@ -446,6 +448,22 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
// 验证最高版本号格式(空字符串=禁用,或合法 semver)
|
||||
if req.MaxClaudeCodeVersion != "" {
|
||||
if !semverPattern.MatchString(req.MaxClaudeCodeVersion) {
|
||||
response.Error(c, http.StatusBadRequest, "max_claude_code_version must be empty or a valid semver (e.g. 3.0.0)")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 交叉验证:如果同时设置了最低和最高版本号,最高版本号必须 >= 最低版本号
|
||||
if req.MinClaudeCodeVersion != "" && req.MaxClaudeCodeVersion != "" {
|
||||
if service.CompareVersions(req.MaxClaudeCodeVersion, req.MinClaudeCodeVersion) < 0 {
|
||||
response.Error(c, http.StatusBadRequest, "max_claude_code_version must be greater than or equal to min_claude_code_version")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
settings := &service.SystemSettings{
|
||||
RegistrationEnabled: req.RegistrationEnabled,
|
||||
EmailVerifyEnabled: req.EmailVerifyEnabled,
|
||||
@@ -492,6 +510,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
EnableIdentityPatch: req.EnableIdentityPatch,
|
||||
IdentityPatchPrompt: req.IdentityPatchPrompt,
|
||||
MinClaudeCodeVersion: req.MinClaudeCodeVersion,
|
||||
MaxClaudeCodeVersion: req.MaxClaudeCodeVersion,
|
||||
AllowUngroupedKeyScheduling: req.AllowUngroupedKeyScheduling,
|
||||
BackendModeEnabled: req.BackendModeEnabled,
|
||||
OpsMonitoringEnabled: func() bool {
|
||||
@@ -592,6 +611,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
OpsQueryModeDefault: updatedSettings.OpsQueryModeDefault,
|
||||
OpsMetricsIntervalSeconds: updatedSettings.OpsMetricsIntervalSeconds,
|
||||
MinClaudeCodeVersion: updatedSettings.MinClaudeCodeVersion,
|
||||
MaxClaudeCodeVersion: updatedSettings.MaxClaudeCodeVersion,
|
||||
AllowUngroupedKeyScheduling: updatedSettings.AllowUngroupedKeyScheduling,
|
||||
BackendModeEnabled: updatedSettings.BackendModeEnabled,
|
||||
})
|
||||
@@ -748,6 +768,9 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.MinClaudeCodeVersion != after.MinClaudeCodeVersion {
|
||||
changed = append(changed, "min_claude_code_version")
|
||||
}
|
||||
if before.MaxClaudeCodeVersion != after.MaxClaudeCodeVersion {
|
||||
changed = append(changed, "max_claude_code_version")
|
||||
}
|
||||
if before.AllowUngroupedKeyScheduling != after.AllowUngroupedKeyScheduling {
|
||||
changed = append(changed, "allow_ungrouped_key_scheduling")
|
||||
}
|
||||
@@ -981,6 +1004,58 @@ func (h *SettingHandler) DeleteAdminAPIKey(c *gin.Context) {
|
||||
response.Success(c, gin.H{"message": "Admin API key deleted"})
|
||||
}
|
||||
|
||||
// GetOverloadCooldownSettings 获取529过载冷却配置
|
||||
// GET /api/v1/admin/settings/overload-cooldown
|
||||
func (h *SettingHandler) GetOverloadCooldownSettings(c *gin.Context) {
|
||||
settings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context())
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, dto.OverloadCooldownSettings{
|
||||
Enabled: settings.Enabled,
|
||||
CooldownMinutes: settings.CooldownMinutes,
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateOverloadCooldownSettingsRequest 更新529过载冷却配置请求
|
||||
type UpdateOverloadCooldownSettingsRequest struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
CooldownMinutes int `json:"cooldown_minutes"`
|
||||
}
|
||||
|
||||
// UpdateOverloadCooldownSettings 更新529过载冷却配置
|
||||
// PUT /api/v1/admin/settings/overload-cooldown
|
||||
func (h *SettingHandler) UpdateOverloadCooldownSettings(c *gin.Context) {
|
||||
var req UpdateOverloadCooldownSettingsRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
settings := &service.OverloadCooldownSettings{
|
||||
Enabled: req.Enabled,
|
||||
CooldownMinutes: req.CooldownMinutes,
|
||||
}
|
||||
|
||||
if err := h.settingService.SetOverloadCooldownSettings(c.Request.Context(), settings); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
updatedSettings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context())
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, dto.OverloadCooldownSettings{
|
||||
Enabled: updatedSettings.Enabled,
|
||||
CooldownMinutes: updatedSettings.CooldownMinutes,
|
||||
})
|
||||
}
|
||||
|
||||
// GetStreamTimeoutSettings 获取流超时处理配置
|
||||
// GET /api/v1/admin/settings/stream-timeout
|
||||
func (h *SettingHandler) GetStreamTimeoutSettings(c *gin.Context) {
|
||||
|
||||
@@ -77,12 +77,13 @@ func (h *SubscriptionHandler) List(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
status := c.Query("status")
|
||||
platform := c.Query("platform")
|
||||
|
||||
// Parse sorting parameters
|
||||
sortBy := c.DefaultQuery("sort_by", "created_at")
|
||||
sortOrder := c.DefaultQuery("sort_order", "desc")
|
||||
|
||||
subscriptions, pagination, err := h.subscriptionService.List(c.Request.Context(), page, pageSize, userID, groupID, status, sortBy, sortOrder)
|
||||
subscriptions, pagination, err := h.subscriptionService.List(c.Request.Context(), page, pageSize, userID, groupID, status, platform, sortBy, sortOrder)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
|
||||
@@ -75,6 +75,7 @@ type UpdateBalanceRequest struct {
|
||||
// - role: filter by user role
|
||||
// - search: search in email, username
|
||||
// - attr[{id}]: filter by custom attribute value, e.g. attr[1]=company
|
||||
// - group_name: fuzzy filter by allowed group name
|
||||
func (h *UserHandler) List(c *gin.Context) {
|
||||
page, pageSize := response.ParsePagination(c)
|
||||
|
||||
@@ -89,6 +90,7 @@ func (h *UserHandler) List(c *gin.Context) {
|
||||
Status: c.Query("status"),
|
||||
Role: c.Query("role"),
|
||||
Search: search,
|
||||
GroupName: strings.TrimSpace(c.Query("group_name")),
|
||||
Attributes: parseAttributeFilters(c),
|
||||
}
|
||||
if raw, ok := c.GetQuery("include_subscriptions"); ok {
|
||||
@@ -366,3 +368,35 @@ func (h *UserHandler) GetBalanceHistory(c *gin.Context) {
|
||||
"total_recharged": totalRecharged,
|
||||
})
|
||||
}
|
||||
|
||||
// ReplaceGroupRequest represents the request to replace a user's exclusive group
|
||||
type ReplaceGroupRequest struct {
|
||||
OldGroupID int64 `json:"old_group_id" binding:"required,gt=0"`
|
||||
NewGroupID int64 `json:"new_group_id" binding:"required,gt=0"`
|
||||
}
|
||||
|
||||
// ReplaceGroup handles replacing a user's exclusive group
|
||||
// POST /api/v1/admin/users/:id/replace-group
|
||||
func (h *UserHandler) ReplaceGroup(c *gin.Context) {
|
||||
userID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid user ID")
|
||||
return
|
||||
}
|
||||
|
||||
var req ReplaceGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
result, err := h.adminService.ReplaceUserGroup(c.Request.Context(), userID, req.OldGroupID, req.NewGroupID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, gin.H{
|
||||
"migrated_keys": result.MigratedKeys,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -143,6 +143,8 @@ func GroupFromServiceAdmin(g *service.Group) *AdminGroup {
|
||||
SimulateClaudeMaxEnabled: g.SimulateClaudeMaxEnabled,
|
||||
SupportedModelScopes: g.SupportedModelScopes,
|
||||
AccountCount: g.AccountCount,
|
||||
ActiveAccountCount: g.ActiveAccountCount,
|
||||
RateLimitedAccountCount: g.RateLimitedAccountCount,
|
||||
SortOrder: g.SortOrder,
|
||||
}
|
||||
if len(g.AccountGroups) > 0 {
|
||||
@@ -266,9 +268,14 @@ func AccountFromServiceShallow(a *service.Account) *Account {
|
||||
}
|
||||
|
||||
// 客户端亲和调度(Anthropic 和 Antigravity 账号)
|
||||
if a.IsClientAffinityEnabled() {
|
||||
if a.IsAffinityEnabled() {
|
||||
enabled := true
|
||||
out.ClientAffinityEnabled = &enabled
|
||||
allow := a.IsAffinityAllowSwitch()
|
||||
out.AffinityAllowSwitch = &allow
|
||||
if pinnedUsers := a.GetPinnedUsers(); len(pinnedUsers) > 0 {
|
||||
out.PinnedUserIDs = pinnedUsers
|
||||
}
|
||||
}
|
||||
|
||||
// 提取账号配额限制(apikey / bedrock 类型有效)
|
||||
@@ -528,6 +535,7 @@ func usageLogFromServiceUser(l *service.UsageLog) UsageLog {
|
||||
AccountID: l.AccountID,
|
||||
RequestID: l.RequestID,
|
||||
Model: l.Model,
|
||||
UpstreamModel: l.UpstreamModel,
|
||||
ServiceTier: l.ServiceTier,
|
||||
ReasoningEffort: l.ReasoningEffort,
|
||||
InboundEndpoint: l.InboundEndpoint,
|
||||
|
||||
@@ -79,6 +79,7 @@ type SystemSettings struct {
|
||||
OpsMetricsIntervalSeconds int `json:"ops_metrics_interval_seconds"`
|
||||
|
||||
MinClaudeCodeVersion string `json:"min_claude_code_version"`
|
||||
MaxClaudeCodeVersion string `json:"max_claude_code_version"`
|
||||
|
||||
// 分组隔离
|
||||
AllowUngroupedKeyScheduling bool `json:"allow_ungrouped_key_scheduling"`
|
||||
@@ -167,6 +168,12 @@ type ListSoraS3ProfilesResponse struct {
|
||||
Items []SoraS3Profile `json:"items"`
|
||||
}
|
||||
|
||||
// OverloadCooldownSettings 529过载冷却配置 DTO
|
||||
type OverloadCooldownSettings struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
CooldownMinutes int `json:"cooldown_minutes"`
|
||||
}
|
||||
|
||||
// StreamTimeoutSettings 流超时处理配置 DTO
|
||||
type StreamTimeoutSettings struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
|
||||
@@ -124,9 +124,11 @@ type AdminGroup struct {
|
||||
DefaultMappedModel string `json:"default_mapped_model"`
|
||||
|
||||
// 支持的模型系列(仅 antigravity 平台使用)
|
||||
SupportedModelScopes []string `json:"supported_model_scopes"`
|
||||
AccountGroups []AccountGroup `json:"account_groups,omitempty"`
|
||||
AccountCount int64 `json:"account_count,omitempty"`
|
||||
SupportedModelScopes []string `json:"supported_model_scopes"`
|
||||
AccountGroups []AccountGroup `json:"account_groups,omitempty"`
|
||||
AccountCount int64 `json:"account_count,omitempty"`
|
||||
ActiveAccountCount int64 `json:"active_account_count,omitempty"`
|
||||
RateLimitedAccountCount int64 `json:"rate_limited_account_count,omitempty"`
|
||||
|
||||
// 分组排序
|
||||
SortOrder int `json:"sort_order"`
|
||||
@@ -201,6 +203,15 @@ type Account struct {
|
||||
// 启用后新会话会优先调度到客户端之前使用过的账号
|
||||
ClientAffinityEnabled *bool `json:"client_affinity_enabled,omitempty"`
|
||||
|
||||
// 亲和允许切换(默认 true)
|
||||
AffinityAllowSwitch *bool `json:"affinity_allow_switch,omitempty"`
|
||||
|
||||
// 亲和用户数量(admin 列表端点注入)
|
||||
AffinityUserCount *int64 `json:"affinity_user_count,omitempty"`
|
||||
|
||||
// 指定亲和用户 ID 列表
|
||||
PinnedUserIDs []int64 `json:"pinned_user_ids,omitempty"`
|
||||
|
||||
// 亲和客户端数据(仅 admin 列表端点注入,不由 mapper 填充)
|
||||
AffinityClientCount *int64 `json:"affinity_client_count,omitempty"`
|
||||
AffinityClients []string `json:"affinity_clients,omitempty"`
|
||||
@@ -342,6 +353,9 @@ type UsageLog struct {
|
||||
AccountID int64 `json:"account_id"`
|
||||
RequestID string `json:"request_id"`
|
||||
Model string `json:"model"`
|
||||
// UpstreamModel is the actual model sent to the upstream provider after mapping.
|
||||
// Omitted when no mapping was applied (requested model was used as-is).
|
||||
UpstreamModel *string `json:"upstream_model,omitempty"`
|
||||
// ServiceTier records the OpenAI service tier used for billing, e.g. "priority" / "flex".
|
||||
ServiceTier *string `json:"service_tier,omitempty"`
|
||||
// ReasoningEffort is the request's reasoning effort level.
|
||||
|
||||
@@ -291,7 +291,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
}
|
||||
|
||||
for {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, "") // Gemini 不使用会话限制
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
if err != nil {
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
||||
@@ -501,8 +501,16 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
|
||||
for {
|
||||
// 选择支持该模型的账号
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), currentAPIKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, parsedReq.MetadataUserID)
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), currentAPIKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, parsedReq.MetadataUserID, subject.UserID)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrAffinityNoSwitch) {
|
||||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "Affinity account unavailable and switching is disabled", streamStarted)
|
||||
return
|
||||
}
|
||||
if errors.Is(err, service.ErrAffinityLimitExceeded) {
|
||||
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "api_error", "Affinity client limit exceeded", streamStarted)
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
||||
return
|
||||
@@ -1222,6 +1230,10 @@ func (h *GatewayHandler) handleFailoverExhausted(c *gin.Context, failoverErr *se
|
||||
}
|
||||
}
|
||||
|
||||
// 记录原始上游状态码,以便 ops 错误日志捕获真实的上游错误
|
||||
upstreamMsg := service.ExtractUpstreamErrorMessage(responseBody)
|
||||
service.SetOpsUpstreamError(c, statusCode, upstreamMsg, "")
|
||||
|
||||
// 使用默认的错误映射
|
||||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||
@@ -1230,6 +1242,7 @@ func (h *GatewayHandler) handleFailoverExhausted(c *gin.Context, failoverErr *se
|
||||
// handleFailoverExhaustedSimple 简化版本,用于没有响应体的情况
|
||||
func (h *GatewayHandler) handleFailoverExhaustedSimple(c *gin.Context, statusCode int, streamStarted bool) {
|
||||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||||
service.SetOpsUpstreamError(c, statusCode, errMsg, "")
|
||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||
}
|
||||
|
||||
@@ -1279,7 +1292,7 @@ func (h *GatewayHandler) ensureForwardErrorResponse(c *gin.Context, streamStarte
|
||||
return true
|
||||
}
|
||||
|
||||
// checkClaudeCodeVersion 检查 Claude Code 客户端版本是否满足最低要求
|
||||
// checkClaudeCodeVersion 检查 Claude Code 客户端版本是否满足版本要求
|
||||
// 仅对已识别的 Claude Code 客户端执行,count_tokens 路径除外
|
||||
func (h *GatewayHandler) checkClaudeCodeVersion(c *gin.Context) bool {
|
||||
ctx := c.Request.Context()
|
||||
@@ -1292,8 +1305,8 @@ func (h *GatewayHandler) checkClaudeCodeVersion(c *gin.Context) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
minVersion := h.settingService.GetMinClaudeCodeVersion(ctx)
|
||||
if minVersion == "" {
|
||||
minVersion, maxVersion := h.settingService.GetClaudeCodeVersionBounds(ctx)
|
||||
if minVersion == "" && maxVersion == "" {
|
||||
return true // 未设置,不检查
|
||||
}
|
||||
|
||||
@@ -1304,13 +1317,22 @@ func (h *GatewayHandler) checkClaudeCodeVersion(c *gin.Context) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
if service.CompareVersions(clientVersion, minVersion) < 0 {
|
||||
if minVersion != "" && service.CompareVersions(clientVersion, minVersion) < 0 {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
||||
fmt.Sprintf("Your Claude Code version (%s) is below the minimum required version (%s). Please update: npm update -g @anthropic-ai/claude-code",
|
||||
clientVersion, minVersion))
|
||||
return false
|
||||
}
|
||||
|
||||
if maxVersion != "" && service.CompareVersions(clientVersion, maxVersion) > 0 {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
||||
fmt.Sprintf("Your Claude Code version (%s) exceeds the maximum allowed version (%s). "+
|
||||
"Please downgrade: npm install -g @anthropic-ai/claude-code@%s && "+
|
||||
"set CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 to prevent auto-upgrade",
|
||||
clientVersion, maxVersion, maxVersion))
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
|
||||
@@ -76,7 +76,7 @@ func (f *fakeGroupRepo) ListActiveByPlatform(context.Context, string) ([]service
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeGroupRepo) ExistsByName(context.Context, string) (bool, error) { return false, nil }
|
||||
func (f *fakeGroupRepo) GetAccountCount(context.Context, int64) (int64, error) { return 0, nil }
|
||||
func (f *fakeGroupRepo) GetAccountCount(context.Context, int64) (int64, int64, error) { return 0, 0, nil }
|
||||
func (f *fakeGroupRepo) DeleteAccountGroupsByGroupID(context.Context, int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
@@ -136,7 +136,7 @@ func validClaudeCodeBodyJSON() []byte {
|
||||
return []byte(`{
|
||||
"model":"claude-3-5-sonnet-20241022",
|
||||
"system":[{"text":"You are Claude Code, Anthropic's official CLI for Claude."}],
|
||||
"metadata":{"user_id":"user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_abc-123"}
|
||||
"metadata":{"user_id":"user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"}
|
||||
}`)
|
||||
}
|
||||
|
||||
@@ -190,7 +190,7 @@ func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing
|
||||
System: []any{
|
||||
map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."},
|
||||
},
|
||||
MetadataUserID: "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_abc-123",
|
||||
MetadataUserID: "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
|
||||
}
|
||||
|
||||
// body 非法 JSON,如果函数复用 parsedReq 成功则仍应判定为 Claude Code。
|
||||
@@ -209,7 +209,7 @@ func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing
|
||||
"system": []any{
|
||||
map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."},
|
||||
},
|
||||
"metadata": map[string]any{"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_abc-123"},
|
||||
"metadata": map[string]any{"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"},
|
||||
})
|
||||
|
||||
SetClaudeCodeClientContext(c, []byte(`{invalid`), nil)
|
||||
|
||||
@@ -352,7 +352,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
}
|
||||
|
||||
for {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "") // Gemini 不使用会话限制
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
if err != nil {
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
|
||||
@@ -593,6 +593,10 @@ func (h *GatewayHandler) handleGeminiFailoverExhausted(c *gin.Context, failoverE
|
||||
}
|
||||
}
|
||||
|
||||
// 记录原始上游状态码,以便 ops 错误日志捕获真实的上游错误
|
||||
upstreamMsg := service.ExtractUpstreamErrorMessage(responseBody)
|
||||
service.SetOpsUpstreamError(c, statusCode, upstreamMsg, "")
|
||||
|
||||
// 使用默认的错误映射
|
||||
status, message := mapGeminiUpstreamError(statusCode)
|
||||
googleError(c, status, message)
|
||||
|
||||
@@ -1435,6 +1435,10 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE
|
||||
}
|
||||
}
|
||||
|
||||
// 记录原始上游状态码,以便 ops 错误日志捕获真实的上游错误
|
||||
upstreamMsg := service.ExtractUpstreamErrorMessage(responseBody)
|
||||
service.SetOpsUpstreamError(c, statusCode, upstreamMsg, "")
|
||||
|
||||
// 使用默认的错误映射
|
||||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||
@@ -1443,6 +1447,7 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE
|
||||
// handleFailoverExhaustedSimple 简化版本,用于没有响应体的情况
|
||||
func (h *OpenAIGatewayHandler) handleFailoverExhaustedSimple(c *gin.Context, statusCode int, streamStarted bool) {
|
||||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||||
service.SetOpsUpstreamError(c, statusCode, errMsg, "")
|
||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||
}
|
||||
|
||||
|
||||
@@ -945,6 +945,9 @@ func (r *stubUserRepoForHandler) ExistsByEmail(context.Context, string) (bool, e
|
||||
func (r *stubUserRepoForHandler) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (r *stubUserRepoForHandler) RemoveGroupFromUserAllowedGroups(context.Context, int64, int64) error {
|
||||
return nil
|
||||
}
|
||||
func (r *stubUserRepoForHandler) UpdateTotpSecret(context.Context, int64, *string) error { return nil }
|
||||
func (r *stubUserRepoForHandler) EnableTotp(context.Context, int64) error { return nil }
|
||||
func (r *stubUserRepoForHandler) DisableTotp(context.Context, int64) error { return nil }
|
||||
@@ -1020,6 +1023,20 @@ func (r *stubAPIKeyRepoForHandler) SearchAPIKeys(context.Context, int64, string,
|
||||
func (r *stubAPIKeyRepoForHandler) ClearGroupIDByGroupID(context.Context, int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (r *stubAPIKeyRepoForHandler) UpdateGroupIDByUserAndGroup(_ context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
|
||||
var updated int64
|
||||
for id, key := range r.keys {
|
||||
if key.UserID != userID || key.GroupID == nil || *key.GroupID != oldGroupID {
|
||||
continue
|
||||
}
|
||||
clone := *key
|
||||
gid := newGroupID
|
||||
clone.GroupID = &gid
|
||||
r.keys[id] = &clone
|
||||
updated++
|
||||
}
|
||||
return updated, nil
|
||||
}
|
||||
func (r *stubAPIKeyRepoForHandler) CountByGroupID(context.Context, int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
@@ -225,7 +225,7 @@ func (h *SoraGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
var lastFailoverHeaders http.Header
|
||||
|
||||
for {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionHash, reqModel, failedAccountIDs, "")
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionHash, reqModel, failedAccountIDs, "", int64(0))
|
||||
if err != nil {
|
||||
reqLog.Warn("sora.account_select_failed",
|
||||
zap.Error(err),
|
||||
@@ -484,6 +484,9 @@ func (h *SoraGatewayHandler) handleConcurrencyError(c *gin.Context, err error, s
|
||||
}
|
||||
|
||||
func (h *SoraGatewayHandler) handleFailoverExhausted(c *gin.Context, statusCode int, responseHeaders http.Header, responseBody []byte, streamStarted bool) {
|
||||
upstreamMsg := service.ExtractUpstreamErrorMessage(responseBody)
|
||||
service.SetOpsUpstreamError(c, statusCode, upstreamMsg, "")
|
||||
|
||||
status, errType, errMsg := h.mapUpstreamError(statusCode, responseHeaders, responseBody)
|
||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||
}
|
||||
|
||||
@@ -273,8 +273,8 @@ func (r *stubGroupRepo) ListActiveByPlatform(ctx context.Context, platform strin
|
||||
func (r *stubGroupRepo) ExistsByName(ctx context.Context, name string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
func (r *stubGroupRepo) GetAccountCount(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, nil
|
||||
func (r *stubGroupRepo) GetAccountCount(ctx context.Context, groupID int64) (int64, int64, error) {
|
||||
return 0, 0, nil
|
||||
}
|
||||
func (r *stubGroupRepo) DeleteAccountGroupsByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, nil
|
||||
@@ -345,6 +345,12 @@ func (s *stubUsageLogRepo) GetUpstreamEndpointStatsWithFilters(ctx context.Conte
|
||||
func (s *stubUsageLogRepo) GetGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) ([]usagestats.GroupStat, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *stubUsageLogRepo) GetUserBreakdownStats(ctx context.Context, startTime, endTime time.Time, dim usagestats.UserBreakdownDimension, limit int) ([]usagestats.UserBreakdownItem, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *stubUsageLogRepo) GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *stubUsageLogRepo) GetAPIKeyUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.APIKeyUsageTrendPoint, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -323,7 +323,7 @@ func (h *SoraVideosHandler) selectAccount(c *gin.Context, model string) (*servic
|
||||
}
|
||||
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(
|
||||
c.Request.Context(), apiKey.GroupID, "", model, nil, "",
|
||||
c.Request.Context(), apiKey.GroupID, "", model, nil, "", int64(0),
|
||||
)
|
||||
if err != nil {
|
||||
soraErrorResponse(c, http.StatusServiceUnavailable, "server_error", "No available accounts")
|
||||
|
||||
@@ -228,9 +228,18 @@ type Client struct {
|
||||
httpClient *http.Client
|
||||
}
|
||||
|
||||
const (
|
||||
// proxyDialTimeout 代理 TCP 连接超时(含代理握手),代理不通时快速失败
|
||||
proxyDialTimeout = 5 * time.Second
|
||||
// proxyTLSHandshakeTimeout 代理 TLS 握手超时
|
||||
proxyTLSHandshakeTimeout = 5 * time.Second
|
||||
// clientTimeout 整体请求超时(含连接、发送、等待响应、读取 body)
|
||||
clientTimeout = 10 * time.Second
|
||||
)
|
||||
|
||||
func NewClient(proxyURL string) (*Client, error) {
|
||||
client := &http.Client{
|
||||
Timeout: 30 * time.Second,
|
||||
Timeout: clientTimeout,
|
||||
}
|
||||
|
||||
_, parsed, err := proxyurl.Parse(proxyURL)
|
||||
@@ -238,7 +247,12 @@ func NewClient(proxyURL string) (*Client, error) {
|
||||
return nil, err
|
||||
}
|
||||
if parsed != nil {
|
||||
transport := &http.Transport{}
|
||||
transport := &http.Transport{
|
||||
DialContext: (&net.Dialer{
|
||||
Timeout: proxyDialTimeout,
|
||||
}).DialContext,
|
||||
TLSHandshakeTimeout: proxyTLSHandshakeTimeout,
|
||||
}
|
||||
if err := proxyutil.ConfigureTransportProxy(transport, parsed); err != nil {
|
||||
return nil, fmt.Errorf("configure proxy: %w", err)
|
||||
}
|
||||
@@ -250,8 +264,8 @@ func NewClient(proxyURL string) (*Client, error) {
|
||||
}, nil
|
||||
}
|
||||
|
||||
// isConnectionError 判断是否为连接错误(网络超时、DNS 失败、连接拒绝)
|
||||
func isConnectionError(err error) bool {
|
||||
// IsConnectionError 判断是否为连接错误(网络超时、DNS 失败、连接拒绝)
|
||||
func IsConnectionError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
@@ -276,7 +290,7 @@ func isConnectionError(err error) bool {
|
||||
// shouldFallbackToNextURL 判断是否应切换到下一个 URL
|
||||
// 与 Antigravity-Manager 保持一致:连接错误、429、408、404、5xx 触发 URL 降级
|
||||
func shouldFallbackToNextURL(err error, statusCode int) bool {
|
||||
if isConnectionError(err) {
|
||||
if IsConnectionError(err) {
|
||||
return true
|
||||
}
|
||||
return statusCode == http.StatusTooManyRequests ||
|
||||
|
||||
@@ -274,8 +274,8 @@ func TestNewClient_无代理(t *testing.T) {
|
||||
if client.httpClient == nil {
|
||||
t.Fatal("httpClient 为 nil")
|
||||
}
|
||||
if client.httpClient.Timeout != 30*time.Second {
|
||||
t.Errorf("Timeout 不匹配: got %v, want 30s", client.httpClient.Timeout)
|
||||
if client.httpClient.Timeout != clientTimeout {
|
||||
t.Errorf("Timeout 不匹配: got %v, want %v", client.httpClient.Timeout, clientTimeout)
|
||||
}
|
||||
// 无代理时 Transport 应为 nil(使用默认)
|
||||
if client.httpClient.Transport != nil {
|
||||
@@ -322,11 +322,11 @@ func TestNewClient_无效代理URL(t *testing.T) {
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// isConnectionError
|
||||
// IsConnectionError
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestIsConnectionError_nil(t *testing.T) {
|
||||
if isConnectionError(nil) {
|
||||
if IsConnectionError(nil) {
|
||||
t.Error("nil 错误不应判定为连接错误")
|
||||
}
|
||||
}
|
||||
@@ -338,7 +338,7 @@ func TestIsConnectionError_超时错误(t *testing.T) {
|
||||
Net: "tcp",
|
||||
Err: &timeoutError{},
|
||||
}
|
||||
if !isConnectionError(err) {
|
||||
if !IsConnectionError(err) {
|
||||
t.Error("超时错误应判定为连接错误")
|
||||
}
|
||||
}
|
||||
@@ -356,7 +356,7 @@ func TestIsConnectionError_netOpError(t *testing.T) {
|
||||
Net: "tcp",
|
||||
Err: fmt.Errorf("connection refused"),
|
||||
}
|
||||
if !isConnectionError(err) {
|
||||
if !IsConnectionError(err) {
|
||||
t.Error("net.OpError 应判定为连接错误")
|
||||
}
|
||||
}
|
||||
@@ -367,14 +367,14 @@ func TestIsConnectionError_urlError(t *testing.T) {
|
||||
URL: "https://example.com",
|
||||
Err: fmt.Errorf("some error"),
|
||||
}
|
||||
if !isConnectionError(err) {
|
||||
if !IsConnectionError(err) {
|
||||
t.Error("url.Error 应判定为连接错误")
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsConnectionError_普通错误(t *testing.T) {
|
||||
err := fmt.Errorf("some random error")
|
||||
if isConnectionError(err) {
|
||||
if IsConnectionError(err) {
|
||||
t.Error("普通错误不应判定为连接错误")
|
||||
}
|
||||
}
|
||||
@@ -386,7 +386,7 @@ func TestIsConnectionError_包装的netOpError(t *testing.T) {
|
||||
Err: fmt.Errorf("connection refused"),
|
||||
}
|
||||
err := fmt.Errorf("wrapping: %w", inner)
|
||||
if !isConnectionError(err) {
|
||||
if !IsConnectionError(err) {
|
||||
t.Error("被包装的 net.OpError 应判定为连接错误")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -275,21 +275,6 @@ func filterOpenCodePrompt(text string) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// systemBlockFilterPrefixes 需要从 system 中过滤的文本前缀列表
|
||||
var systemBlockFilterPrefixes = []string{
|
||||
"x-anthropic-billing-header",
|
||||
}
|
||||
|
||||
// filterSystemBlockByPrefix 如果文本匹配过滤前缀,返回空字符串
|
||||
func filterSystemBlockByPrefix(text string) string {
|
||||
for _, prefix := range systemBlockFilterPrefixes {
|
||||
if strings.HasPrefix(text, prefix) {
|
||||
return ""
|
||||
}
|
||||
}
|
||||
return text
|
||||
}
|
||||
|
||||
// buildSystemInstruction 构建 systemInstruction(与 Antigravity-Manager 保持一致)
|
||||
func buildSystemInstruction(system json.RawMessage, modelName string, opts TransformOptions, tools []ClaudeTool) *GeminiContent {
|
||||
var parts []GeminiPart
|
||||
@@ -306,8 +291,8 @@ func buildSystemInstruction(system json.RawMessage, modelName string, opts Trans
|
||||
if strings.Contains(sysStr, "You are Antigravity") {
|
||||
userHasAntigravityIdentity = true
|
||||
}
|
||||
// 过滤 OpenCode 默认提示词和黑名单前缀
|
||||
filtered := filterSystemBlockByPrefix(filterOpenCodePrompt(sysStr))
|
||||
// 过滤 OpenCode 默认提示词
|
||||
filtered := filterOpenCodePrompt(sysStr)
|
||||
if filtered != "" {
|
||||
userSystemParts = append(userSystemParts, GeminiPart{Text: filtered})
|
||||
}
|
||||
@@ -321,8 +306,8 @@ func buildSystemInstruction(system json.RawMessage, modelName string, opts Trans
|
||||
if strings.Contains(block.Text, "You are Antigravity") {
|
||||
userHasAntigravityIdentity = true
|
||||
}
|
||||
// 过滤 OpenCode 默认提示词和黑名单前缀
|
||||
filtered := filterSystemBlockByPrefix(filterOpenCodePrompt(block.Text))
|
||||
// 过滤 OpenCode 默认提示词
|
||||
filtered := filterOpenCodePrompt(block.Text)
|
||||
if filtered != "" {
|
||||
userSystemParts = append(userSystemParts, GeminiPart{Text: filtered})
|
||||
}
|
||||
|
||||
@@ -2,7 +2,10 @@ package antigravity
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestBuildParts_ThinkingBlockWithoutSignature 测试thinking block无signature时的处理
|
||||
@@ -349,3 +352,51 @@ func TestBuildGenerationConfig_ThinkingDynamicBudget(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTransformClaudeToGeminiWithOptions_PreservesBillingHeaderSystemBlock(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
system json.RawMessage
|
||||
}{
|
||||
{
|
||||
name: "system array",
|
||||
system: json.RawMessage(`[{"type":"text","text":"x-anthropic-billing-header keep"}]`),
|
||||
},
|
||||
{
|
||||
name: "system string",
|
||||
system: json.RawMessage(`"x-anthropic-billing-header keep"`),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
claudeReq := &ClaudeRequest{
|
||||
Model: "claude-3-5-sonnet-latest",
|
||||
System: tt.system,
|
||||
Messages: []ClaudeMessage{
|
||||
{
|
||||
Role: "user",
|
||||
Content: json.RawMessage(`[{"type":"text","text":"hello"}]`),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
body, err := TransformClaudeToGeminiWithOptions(claudeReq, "project-1", "gemini-2.5-flash", DefaultTransformOptions())
|
||||
require.NoError(t, err)
|
||||
|
||||
var req V1InternalRequest
|
||||
require.NoError(t, json.Unmarshal(body, &req))
|
||||
require.NotNil(t, req.Request.SystemInstruction)
|
||||
|
||||
found := false
|
||||
for _, part := range req.Request.SystemInstruction.Parts {
|
||||
if strings.Contains(part.Text, "x-anthropic-billing-header keep") {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
require.True(t, found, "转换后的 systemInstruction 应保留 x-anthropic-billing-header 内容")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1008,3 +1008,114 @@ func TestAnthropicToResponses_ImageEmptyMediaType(t *testing.T) {
|
||||
// Should default to image/png when media_type is empty.
|
||||
assert.Equal(t, "data:image/png;base64,iVBOR", parts[0].ImageURL)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// normalizeToolParameters tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestNormalizeToolParameters(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input json.RawMessage
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "nil input",
|
||||
input: nil,
|
||||
expected: `{"type":"object","properties":{}}`,
|
||||
},
|
||||
{
|
||||
name: "empty input",
|
||||
input: json.RawMessage(``),
|
||||
expected: `{"type":"object","properties":{}}`,
|
||||
},
|
||||
{
|
||||
name: "null input",
|
||||
input: json.RawMessage(`null`),
|
||||
expected: `{"type":"object","properties":{}}`,
|
||||
},
|
||||
{
|
||||
name: "object without properties",
|
||||
input: json.RawMessage(`{"type":"object"}`),
|
||||
expected: `{"type":"object","properties":{}}`,
|
||||
},
|
||||
{
|
||||
name: "object with properties",
|
||||
input: json.RawMessage(`{"type":"object","properties":{"city":{"type":"string"}}}`),
|
||||
expected: `{"type":"object","properties":{"city":{"type":"string"}}}`,
|
||||
},
|
||||
{
|
||||
name: "non-object type",
|
||||
input: json.RawMessage(`{"type":"string"}`),
|
||||
expected: `{"type":"string"}`,
|
||||
},
|
||||
{
|
||||
name: "object with additional fields preserved",
|
||||
input: json.RawMessage(`{"type":"object","required":["name"]}`),
|
||||
expected: `{"type":"object","required":["name"],"properties":{}}`,
|
||||
},
|
||||
{
|
||||
name: "invalid JSON passthrough",
|
||||
input: json.RawMessage(`not json`),
|
||||
expected: `not json`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := normalizeToolParameters(tt.input)
|
||||
if tt.name == "invalid JSON passthrough" {
|
||||
assert.Equal(t, tt.expected, string(result))
|
||||
} else {
|
||||
assert.JSONEq(t, tt.expected, string(result))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_ToolWithoutProperties(t *testing.T) {
|
||||
req := &AnthropicRequest{
|
||||
Model: "gpt-5.2",
|
||||
MaxTokens: 1024,
|
||||
Messages: []AnthropicMessage{
|
||||
{Role: "user", Content: json.RawMessage(`"Hello"`)},
|
||||
},
|
||||
Tools: []AnthropicTool{
|
||||
{Name: "mcp__pencil__get_style_guide_tags", Description: "Get style tags", InputSchema: json.RawMessage(`{"type":"object"}`)},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := AnthropicToResponses(req)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Len(t, resp.Tools, 1)
|
||||
assert.Equal(t, "function", resp.Tools[0].Type)
|
||||
assert.Equal(t, "mcp__pencil__get_style_guide_tags", resp.Tools[0].Name)
|
||||
|
||||
// Parameters must have "properties" field after normalization.
|
||||
var params map[string]json.RawMessage
|
||||
require.NoError(t, json.Unmarshal(resp.Tools[0].Parameters, ¶ms))
|
||||
assert.Contains(t, params, "properties")
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_ToolWithNilSchema(t *testing.T) {
|
||||
req := &AnthropicRequest{
|
||||
Model: "gpt-5.2",
|
||||
MaxTokens: 1024,
|
||||
Messages: []AnthropicMessage{
|
||||
{Role: "user", Content: json.RawMessage(`"Hello"`)},
|
||||
},
|
||||
Tools: []AnthropicTool{
|
||||
{Name: "simple_tool", Description: "A tool"},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := AnthropicToResponses(req)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Len(t, resp.Tools, 1)
|
||||
var params map[string]json.RawMessage
|
||||
require.NoError(t, json.Unmarshal(resp.Tools[0].Parameters, ¶ms))
|
||||
assert.JSONEq(t, `"object"`, string(params["type"]))
|
||||
assert.JSONEq(t, `{}`, string(params["properties"]))
|
||||
}
|
||||
|
||||
@@ -409,8 +409,41 @@ func convertAnthropicToolsToResponses(tools []AnthropicTool) []ResponsesTool {
|
||||
Type: "function",
|
||||
Name: t.Name,
|
||||
Description: t.Description,
|
||||
Parameters: t.InputSchema,
|
||||
Parameters: normalizeToolParameters(t.InputSchema),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// normalizeToolParameters ensures the tool parameter schema is valid for
|
||||
// OpenAI's Responses API, which requires "properties" on object schemas.
|
||||
//
|
||||
// - nil/empty → {"type":"object","properties":{}}
|
||||
// - type=object without properties → adds "properties": {}
|
||||
// - otherwise → returned unchanged
|
||||
func normalizeToolParameters(schema json.RawMessage) json.RawMessage {
|
||||
if len(schema) == 0 || string(schema) == "null" {
|
||||
return json.RawMessage(`{"type":"object","properties":{}}`)
|
||||
}
|
||||
|
||||
var m map[string]json.RawMessage
|
||||
if err := json.Unmarshal(schema, &m); err != nil {
|
||||
return schema
|
||||
}
|
||||
|
||||
typ := m["type"]
|
||||
if string(typ) != `"object"` {
|
||||
return schema
|
||||
}
|
||||
|
||||
if _, ok := m["properties"]; ok {
|
||||
return schema
|
||||
}
|
||||
|
||||
m["properties"] = json.RawMessage(`{}`)
|
||||
out, err := json.Marshal(m)
|
||||
if err != nil {
|
||||
return schema
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -17,6 +17,7 @@ package httpclient
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -32,6 +33,8 @@ const (
|
||||
defaultMaxIdleConns = 100 // 最大空闲连接数
|
||||
defaultMaxIdleConnsPerHost = 10 // 每个主机最大空闲连接数
|
||||
defaultIdleConnTimeout = 90 * time.Second // 空闲连接超时时间(建议小于上游 LB 超时)
|
||||
defaultDialTimeout = 5 * time.Second // TCP 连接超时(含代理握手),代理不通时快速失败
|
||||
defaultTLSHandshakeTimeout = 5 * time.Second // TLS 握手超时
|
||||
validatedHostTTL = 30 * time.Second // DNS Rebinding 校验缓存 TTL
|
||||
)
|
||||
|
||||
@@ -107,6 +110,10 @@ func buildTransport(opts Options) (*http.Transport, error) {
|
||||
}
|
||||
|
||||
transport := &http.Transport{
|
||||
DialContext: (&net.Dialer{
|
||||
Timeout: defaultDialTimeout,
|
||||
}).DialContext,
|
||||
TLSHandshakeTimeout: defaultTLSHandshakeTimeout,
|
||||
MaxIdleConns: maxIdleConns,
|
||||
MaxIdleConnsPerHost: maxIdleConnsPerHost,
|
||||
MaxConnsPerHost: opts.MaxConnsPerHost, // 0 表示无限制
|
||||
|
||||
@@ -10,7 +10,13 @@ import (
|
||||
)
|
||||
|
||||
func TestInit_DualOutput(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
// Use os.MkdirTemp instead of t.TempDir to avoid cleanup failures
|
||||
// when lumberjack holds file handles on Windows.
|
||||
tmpDir, err := os.MkdirTemp("", "logger-test-*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = os.RemoveAll(tmpDir) })
|
||||
logPath := filepath.Join(tmpDir, "logs", "sub2api.log")
|
||||
|
||||
origStdout := os.Stdout
|
||||
@@ -57,7 +63,9 @@ func TestInit_DualOutput(t *testing.T) {
|
||||
|
||||
L().Info("dual-output-info")
|
||||
L().Warn("dual-output-warn")
|
||||
Sync()
|
||||
|
||||
// Skip Sync() — on Windows, fsync on pipes deadlocks (FlushFileBuffers).
|
||||
// The log data is already in the pipe buffer; closing writers is sufficient.
|
||||
|
||||
_ = stdoutW.Close()
|
||||
_ = stderrW.Close()
|
||||
@@ -166,7 +174,9 @@ func TestInit_CallerShouldPointToCallsite(t *testing.T) {
|
||||
}
|
||||
|
||||
L().Info("caller-check")
|
||||
Sync()
|
||||
// Skip Sync() — on Windows, fsync on pipes deadlocks (FlushFileBuffers).
|
||||
os.Stdout = origStdout
|
||||
os.Stderr = origStderr
|
||||
_ = stdoutW.Close()
|
||||
logBytes, _ := io.ReadAll(stdoutR)
|
||||
|
||||
|
||||
@@ -77,7 +77,7 @@ func TestStdLogBridgeRoutesLevels(t *testing.T) {
|
||||
log.Printf("service started")
|
||||
log.Printf("Warning: queue full")
|
||||
log.Printf("Forward request failed: timeout")
|
||||
Sync()
|
||||
// Skip Sync() — on Windows, fsync on pipes deadlocks (FlushFileBuffers).
|
||||
|
||||
_ = stdoutW.Close()
|
||||
_ = stderrW.Close()
|
||||
@@ -139,7 +139,7 @@ func TestLegacyPrintfRoutesLevels(t *testing.T) {
|
||||
LegacyPrintf("service.test", "request started")
|
||||
LegacyPrintf("service.test", "Warning: queue full")
|
||||
LegacyPrintf("service.test", "forward failed: timeout")
|
||||
Sync()
|
||||
// Skip Sync() — on Windows, fsync on pipes deadlocks (FlushFileBuffers).
|
||||
|
||||
_ = stdoutW.Close()
|
||||
_ = stderrW.Close()
|
||||
|
||||
@@ -47,6 +47,15 @@ func Created(c *gin.Context, data any) {
|
||||
})
|
||||
}
|
||||
|
||||
// Accepted 返回异步接受响应 (HTTP 202)
|
||||
func Accepted(c *gin.Context, data any) {
|
||||
c.JSON(http.StatusAccepted, Response{
|
||||
Code: 0,
|
||||
Message: "accepted",
|
||||
Data: data,
|
||||
})
|
||||
}
|
||||
|
||||
// Error 返回错误响应
|
||||
func Error(c *gin.Context, statusCode int, message string) {
|
||||
c.JSON(statusCode, Response{
|
||||
|
||||
@@ -3,6 +3,28 @@ package usagestats
|
||||
|
||||
import "time"
|
||||
|
||||
const (
|
||||
ModelSourceRequested = "requested"
|
||||
ModelSourceUpstream = "upstream"
|
||||
ModelSourceMapping = "mapping"
|
||||
)
|
||||
|
||||
func IsValidModelSource(source string) bool {
|
||||
switch source {
|
||||
case ModelSourceRequested, ModelSourceUpstream, ModelSourceMapping:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func NormalizeModelSource(source string) string {
|
||||
if IsValidModelSource(source) {
|
||||
return source
|
||||
}
|
||||
return ModelSourceRequested
|
||||
}
|
||||
|
||||
// DashboardStats 仪表盘统计
|
||||
type DashboardStats struct {
|
||||
// 用户统计
|
||||
@@ -90,6 +112,13 @@ type EndpointStat struct {
|
||||
ActualCost float64 `json:"actual_cost"` // 实际扣除
|
||||
}
|
||||
|
||||
// GroupUsageSummary represents today's and cumulative cost for a single group.
|
||||
type GroupUsageSummary struct {
|
||||
GroupID int64 `json:"group_id"`
|
||||
TodayCost float64 `json:"today_cost"`
|
||||
TotalCost float64 `json:"total_cost"`
|
||||
}
|
||||
|
||||
// GroupStat represents usage statistics for a single group
|
||||
type GroupStat struct {
|
||||
GroupID int64 `json:"group_id"`
|
||||
@@ -129,6 +158,25 @@ type UserSpendingRankingResponse struct {
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
}
|
||||
|
||||
// UserBreakdownItem represents per-user usage breakdown within a dimension (group, model, endpoint).
|
||||
type UserBreakdownItem struct {
|
||||
UserID int64 `json:"user_id"`
|
||||
Email string `json:"email"`
|
||||
Requests int64 `json:"requests"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
Cost float64 `json:"cost"` // 标准计费
|
||||
ActualCost float64 `json:"actual_cost"` // 实际扣除
|
||||
}
|
||||
|
||||
// UserBreakdownDimension specifies the dimension to filter for user breakdown.
|
||||
type UserBreakdownDimension struct {
|
||||
GroupID int64 // filter by group_id (>0 to enable)
|
||||
Model string // filter by model name (non-empty to enable)
|
||||
ModelType string // "requested", "upstream", or "mapping"
|
||||
Endpoint string // filter by endpoint value (non-empty to enable)
|
||||
EndpointType string // "inbound", "upstream", or "path"
|
||||
}
|
||||
|
||||
// APIKeyUsageTrendPoint represents API key usage trend data point
|
||||
type APIKeyUsageTrendPoint struct {
|
||||
Date string `json:"date"`
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
package usagestats
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestIsValidModelSource(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
source string
|
||||
want bool
|
||||
}{
|
||||
{name: "requested", source: ModelSourceRequested, want: true},
|
||||
{name: "upstream", source: ModelSourceUpstream, want: true},
|
||||
{name: "mapping", source: ModelSourceMapping, want: true},
|
||||
{name: "invalid", source: "foobar", want: false},
|
||||
{name: "empty", source: "", want: false},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := IsValidModelSource(tc.source); got != tc.want {
|
||||
t.Fatalf("IsValidModelSource(%q)=%v want %v", tc.source, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeModelSource(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
source string
|
||||
want string
|
||||
}{
|
||||
{name: "requested", source: ModelSourceRequested, want: ModelSourceRequested},
|
||||
{name: "upstream", source: ModelSourceUpstream, want: ModelSourceUpstream},
|
||||
{name: "mapping", source: ModelSourceMapping, want: ModelSourceMapping},
|
||||
{name: "invalid falls back", source: "foobar", want: ModelSourceRequested},
|
||||
{name: "empty falls back", source: "", want: ModelSourceRequested},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := NormalizeModelSource(tc.source); got != tc.want {
|
||||
t.Fatalf("NormalizeModelSource(%q)=%q want %q", tc.source, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -56,6 +56,7 @@ var schedulerNeutralExtraKeyPrefixes = []string{
|
||||
"codex_secondary_",
|
||||
"codex_5h_",
|
||||
"codex_7d_",
|
||||
"passive_usage_",
|
||||
}
|
||||
|
||||
var schedulerNeutralExtraKeys = map[string]struct{}{
|
||||
@@ -473,7 +474,9 @@ func (r *accountRepository) ListWithFilters(ctx context.Context, params paginati
|
||||
if search != "" {
|
||||
q = q.Where(dbaccount.NameContainsFold(search))
|
||||
}
|
||||
if groupID > 0 {
|
||||
if groupID == service.AccountListGroupUngrouped {
|
||||
q = q.Where(dbaccount.Not(dbaccount.HasAccountGroups()))
|
||||
} else if groupID > 0 {
|
||||
q = q.Where(dbaccount.HasAccountGroupsWith(dbaccountgroup.GroupIDEQ(groupID)))
|
||||
}
|
||||
|
||||
|
||||
@@ -214,6 +214,7 @@ func (s *AccountRepoSuite) TestListWithFilters() {
|
||||
accType string
|
||||
status string
|
||||
search string
|
||||
groupID int64
|
||||
wantCount int
|
||||
validate func(accounts []service.Account)
|
||||
}{
|
||||
@@ -265,6 +266,21 @@ func (s *AccountRepoSuite) TestListWithFilters() {
|
||||
s.Require().Contains(accounts[0].Name, "alpha")
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "filter_by_ungrouped",
|
||||
setup: func(client *dbent.Client) {
|
||||
group := mustCreateGroup(s.T(), client, &service.Group{Name: "g-ungrouped"})
|
||||
grouped := mustCreateAccount(s.T(), client, &service.Account{Name: "grouped-account"})
|
||||
mustCreateAccount(s.T(), client, &service.Account{Name: "ungrouped-account"})
|
||||
mustBindAccountToGroup(s.T(), client, grouped.ID, group.ID, 1)
|
||||
},
|
||||
groupID: service.AccountListGroupUngrouped,
|
||||
wantCount: 1,
|
||||
validate: func(accounts []service.Account) {
|
||||
s.Require().Equal("ungrouped-account", accounts[0].Name)
|
||||
s.Require().Empty(accounts[0].GroupIDs)
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -277,7 +293,7 @@ func (s *AccountRepoSuite) TestListWithFilters() {
|
||||
|
||||
tt.setup(client)
|
||||
|
||||
accounts, _, err := repo.ListWithFilters(ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, tt.platform, tt.accType, tt.status, tt.search, 0)
|
||||
accounts, _, err := repo.ListWithFilters(ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, tt.platform, tt.accType, tt.status, tt.search, tt.groupID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(accounts, tt.wantCount)
|
||||
if tt.validate != nil {
|
||||
|
||||
@@ -410,6 +410,16 @@ func (r *apiKeyRepository) ClearGroupIDByGroupID(ctx context.Context, groupID in
|
||||
return int64(n), err
|
||||
}
|
||||
|
||||
// UpdateGroupIDByUserAndGroup 将用户下绑定 oldGroupID 的所有 Key 迁移到 newGroupID
|
||||
func (r *apiKeyRepository) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
|
||||
client := clientFromContext(ctx, r.client)
|
||||
n, err := client.APIKey.Update().
|
||||
Where(apikey.UserIDEQ(userID), apikey.GroupIDEQ(oldGroupID), apikey.DeletedAtIsNil()).
|
||||
SetGroupID(newGroupID).
|
||||
Save(ctx)
|
||||
return int64(n), err
|
||||
}
|
||||
|
||||
// CountByGroupID 获取分组的 API Key 数量
|
||||
func (r *apiKeyRepository) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
count, err := r.activeQuery().Where(apikey.GroupIDEQ(groupID)).Count(ctx)
|
||||
|
||||
@@ -57,6 +57,7 @@ func NewS3BackupStoreFactory() service.BackupObjectStoreFactory {
|
||||
|
||||
func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader, contentType string) (int64, error) {
|
||||
// 读取全部内容以获取大小(S3 PutObject 需要知道内容长度)
|
||||
// 注意:阿里云 OSS 不兼容 s3manager 分片上传的签名方式,因此使用 PutObject
|
||||
data, err := io.ReadAll(body)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("read body: %w", err)
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
_ "embed"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -12,9 +13,9 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
stickySessionPrefix = "sticky_session:"
|
||||
clientAffinityPrefix = "client_affinity:"
|
||||
clientAffinityReversePrefix = "client_affinity_rev:"
|
||||
stickySessionPrefix = "sticky_session:"
|
||||
affinityKeyPrefix = "affinity:"
|
||||
affinityRevKeyPrefix = "affinity_rev:"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -30,6 +31,8 @@ var (
|
||||
getAffinityClientsWithScoresLua string
|
||||
//go:embed lua/clear_account_affinity.lua
|
||||
clearAccountAffinityLua string
|
||||
//go:embed lua/get_affinity_multi_count.lua
|
||||
getAffinityMultiCountLua string
|
||||
|
||||
getAffinityScript = redis.NewScript(getAffinityLua)
|
||||
updateAffinityScript = redis.NewScript(updateAffinityLua)
|
||||
@@ -37,6 +40,7 @@ var (
|
||||
getAffinityClientsScript = redis.NewScript(getAffinityClientsLua)
|
||||
getAffinityClientsWithScoresScript = redis.NewScript(getAffinityClientsWithScoresLua)
|
||||
clearAccountAffinityScript = redis.NewScript(clearAccountAffinityLua)
|
||||
getAffinityMultiCountScript = redis.NewScript(getAffinityMultiCountLua)
|
||||
)
|
||||
|
||||
type gatewayCache struct {
|
||||
@@ -84,20 +88,39 @@ func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
}
|
||||
|
||||
// buildAffinityKey 构建正向亲和 key(client → accounts)
|
||||
// 格式: client_affinity:{groupID}:{clientID}
|
||||
func buildAffinityKey(groupID int64, clientID string) string {
|
||||
return fmt.Sprintf("%s%d:%s", clientAffinityPrefix, groupID, clientID)
|
||||
// buildAffinityKey 构建正向亲和 key(member → accounts)
|
||||
// 格式: affinity:{groupID}:{userID}/{clientID}
|
||||
func buildAffinityKey(groupID int64, userID int64, clientID string) string {
|
||||
return fmt.Sprintf("%s%d:%s", affinityKeyPrefix, groupID, buildAffinityMember(userID, clientID))
|
||||
}
|
||||
|
||||
// buildAffinityReverseKey 构建反向亲和 key(account → clients)
|
||||
// 格式: client_affinity_rev:{groupID}:{accountID}
|
||||
// buildAffinityReverseKey 构建反向亲和 key(account → members)
|
||||
// 格式: affinity_rev:{groupID}:{accountID}
|
||||
func buildAffinityReverseKey(groupID int64, accountID int64) string {
|
||||
return fmt.Sprintf("%s%d:%d", clientAffinityReversePrefix, groupID, accountID)
|
||||
return fmt.Sprintf("%s%d:%d", affinityRevKeyPrefix, groupID, accountID)
|
||||
}
|
||||
|
||||
func (c *gatewayCache) GetClientAffinityAccounts(ctx context.Context, groupID int64, clientID string, ttl time.Duration) ([]int64, error) {
|
||||
key := buildAffinityKey(groupID, clientID)
|
||||
// buildAffinityMember 构建亲和成员标识
|
||||
// 格式: {userID}/{clientID}
|
||||
func buildAffinityMember(userID int64, clientID string) string {
|
||||
return fmt.Sprintf("%d/%s", userID, clientID)
|
||||
}
|
||||
|
||||
// parseAffinityMember 解析亲和成员标识为 userID 和 clientID
|
||||
func parseAffinityMember(member string) (userID int64, clientID string) {
|
||||
idx := strings.IndexByte(member, '/')
|
||||
if idx < 0 {
|
||||
// 兼容旧格式(纯 clientID,无 userID 前缀)
|
||||
return 0, member
|
||||
}
|
||||
userID, _ = strconv.ParseInt(member[:idx], 10, 64)
|
||||
clientID = member[idx+1:]
|
||||
return userID, clientID
|
||||
}
|
||||
|
||||
// GetAffinityAccounts 获取亲和账号列表(按最近使用降序),同时清理过期成员
|
||||
func (c *gatewayCache) GetAffinityAccounts(ctx context.Context, groupID int64, userID int64, clientID string, ttl time.Duration) ([]int64, error) {
|
||||
key := buildAffinityKey(groupID, userID, clientID)
|
||||
now := time.Now().Unix()
|
||||
expireThreshold := now - int64(ttl.Seconds())
|
||||
|
||||
@@ -120,19 +143,21 @@ func (c *gatewayCache) GetClientAffinityAccounts(ctx context.Context, groupID in
|
||||
return accountIDs, nil
|
||||
}
|
||||
|
||||
func (c *gatewayCache) UpdateClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64, ttl time.Duration) error {
|
||||
fwdKey := buildAffinityKey(groupID, clientID)
|
||||
// UpdateAffinity 添加/更新亲和关系(更新 score 为当前时间戳,刷新 key TTL)
|
||||
func (c *gatewayCache) UpdateAffinity(ctx context.Context, groupID int64, userID int64, clientID string, accountID int64, ttl time.Duration) error {
|
||||
fwdKey := buildAffinityKey(groupID, userID, clientID)
|
||||
revKey := buildAffinityReverseKey(groupID, accountID)
|
||||
now := time.Now().Unix()
|
||||
ttlSeconds := int64(ttl.Seconds())
|
||||
expireThreshold := now - ttlSeconds
|
||||
|
||||
member := buildAffinityMember(userID, clientID)
|
||||
return updateAffinityScript.Run(ctx, c.rdb, []string{fwdKey, revKey},
|
||||
now, ttlSeconds, accountID, expireThreshold, clientID,
|
||||
now, ttlSeconds, accountID, expireThreshold, member,
|
||||
).Err()
|
||||
}
|
||||
|
||||
// GetAccountAffinityCountBatch 批量获取账号的亲和客户端数量(惰性清理过期成员)
|
||||
// GetAccountAffinityCountBatch 批量获取账号的亲和成员数量(惰性清理过期成员)
|
||||
func (c *gatewayCache) GetAccountAffinityCountBatch(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error) {
|
||||
if len(accountIDs) == 0 {
|
||||
return map[int64]int64{}, nil
|
||||
@@ -162,8 +187,9 @@ func (c *gatewayCache) GetAccountAffinityCountBatch(ctx context.Context, groupID
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// GetAccountAffinityClientsBatch 批量获取每个账号跨所有分组的亲和客户端列表(去重)。
|
||||
// GetAccountAffinityClientsBatch 批量获取每个账号跨所有分组的亲和成员列表(去重)。
|
||||
// accountGroups: map[accountID][]groupID,对每个 (groupID, accountID) 组合查询反向索引。
|
||||
// 返回值成员格式为 {userID}/{clientID}。
|
||||
func (c *gatewayCache) GetAccountAffinityClientsBatch(ctx context.Context, accountGroups map[int64][]int64, ttl time.Duration) (map[int64][]string, error) {
|
||||
if len(accountGroups) == 0 {
|
||||
return map[int64][]string{}, nil
|
||||
@@ -197,21 +223,21 @@ func (c *gatewayCache) GetAccountAffinityClientsBatch(ctx context.Context, accou
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 合并结果:同一个 accountID 跨多个 group 的 clientID 去重
|
||||
// 合并结果:同一个 accountID 跨多个 group 的成员去重
|
||||
result := make(map[int64][]string, len(accountGroups))
|
||||
seen := make(map[int64]map[string]struct{}, len(accountGroups))
|
||||
for i, q := range queries {
|
||||
clients, _ := cmds[i].StringSlice()
|
||||
if len(clients) == 0 {
|
||||
members, _ := cmds[i].StringSlice()
|
||||
if len(members) == 0 {
|
||||
continue
|
||||
}
|
||||
if seen[q.accountID] == nil {
|
||||
seen[q.accountID] = make(map[string]struct{})
|
||||
}
|
||||
for _, clientID := range clients {
|
||||
if _, exists := seen[q.accountID][clientID]; !exists {
|
||||
seen[q.accountID][clientID] = struct{}{}
|
||||
result[q.accountID] = append(result[q.accountID], clientID)
|
||||
for _, member := range members {
|
||||
if _, exists := seen[q.accountID][member]; !exists {
|
||||
seen[q.accountID][member] = struct{}{}
|
||||
result[q.accountID] = append(result[q.accountID], member)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -245,25 +271,32 @@ func (c *gatewayCache) GetAccountAffinityClientsWithScores(
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 合并跨组结果,同一 clientID 取最近的 lastActive
|
||||
seen := make(map[string]int64) // clientID → max timestamp
|
||||
// 合并跨组结果,同一 member 取最近的 lastActive
|
||||
type memberInfo struct {
|
||||
userID int64
|
||||
clientID string
|
||||
ts int64
|
||||
}
|
||||
seen := make(map[string]*memberInfo) // member string → info
|
||||
for _, cmd := range cmds {
|
||||
vals, _ := cmd.StringSlice()
|
||||
// vals 格式: [clientID1, score1, clientID2, score2, ...]
|
||||
// vals 格式: [member1, score1, member2, score2, ...]
|
||||
for j := 0; j+1 < len(vals); j += 2 {
|
||||
clientID := vals[j]
|
||||
member := vals[j]
|
||||
ts, _ := strconv.ParseInt(vals[j+1], 10, 64)
|
||||
if existing, ok := seen[clientID]; !ok || ts > existing {
|
||||
seen[clientID] = ts
|
||||
if existing, ok := seen[member]; !ok || ts > existing.ts {
|
||||
uid, cid := parseAffinityMember(member)
|
||||
seen[member] = &memberInfo{userID: uid, clientID: cid, ts: ts}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result := make([]service.AffinityClient, 0, len(seen))
|
||||
for clientID, ts := range seen {
|
||||
for _, info := range seen {
|
||||
result = append(result, service.AffinityClient{
|
||||
ClientID: clientID,
|
||||
LastActive: time.Unix(ts, 0),
|
||||
UserID: info.userID,
|
||||
ClientID: info.clientID,
|
||||
LastActive: time.Unix(info.ts, 0),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -274,8 +307,8 @@ func (c *gatewayCache) GetAccountAffinityClientsWithScores(
|
||||
}
|
||||
|
||||
// ClearAccountAffinity 清除指定账号在所有分组的亲和记录(正向+反向索引)。
|
||||
// 对每个 groupID 执行 Lua 脚本:读取反向索引获取所有客户端,
|
||||
// 从每个客户端的正向索引中移除该账号,然后删除反向索引。
|
||||
// 对每个 groupID 执行 Lua 脚本:读取反向索引获取所有成员,
|
||||
// 从每个成员的正向索引中移除该账号,然后删除反向索引。
|
||||
func (c *gatewayCache) ClearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) error {
|
||||
if len(groupIDs) == 0 {
|
||||
return nil
|
||||
@@ -294,3 +327,37 @@ func (c *gatewayCache) ClearAccountAffinity(ctx context.Context, accountID int64
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAffinityMultiCount 获取账号的多维度亲和计数。
|
||||
// 返回: uniqueUsers(独立用户数), uniqueClients(独立客户端数), perUserClients(目标用户的客户端数)
|
||||
func (c *gatewayCache) GetAffinityMultiCount(
|
||||
ctx context.Context,
|
||||
groupID int64,
|
||||
accountID int64,
|
||||
targetUserID int64,
|
||||
ttl time.Duration,
|
||||
) (users, clients, perUser int64, err error) {
|
||||
key := buildAffinityReverseKey(groupID, accountID)
|
||||
now := time.Now().Unix()
|
||||
expireThreshold := now - int64(ttl.Seconds())
|
||||
|
||||
targetUserStr := ""
|
||||
if targetUserID > 0 {
|
||||
targetUserStr = strconv.FormatInt(targetUserID, 10)
|
||||
}
|
||||
|
||||
result, err := getAffinityMultiCountScript.Run(ctx, c.rdb, []string{key}, expireThreshold, targetUserStr).Int64Slice()
|
||||
if err != nil {
|
||||
if err == redis.Nil {
|
||||
return 0, 0, 0, nil
|
||||
}
|
||||
return 0, 0, 0, err
|
||||
}
|
||||
|
||||
if len(result) < 4 {
|
||||
return 0, 0, 0, nil
|
||||
}
|
||||
|
||||
// result: {totalMembers, uniqueUsers, uniqueClients, perUserClients}
|
||||
return result[1], result[2], result[3], nil
|
||||
}
|
||||
|
||||
@@ -89,8 +89,9 @@ func (r *groupRepository) GetByID(ctx context.Context, id int64) (*service.Group
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
count, _ := r.GetAccountCount(ctx, out.ID)
|
||||
out.AccountCount = count
|
||||
total, active, _ := r.GetAccountCount(ctx, out.ID)
|
||||
out.AccountCount = total
|
||||
out.ActiveAccountCount = active
|
||||
return out, nil
|
||||
}
|
||||
|
||||
@@ -258,7 +259,10 @@ func (r *groupRepository) ListWithFilters(ctx context.Context, params pagination
|
||||
counts, err := r.loadAccountCounts(ctx, groupIDs)
|
||||
if err == nil {
|
||||
for i := range outGroups {
|
||||
outGroups[i].AccountCount = counts[outGroups[i].ID]
|
||||
c := counts[outGroups[i].ID]
|
||||
outGroups[i].AccountCount = c.Total
|
||||
outGroups[i].ActiveAccountCount = c.Active
|
||||
outGroups[i].RateLimitedAccountCount = c.RateLimited
|
||||
}
|
||||
}
|
||||
|
||||
@@ -285,7 +289,10 @@ func (r *groupRepository) ListActive(ctx context.Context) ([]service.Group, erro
|
||||
counts, err := r.loadAccountCounts(ctx, groupIDs)
|
||||
if err == nil {
|
||||
for i := range outGroups {
|
||||
outGroups[i].AccountCount = counts[outGroups[i].ID]
|
||||
c := counts[outGroups[i].ID]
|
||||
outGroups[i].AccountCount = c.Total
|
||||
outGroups[i].ActiveAccountCount = c.Active
|
||||
outGroups[i].RateLimitedAccountCount = c.RateLimited
|
||||
}
|
||||
}
|
||||
|
||||
@@ -312,7 +319,10 @@ func (r *groupRepository) ListActiveByPlatform(ctx context.Context, platform str
|
||||
counts, err := r.loadAccountCounts(ctx, groupIDs)
|
||||
if err == nil {
|
||||
for i := range outGroups {
|
||||
outGroups[i].AccountCount = counts[outGroups[i].ID]
|
||||
c := counts[outGroups[i].ID]
|
||||
outGroups[i].AccountCount = c.Total
|
||||
outGroups[i].ActiveAccountCount = c.Active
|
||||
outGroups[i].RateLimitedAccountCount = c.RateLimited
|
||||
}
|
||||
}
|
||||
|
||||
@@ -371,12 +381,20 @@ func (r *groupRepository) ExistsByIDs(ctx context.Context, ids []int64) (map[int
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (r *groupRepository) GetAccountCount(ctx context.Context, groupID int64) (int64, error) {
|
||||
var count int64
|
||||
if err := scanSingleRow(ctx, r.sql, "SELECT COUNT(*) FROM account_groups WHERE group_id = $1", []any{groupID}, &count); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return count, nil
|
||||
func (r *groupRepository) GetAccountCount(ctx context.Context, groupID int64) (total int64, active int64, err error) {
|
||||
var rateLimited int64
|
||||
err = scanSingleRow(ctx, r.sql,
|
||||
`SELECT COUNT(*),
|
||||
COUNT(*) FILTER (WHERE a.status = 'active' AND a.schedulable = true),
|
||||
COUNT(*) FILTER (WHERE a.status = 'active' AND (
|
||||
a.rate_limit_reset_at > NOW() OR
|
||||
a.overload_until > NOW() OR
|
||||
a.temp_unschedulable_until > NOW()
|
||||
))
|
||||
FROM account_groups ag JOIN accounts a ON a.id = ag.account_id
|
||||
WHERE ag.group_id = $1`,
|
||||
[]any{groupID}, &total, &active, &rateLimited)
|
||||
return
|
||||
}
|
||||
|
||||
func (r *groupRepository) DeleteAccountGroupsByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
@@ -502,15 +520,32 @@ func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64,
|
||||
return affectedUserIDs, nil
|
||||
}
|
||||
|
||||
func (r *groupRepository) loadAccountCounts(ctx context.Context, groupIDs []int64) (counts map[int64]int64, err error) {
|
||||
counts = make(map[int64]int64, len(groupIDs))
|
||||
type groupAccountCounts struct {
|
||||
Total int64
|
||||
Active int64
|
||||
RateLimited int64
|
||||
}
|
||||
|
||||
func (r *groupRepository) loadAccountCounts(ctx context.Context, groupIDs []int64) (counts map[int64]groupAccountCounts, err error) {
|
||||
counts = make(map[int64]groupAccountCounts, len(groupIDs))
|
||||
if len(groupIDs) == 0 {
|
||||
return counts, nil
|
||||
}
|
||||
|
||||
rows, err := r.sql.QueryContext(
|
||||
ctx,
|
||||
"SELECT group_id, COUNT(*) FROM account_groups WHERE group_id = ANY($1) GROUP BY group_id",
|
||||
`SELECT ag.group_id,
|
||||
COUNT(*) AS total,
|
||||
COUNT(*) FILTER (WHERE a.status = 'active' AND a.schedulable = true) AS active,
|
||||
COUNT(*) FILTER (WHERE a.status = 'active' AND (
|
||||
a.rate_limit_reset_at > NOW() OR
|
||||
a.overload_until > NOW() OR
|
||||
a.temp_unschedulable_until > NOW()
|
||||
)) AS rate_limited
|
||||
FROM account_groups ag
|
||||
JOIN accounts a ON a.id = ag.account_id
|
||||
WHERE ag.group_id = ANY($1)
|
||||
GROUP BY ag.group_id`,
|
||||
pq.Array(groupIDs),
|
||||
)
|
||||
if err != nil {
|
||||
@@ -525,11 +560,11 @@ func (r *groupRepository) loadAccountCounts(ctx context.Context, groupIDs []int6
|
||||
|
||||
for rows.Next() {
|
||||
var groupID int64
|
||||
var count int64
|
||||
if err = rows.Scan(&groupID, &count); err != nil {
|
||||
var c groupAccountCounts
|
||||
if err = rows.Scan(&groupID, &c.Total, &c.Active, &c.RateLimited); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
counts[groupID] = count
|
||||
counts[groupID] = c
|
||||
}
|
||||
if err = rows.Err(); err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -603,7 +603,7 @@ func (s *GroupRepoSuite) TestGetAccountCount() {
|
||||
_, err = s.tx.ExecContext(s.ctx, "INSERT INTO account_groups (account_id, group_id, priority, created_at) VALUES ($1, $2, $3, NOW())", a2, group.ID, 2)
|
||||
s.Require().NoError(err)
|
||||
|
||||
count, err := s.repo.GetAccountCount(s.ctx, group.ID)
|
||||
count, _, err := s.repo.GetAccountCount(s.ctx, group.ID)
|
||||
s.Require().NoError(err, "GetAccountCount")
|
||||
s.Require().Equal(int64(2), count)
|
||||
}
|
||||
@@ -619,7 +619,7 @@ func (s *GroupRepoSuite) TestGetAccountCount_Empty() {
|
||||
}
|
||||
s.Require().NoError(s.repo.Create(s.ctx, group))
|
||||
|
||||
count, err := s.repo.GetAccountCount(s.ctx, group.ID)
|
||||
count, _, err := s.repo.GetAccountCount(s.ctx, group.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().Zero(count)
|
||||
}
|
||||
@@ -651,7 +651,7 @@ func (s *GroupRepoSuite) TestDeleteAccountGroupsByGroupID() {
|
||||
s.Require().NoError(err, "DeleteAccountGroupsByGroupID")
|
||||
s.Require().Equal(int64(1), affected, "expected 1 affected row")
|
||||
|
||||
count, err := s.repo.GetAccountCount(s.ctx, g.ID)
|
||||
count, _, err := s.repo.GetAccountCount(s.ctx, g.ID)
|
||||
s.Require().NoError(err, "GetAccountCount")
|
||||
s.Require().Equal(int64(0), count, "expected 0 account groups")
|
||||
}
|
||||
@@ -692,7 +692,7 @@ func (s *GroupRepoSuite) TestDeleteAccountGroupsByGroupID_MultipleAccounts() {
|
||||
s.Require().NoError(err)
|
||||
s.Require().Equal(int64(3), affected)
|
||||
|
||||
count, _ := s.repo.GetAccountCount(s.ctx, g.ID)
|
||||
count, _, _ := s.repo.GetAccountCount(s.ctx, g.ID)
|
||||
s.Require().Zero(count)
|
||||
}
|
||||
|
||||
|
||||
@@ -1,21 +1,21 @@
|
||||
-- 清除单个账号在指定分组的所有亲和记录(正向+反向)
|
||||
-- KEYS[1] = client_affinity_rev:{groupID}:{accountID} (反向索引)
|
||||
-- KEYS[1] = affinity_rev:{groupID}:{accountID} (反向索引)
|
||||
-- ARGV[1] = groupID (用于构建正向 key)
|
||||
-- ARGV[2] = accountID (正向索引中要移除的成员)
|
||||
-- 返回: 清理的客户端数量
|
||||
-- 返回: 清理的成员数量
|
||||
local rev_key = KEYS[1]
|
||||
local group_id = ARGV[1]
|
||||
local account_id = ARGV[2]
|
||||
|
||||
-- 获取反向索引中所有客户端 ID
|
||||
local clients = redis.call('ZRANGE', rev_key, 0, -1)
|
||||
if #clients == 0 then
|
||||
-- 获取反向索引中所有成员 ({userID}/{clientID})
|
||||
local members = redis.call('ZRANGE', rev_key, 0, -1)
|
||||
if #members == 0 then
|
||||
return 0
|
||||
end
|
||||
|
||||
-- 从每个客户端的正向索引中移除该账号
|
||||
for _, client_id in ipairs(clients) do
|
||||
local fwd_key = 'client_affinity:' .. group_id .. ':' .. client_id
|
||||
-- 从每个成员的正向索引中移除该账号
|
||||
for _, member in ipairs(members) do
|
||||
local fwd_key = 'affinity:' .. group_id .. ':' .. member
|
||||
redis.call('ZREM', fwd_key, account_id)
|
||||
-- 如果正向索引为空,删除 key
|
||||
if redis.call('ZCARD', fwd_key) == 0 then
|
||||
@@ -26,4 +26,4 @@ end
|
||||
-- 删除反向索引
|
||||
redis.call('DEL', rev_key)
|
||||
|
||||
return #clients
|
||||
return #members
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
-- 清理过期成员后返回亲和账号列表(按最近使用降序)
|
||||
-- KEYS[1] = client_affinity:{groupID}:{clientID}
|
||||
-- KEYS[1] = affinity:{groupID}:{userID}/{clientID}
|
||||
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
|
||||
return redis.call('ZREVRANGE', KEYS[1], 0, -1)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
-- 清理过期成员后返回反向索引的 clientID 列表(按最近使用降序)
|
||||
-- KEYS[1] = client_affinity_rev:{groupID}:{accountID}
|
||||
-- 清理过期成员后返回反向索引的成员列表(按最近使用降序)
|
||||
-- 成员格式: {userID}/{clientID}
|
||||
-- KEYS[1] = affinity_rev:{groupID}:{accountID}
|
||||
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
|
||||
return redis.call('ZREVRANGE', KEYS[1], 0, -1)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
-- 清理过期成员后返回反向索引的 clientID 列表及其 score(最后活跃时间戳)
|
||||
-- KEYS[1] = client_affinity_rev:{groupID}:{accountID}
|
||||
-- 清理过期成员后返回反向索引的成员列表及其 score(最后活跃时间戳)
|
||||
-- 成员格式: {userID}/{clientID}
|
||||
-- KEYS[1] = affinity_rev:{groupID}:{accountID}
|
||||
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
|
||||
-- 返回: {clientID1, score1, clientID2, score2, ...}(按最近使用降序)
|
||||
-- 返回: {member1, score1, member2, score2, ...}(按最近使用降序)
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
|
||||
return redis.call('ZREVRANGEBYSCORE', KEYS[1], '+inf', '-inf', 'WITHSCORES')
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
-- 清理过期成员后返回反向索引的成员数量
|
||||
-- KEYS[1] = client_affinity_rev:{groupID}:{accountID}
|
||||
-- KEYS[1] = affinity_rev:{groupID}:{accountID}
|
||||
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
|
||||
return redis.call('ZCARD', KEYS[1])
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
-- 从反向索引解析多维度计数(用户/客户端/每用户客户端)
|
||||
-- KEYS[1] = affinity_rev:{groupID}:{accountID}
|
||||
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
|
||||
-- ARGV[2] = 目标 userID(传 "" 则不计算 perUserClients)
|
||||
-- 返回: {totalMembers, uniqueUsers, uniqueClients, perUserClients}
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
|
||||
|
||||
local members = redis.call('ZRANGE', KEYS[1], 0, -1)
|
||||
local total = #members
|
||||
if total == 0 then
|
||||
return {0, 0, 0, 0}
|
||||
end
|
||||
|
||||
local target_user = ARGV[2]
|
||||
local users = {}
|
||||
local clients = {}
|
||||
local user_count = 0
|
||||
local client_count = 0
|
||||
local per_user_count = 0
|
||||
|
||||
for _, member in ipairs(members) do
|
||||
local sep = string.find(member, '/', 1, true)
|
||||
if sep then
|
||||
local uid = string.sub(member, 1, sep - 1)
|
||||
local cid = string.sub(member, sep + 1)
|
||||
if not users[uid] then
|
||||
users[uid] = true
|
||||
user_count = user_count + 1
|
||||
end
|
||||
if not clients[cid] then
|
||||
clients[cid] = true
|
||||
client_count = client_count + 1
|
||||
end
|
||||
if target_user ~= '' and uid == target_user then
|
||||
per_user_count = per_user_count + 1
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
return {total, user_count, client_count, per_user_count}
|
||||
@@ -1,11 +1,11 @@
|
||||
-- 原子双写正向+反向索引
|
||||
-- KEYS[1] = client_affinity:{groupID}:{clientID} (正向: client → accounts)
|
||||
-- KEYS[2] = client_affinity_rev:{groupID}:{accountID} (反向: account → clients)
|
||||
-- KEYS[1] = affinity:{groupID}:{userID}/{clientID} (正向: member → accounts)
|
||||
-- KEYS[2] = affinity_rev:{groupID}:{accountID} (反向: account → members)
|
||||
-- ARGV[1] = 当前时间戳 (score)
|
||||
-- ARGV[2] = TTL 秒数
|
||||
-- ARGV[3] = accountID (正向索引的成员)
|
||||
-- ARGV[4] = 过期阈值时间戳 (now - ttl)
|
||||
-- ARGV[5] = clientID (反向索引的成员)
|
||||
-- ARGV[5] = {userID}/{clientID} (反向索引的成员)
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[4])
|
||||
redis.call('ZADD', KEYS[1], ARGV[1], ARGV[3])
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||||
|
||||
@@ -40,7 +40,7 @@ func NewProxyExitInfoProber(cfg *config.Config) service.ProxyExitInfoProber {
|
||||
}
|
||||
|
||||
const (
|
||||
defaultProxyProbeTimeout = 30 * time.Second
|
||||
defaultProxyProbeTimeout = 10 * time.Second
|
||||
defaultProxyProbeResponseMaxBytes = int64(1024 * 1024)
|
||||
)
|
||||
|
||||
|
||||
@@ -28,7 +28,7 @@ import (
|
||||
gocache "github.com/patrickmn/go-cache"
|
||||
)
|
||||
|
||||
const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, media_type, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, created_at"
|
||||
const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, media_type, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, created_at"
|
||||
|
||||
var usageLogInsertArgTypes = [...]string{
|
||||
"bigint",
|
||||
@@ -36,6 +36,7 @@ var usageLogInsertArgTypes = [...]string{
|
||||
"bigint",
|
||||
"text",
|
||||
"text",
|
||||
"text",
|
||||
"bigint",
|
||||
"bigint",
|
||||
"integer",
|
||||
@@ -277,6 +278,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor,
|
||||
account_id,
|
||||
request_id,
|
||||
model,
|
||||
upstream_model,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -311,12 +313,12 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor,
|
||||
cache_ttl_overridden,
|
||||
created_at
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5,
|
||||
$6, $7,
|
||||
$8, $9, $10, $11,
|
||||
$12, $13,
|
||||
$14, $15, $16, $17, $18, $19,
|
||||
$20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38
|
||||
$1, $2, $3, $4, $5, $6,
|
||||
$7, $8,
|
||||
$9, $10, $11, $12,
|
||||
$13, $14,
|
||||
$15, $16, $17, $18, $19, $20,
|
||||
$21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39
|
||||
)
|
||||
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
||||
RETURNING id, created_at
|
||||
@@ -707,6 +709,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
account_id,
|
||||
request_id,
|
||||
model,
|
||||
upstream_model,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -742,7 +745,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
created_at
|
||||
) AS (VALUES `)
|
||||
|
||||
args := make([]any, 0, len(keys)*38)
|
||||
args := make([]any, 0, len(keys)*39)
|
||||
argPos := 1
|
||||
for idx, key := range keys {
|
||||
if idx > 0 {
|
||||
@@ -776,6 +779,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
account_id,
|
||||
request_id,
|
||||
model,
|
||||
upstream_model,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -816,6 +820,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
account_id,
|
||||
request_id,
|
||||
model,
|
||||
upstream_model,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -896,6 +901,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
account_id,
|
||||
request_id,
|
||||
model,
|
||||
upstream_model,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -931,7 +937,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
created_at
|
||||
) AS (VALUES `)
|
||||
|
||||
args := make([]any, 0, len(preparedList)*38)
|
||||
args := make([]any, 0, len(preparedList)*39)
|
||||
argPos := 1
|
||||
for idx, prepared := range preparedList {
|
||||
if idx > 0 {
|
||||
@@ -962,6 +968,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
account_id,
|
||||
request_id,
|
||||
model,
|
||||
upstream_model,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -1002,6 +1009,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
account_id,
|
||||
request_id,
|
||||
model,
|
||||
upstream_model,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -1050,6 +1058,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared
|
||||
account_id,
|
||||
request_id,
|
||||
model,
|
||||
upstream_model,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -1084,12 +1093,12 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared
|
||||
cache_ttl_overridden,
|
||||
created_at
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5,
|
||||
$6, $7,
|
||||
$8, $9, $10, $11,
|
||||
$12, $13,
|
||||
$14, $15, $16, $17, $18, $19,
|
||||
$20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38
|
||||
$1, $2, $3, $4, $5, $6,
|
||||
$7, $8,
|
||||
$9, $10, $11, $12,
|
||||
$13, $14,
|
||||
$15, $16, $17, $18, $19, $20,
|
||||
$21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39
|
||||
)
|
||||
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
||||
`, prepared.args...)
|
||||
@@ -1121,6 +1130,7 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
|
||||
reasoningEffort := nullString(log.ReasoningEffort)
|
||||
inboundEndpoint := nullString(log.InboundEndpoint)
|
||||
upstreamEndpoint := nullString(log.UpstreamEndpoint)
|
||||
upstreamModel := nullString(log.UpstreamModel)
|
||||
|
||||
var requestIDArg any
|
||||
if requestID != "" {
|
||||
@@ -1138,6 +1148,7 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
|
||||
log.AccountID,
|
||||
requestIDArg,
|
||||
log.Model,
|
||||
upstreamModel,
|
||||
groupID,
|
||||
subscriptionID,
|
||||
log.InputTokens,
|
||||
@@ -2864,15 +2875,26 @@ func (r *usageLogRepository) getUsageTrendFromAggregates(ctx context.Context, st
|
||||
|
||||
// GetModelStatsWithFilters returns model statistics with optional filters
|
||||
func (r *usageLogRepository) GetModelStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) (results []ModelStat, err error) {
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType, usagestats.ModelSourceRequested)
|
||||
}
|
||||
|
||||
// GetModelStatsWithFiltersBySource returns model statistics with optional filters and model source dimension.
|
||||
// source: requested | upstream | mapping.
|
||||
func (r *usageLogRepository) GetModelStatsWithFiltersBySource(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8, source string) (results []ModelStat, err error) {
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType, source)
|
||||
}
|
||||
|
||||
func (r *usageLogRepository) getModelStatsWithFiltersBySource(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8, source string) (results []ModelStat, err error) {
|
||||
actualCostExpr := "COALESCE(SUM(actual_cost), 0) as actual_cost"
|
||||
// 当仅按 account_id 聚合时,实际费用使用账号倍率(total_cost * account_rate_multiplier)。
|
||||
if accountID > 0 && userID == 0 && apiKeyID == 0 {
|
||||
actualCostExpr = "COALESCE(SUM(total_cost * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost"
|
||||
}
|
||||
modelExpr := resolveModelDimensionExpression(source)
|
||||
|
||||
query := fmt.Sprintf(`
|
||||
SELECT
|
||||
model,
|
||||
%s as model,
|
||||
COUNT(*) as requests,
|
||||
COALESCE(SUM(input_tokens), 0) as input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) as output_tokens,
|
||||
@@ -2883,7 +2905,7 @@ func (r *usageLogRepository) GetModelStatsWithFilters(ctx context.Context, start
|
||||
%s
|
||||
FROM usage_logs
|
||||
WHERE created_at >= $1 AND created_at < $2
|
||||
`, actualCostExpr)
|
||||
`, modelExpr, actualCostExpr)
|
||||
|
||||
args := []any{startTime, endTime}
|
||||
if userID > 0 {
|
||||
@@ -2907,7 +2929,7 @@ func (r *usageLogRepository) GetModelStatsWithFilters(ctx context.Context, start
|
||||
query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1)
|
||||
args = append(args, int16(*billingType))
|
||||
}
|
||||
query += " GROUP BY model ORDER BY total_tokens DESC"
|
||||
query += fmt.Sprintf(" GROUP BY %s ORDER BY total_tokens DESC", modelExpr)
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
@@ -3000,6 +3022,132 @@ func (r *usageLogRepository) GetGroupStatsWithFilters(ctx context.Context, start
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// GetUserBreakdownStats returns per-user usage breakdown within a specific dimension.
|
||||
func (r *usageLogRepository) GetUserBreakdownStats(ctx context.Context, startTime, endTime time.Time, dim usagestats.UserBreakdownDimension, limit int) (results []usagestats.UserBreakdownItem, err error) {
|
||||
query := `
|
||||
SELECT
|
||||
COALESCE(ul.user_id, 0) as user_id,
|
||||
COALESCE(u.email, '') as email,
|
||||
COUNT(*) as requests,
|
||||
COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens), 0) as total_tokens,
|
||||
COALESCE(SUM(ul.total_cost), 0) as cost,
|
||||
COALESCE(SUM(ul.actual_cost), 0) as actual_cost
|
||||
FROM usage_logs ul
|
||||
LEFT JOIN users u ON u.id = ul.user_id
|
||||
WHERE ul.created_at >= $1 AND ul.created_at < $2
|
||||
`
|
||||
args := []any{startTime, endTime}
|
||||
|
||||
if dim.GroupID > 0 {
|
||||
query += fmt.Sprintf(" AND ul.group_id = $%d", len(args)+1)
|
||||
args = append(args, dim.GroupID)
|
||||
}
|
||||
if dim.Model != "" {
|
||||
query += fmt.Sprintf(" AND %s = $%d", resolveModelDimensionExpression(dim.ModelType), len(args)+1)
|
||||
args = append(args, dim.Model)
|
||||
}
|
||||
if dim.Endpoint != "" {
|
||||
col := resolveEndpointColumn(dim.EndpointType)
|
||||
query += fmt.Sprintf(" AND %s = $%d", col, len(args)+1)
|
||||
args = append(args, dim.Endpoint)
|
||||
}
|
||||
|
||||
query += " GROUP BY ul.user_id, u.email ORDER BY actual_cost DESC"
|
||||
if limit > 0 {
|
||||
query += fmt.Sprintf(" LIMIT %d", limit)
|
||||
}
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() {
|
||||
if closeErr := rows.Close(); closeErr != nil && err == nil {
|
||||
err = closeErr
|
||||
results = nil
|
||||
}
|
||||
}()
|
||||
|
||||
results = make([]usagestats.UserBreakdownItem, 0)
|
||||
for rows.Next() {
|
||||
var row usagestats.UserBreakdownItem
|
||||
if err := rows.Scan(
|
||||
&row.UserID,
|
||||
&row.Email,
|
||||
&row.Requests,
|
||||
&row.TotalTokens,
|
||||
&row.Cost,
|
||||
&row.ActualCost,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results = append(results, row)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// GetAllGroupUsageSummary returns today's and cumulative actual_cost for every group.
|
||||
// todayStart is the start-of-day in the caller's timezone (UTC-based).
|
||||
// TODO(perf): This query scans ALL usage_logs rows for total_cost aggregation.
|
||||
// When usage_logs exceeds ~1M rows, consider adding a short-lived cache (30s)
|
||||
// or a materialized view / pre-aggregation table for cumulative costs.
|
||||
func (r *usageLogRepository) GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) {
|
||||
query := `
|
||||
SELECT
|
||||
g.id AS group_id,
|
||||
COALESCE(SUM(ul.actual_cost), 0) AS total_cost,
|
||||
COALESCE(SUM(CASE WHEN ul.created_at >= $1 THEN ul.actual_cost ELSE 0 END), 0) AS today_cost
|
||||
FROM groups g
|
||||
LEFT JOIN usage_logs ul ON ul.group_id = g.id
|
||||
GROUP BY g.id
|
||||
`
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, todayStart)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
var results []usagestats.GroupUsageSummary
|
||||
for rows.Next() {
|
||||
var row usagestats.GroupUsageSummary
|
||||
if err := rows.Scan(&row.GroupID, &row.TotalCost, &row.TodayCost); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results = append(results, row)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// resolveModelDimensionExpression maps model source type to a safe SQL expression.
|
||||
func resolveModelDimensionExpression(modelType string) string {
|
||||
switch usagestats.NormalizeModelSource(modelType) {
|
||||
case usagestats.ModelSourceUpstream:
|
||||
return "COALESCE(NULLIF(TRIM(upstream_model), ''), model)"
|
||||
case usagestats.ModelSourceMapping:
|
||||
return "(model || ' -> ' || COALESCE(NULLIF(TRIM(upstream_model), ''), model))"
|
||||
default:
|
||||
return "model"
|
||||
}
|
||||
}
|
||||
|
||||
// resolveEndpointColumn maps endpoint type to the corresponding DB column name.
|
||||
func resolveEndpointColumn(endpointType string) string {
|
||||
switch endpointType {
|
||||
case "upstream":
|
||||
return "ul.upstream_endpoint"
|
||||
case "path":
|
||||
return "ul.inbound_endpoint || ' -> ' || ul.upstream_endpoint"
|
||||
default:
|
||||
return "ul.inbound_endpoint"
|
||||
}
|
||||
}
|
||||
|
||||
// GetGlobalStats gets usage statistics for all users within a time range
|
||||
func (r *usageLogRepository) GetGlobalStats(ctx context.Context, startTime, endTime time.Time) (*UsageStats, error) {
|
||||
query := `
|
||||
@@ -3740,6 +3888,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
|
||||
accountID int64
|
||||
requestID sql.NullString
|
||||
model string
|
||||
upstreamModel sql.NullString
|
||||
groupID sql.NullInt64
|
||||
subscriptionID sql.NullInt64
|
||||
inputTokens int
|
||||
@@ -3782,6 +3931,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
|
||||
&accountID,
|
||||
&requestID,
|
||||
&model,
|
||||
&upstreamModel,
|
||||
&groupID,
|
||||
&subscriptionID,
|
||||
&inputTokens,
|
||||
@@ -3894,6 +4044,9 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
|
||||
if upstreamEndpoint.Valid {
|
||||
log.UpstreamEndpoint = &upstreamEndpoint.String
|
||||
}
|
||||
if upstreamModel.Valid {
|
||||
log.UpstreamModel = &upstreamModel.String
|
||||
}
|
||||
|
||||
return log, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestResolveEndpointColumn(t *testing.T) {
|
||||
tests := []struct {
|
||||
endpointType string
|
||||
want string
|
||||
}{
|
||||
{"inbound", "ul.inbound_endpoint"},
|
||||
{"upstream", "ul.upstream_endpoint"},
|
||||
{"path", "ul.inbound_endpoint || ' -> ' || ul.upstream_endpoint"},
|
||||
{"", "ul.inbound_endpoint"}, // default
|
||||
{"unknown", "ul.inbound_endpoint"}, // fallback
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.endpointType, func(t *testing.T) {
|
||||
got := resolveEndpointColumn(tc.endpointType)
|
||||
require.Equal(t, tc.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveModelDimensionExpression(t *testing.T) {
|
||||
tests := []struct {
|
||||
modelType string
|
||||
want string
|
||||
}{
|
||||
{usagestats.ModelSourceRequested, "model"},
|
||||
{usagestats.ModelSourceUpstream, "COALESCE(NULLIF(TRIM(upstream_model), ''), model)"},
|
||||
{usagestats.ModelSourceMapping, "(model || ' -> ' || COALESCE(NULLIF(TRIM(upstream_model), ''), model))"},
|
||||
{"", "model"},
|
||||
{"invalid", "model"},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.modelType, func(t *testing.T) {
|
||||
got := resolveModelDimensionExpression(tc.modelType)
|
||||
require.Equal(t, tc.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -44,6 +44,7 @@ func TestUsageLogRepositoryCreateSyncRequestTypeAndLegacyFields(t *testing.T) {
|
||||
log.AccountID,
|
||||
log.RequestID,
|
||||
log.Model,
|
||||
sqlmock.AnyArg(), // upstream_model
|
||||
sqlmock.AnyArg(), // group_id
|
||||
sqlmock.AnyArg(), // subscription_id
|
||||
log.InputTokens,
|
||||
@@ -116,6 +117,7 @@ func TestUsageLogRepositoryCreate_PersistsServiceTier(t *testing.T) {
|
||||
log.Model,
|
||||
sqlmock.AnyArg(),
|
||||
sqlmock.AnyArg(),
|
||||
sqlmock.AnyArg(),
|
||||
log.InputTokens,
|
||||
log.OutputTokens,
|
||||
log.CacheCreationTokens,
|
||||
@@ -353,6 +355,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
int64(30), // account_id
|
||||
sql.NullString{Valid: true, String: "req-1"},
|
||||
"gpt-5", // model
|
||||
sql.NullString{}, // upstream_model
|
||||
sql.NullInt64{}, // group_id
|
||||
sql.NullInt64{}, // subscription_id
|
||||
1, // input_tokens
|
||||
@@ -404,6 +407,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
int64(31),
|
||||
sql.NullString{Valid: true, String: "req-2"},
|
||||
"gpt-5",
|
||||
sql.NullString{},
|
||||
sql.NullInt64{},
|
||||
sql.NullInt64{},
|
||||
1, 2, 3, 4, 5, 6,
|
||||
@@ -445,6 +449,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
int64(32),
|
||||
sql.NullString{Valid: true, String: "req-3"},
|
||||
"gpt-5.4",
|
||||
sql.NullString{},
|
||||
sql.NullInt64{},
|
||||
sql.NullInt64{},
|
||||
1, 2, 3, 4, 5, 6,
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/ent/apikey"
|
||||
dbgroup "github.com/Wei-Shaw/sub2api/ent/group"
|
||||
dbuser "github.com/Wei-Shaw/sub2api/ent/user"
|
||||
"github.com/Wei-Shaw/sub2api/ent/userallowedgroup"
|
||||
"github.com/Wei-Shaw/sub2api/ent/usersubscription"
|
||||
@@ -200,6 +201,12 @@ func (r *userRepository) ListWithFilters(ctx context.Context, params pagination.
|
||||
)
|
||||
}
|
||||
|
||||
if filters.GroupName != "" {
|
||||
q = q.Where(dbuser.HasAllowedGroupsWith(
|
||||
dbgroup.NameContainsFold(filters.GroupName),
|
||||
))
|
||||
}
|
||||
|
||||
// If attribute filters are specified, we need to filter by user IDs first
|
||||
var allowedUserIDs []int64
|
||||
if len(filters.Attributes) > 0 {
|
||||
@@ -453,6 +460,15 @@ func (r *userRepository) RemoveGroupFromAllowedGroups(ctx context.Context, group
|
||||
return int64(affected), nil
|
||||
}
|
||||
|
||||
// RemoveGroupFromUserAllowedGroups 移除单个用户的指定分组权限
|
||||
func (r *userRepository) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
client := clientFromContext(ctx, r.client)
|
||||
_, err := client.UserAllowedGroup.Delete().
|
||||
Where(userallowedgroup.UserIDEQ(userID), userallowedgroup.GroupIDEQ(groupID)).
|
||||
Exec(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *userRepository) GetFirstAdmin(ctx context.Context) (*service.User, error) {
|
||||
m, err := r.client.User.Query().
|
||||
Where(
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"time"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/ent/group"
|
||||
"github.com/Wei-Shaw/sub2api/ent/usersubscription"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -190,7 +191,7 @@ func (r *userSubscriptionRepository) ListByGroupID(ctx context.Context, groupID
|
||||
return userSubscriptionEntitiesToService(subs), paginationResultFromTotal(int64(total), params), nil
|
||||
}
|
||||
|
||||
func (r *userSubscriptionRepository) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
func (r *userSubscriptionRepository) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, platform, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
client := clientFromContext(ctx, r.client)
|
||||
q := client.UserSubscription.Query()
|
||||
if userID != nil {
|
||||
@@ -199,6 +200,9 @@ func (r *userSubscriptionRepository) List(ctx context.Context, params pagination
|
||||
if groupID != nil {
|
||||
q = q.Where(usersubscription.GroupIDEQ(*groupID))
|
||||
}
|
||||
if platform != "" {
|
||||
q = q.Where(usersubscription.HasGroupWith(group.PlatformEQ(platform)))
|
||||
}
|
||||
|
||||
// Status filtering with real-time expiration check
|
||||
now := time.Now()
|
||||
|
||||
@@ -271,7 +271,7 @@ func (s *UserSubscriptionRepoSuite) TestList_NoFilters() {
|
||||
group := s.mustCreateGroup("g-list")
|
||||
s.mustCreateSubscription(user.ID, group.ID, nil)
|
||||
|
||||
subs, page, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, nil, nil, "", "", "")
|
||||
subs, page, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, nil, nil, "", "", "", "")
|
||||
s.Require().NoError(err, "List")
|
||||
s.Require().Len(subs, 1)
|
||||
s.Require().Equal(int64(1), page.Total)
|
||||
@@ -285,7 +285,7 @@ func (s *UserSubscriptionRepoSuite) TestList_FilterByUserID() {
|
||||
s.mustCreateSubscription(user1.ID, group.ID, nil)
|
||||
s.mustCreateSubscription(user2.ID, group.ID, nil)
|
||||
|
||||
subs, _, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, &user1.ID, nil, "", "", "")
|
||||
subs, _, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, &user1.ID, nil, "", "", "", "")
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(subs, 1)
|
||||
s.Require().Equal(user1.ID, subs[0].UserID)
|
||||
@@ -299,7 +299,7 @@ func (s *UserSubscriptionRepoSuite) TestList_FilterByGroupID() {
|
||||
s.mustCreateSubscription(user.ID, g1.ID, nil)
|
||||
s.mustCreateSubscription(user.ID, g2.ID, nil)
|
||||
|
||||
subs, _, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, nil, &g1.ID, "", "", "")
|
||||
subs, _, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, nil, &g1.ID, "", "", "", "")
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(subs, 1)
|
||||
s.Require().Equal(g1.ID, subs[0].GroupID)
|
||||
@@ -320,7 +320,7 @@ func (s *UserSubscriptionRepoSuite) TestList_FilterByStatus() {
|
||||
c.SetExpiresAt(time.Now().Add(-24 * time.Hour))
|
||||
})
|
||||
|
||||
subs, _, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, nil, nil, service.SubscriptionStatusExpired, "", "")
|
||||
subs, _, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, nil, nil, service.SubscriptionStatusExpired, "", "", "")
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(subs, 1)
|
||||
s.Require().Equal(service.SubscriptionStatusExpired, subs[0].Status)
|
||||
|
||||
@@ -536,6 +536,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"purchase_subscription_enabled": false,
|
||||
"purchase_subscription_url": "",
|
||||
"min_claude_code_version": "",
|
||||
"max_claude_code_version": "",
|
||||
"allow_ungrouped_key_scheduling": false,
|
||||
"backend_mode_enabled": false,
|
||||
"custom_menu_items": []
|
||||
@@ -806,6 +807,10 @@ func (r *stubUserRepo) RemoveGroupFromAllowedGroups(ctx context.Context, groupID
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserRepo) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserRepo) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
@@ -923,8 +928,8 @@ func (stubGroupRepo) ExistsByName(ctx context.Context, name string) (bool, error
|
||||
return false, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (stubGroupRepo) GetAccountCount(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
func (stubGroupRepo) GetAccountCount(ctx context.Context, groupID int64) (int64, int64, error) {
|
||||
return 0, 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (stubGroupRepo) DeleteAccountGroupsByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
@@ -1288,7 +1293,7 @@ func (r *stubUserSubscriptionRepo) ListActiveByUserID(ctx context.Context, userI
|
||||
func (stubUserSubscriptionRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
func (stubUserSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
func (stubUserSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, platform, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
func (stubUserSubscriptionRepo) ExistsByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
|
||||
@@ -1508,6 +1513,22 @@ func (r *stubApiKeyRepo) ClearGroupIDByGroupID(ctx context.Context, groupID int6
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
|
||||
var updated int64
|
||||
for id, key := range r.byID {
|
||||
if key.UserID != userID || key.GroupID == nil || *key.GroupID != oldGroupID {
|
||||
continue
|
||||
}
|
||||
clone := *key
|
||||
gid := newGroupID
|
||||
clone.GroupID = &gid
|
||||
r.byID[id] = &clone
|
||||
r.byKey[clone.Key] = &clone
|
||||
updated++
|
||||
}
|
||||
return updated, nil
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
@@ -1636,6 +1657,10 @@ func (r *stubUsageLogRepo) GetGroupStatsWithFilters(ctx context.Context, startTi
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUsageLogRepo) GetUserBreakdownStats(ctx context.Context, startTime, endTime time.Time, dim usagestats.UserBreakdownDimension, limit int) ([]usagestats.UserBreakdownItem, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUsageLogRepo) GetAPIKeyUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.APIKeyUsageTrendPoint, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
@@ -1781,6 +1806,9 @@ func (r *stubUsageLogRepo) GetAccountUsageStats(ctx context.Context, accountID i
|
||||
func (r *stubUsageLogRepo) GetStatsWithFilters(ctx context.Context, filters usagestats.UsageLogFilters) (*usagestats.UsageStats, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (r *stubUsageLogRepo) GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
type stubSettingRepo struct {
|
||||
all map[string]string
|
||||
|
||||
@@ -181,6 +181,10 @@ func (s *stubUserRepo) RemoveGroupFromAllowedGroups(ctx context.Context, groupID
|
||||
panic("unexpected RemoveGroupFromAllowedGroups call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
panic("unexpected RemoveGroupFromUserAllowedGroups call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
panic("unexpected AddGroupToAllowedGroups call")
|
||||
}
|
||||
|
||||
@@ -104,6 +104,9 @@ func (f fakeAPIKeyRepo) ResetRateLimitWindows(ctx context.Context, id int64) err
|
||||
func (f fakeAPIKeyRepo) GetRateLimitData(ctx context.Context, id int64) (*service.APIKeyRateLimitData, error) {
|
||||
return &service.APIKeyRateLimitData{}, nil
|
||||
}
|
||||
func (f fakeAPIKeyRepo) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (f fakeGoogleSubscriptionRepo) Create(ctx context.Context, sub *service.UserSubscription) error {
|
||||
return errors.New("not implemented")
|
||||
@@ -135,7 +138,7 @@ func (f fakeGoogleSubscriptionRepo) ListActiveByUserID(ctx context.Context, user
|
||||
func (f fakeGoogleSubscriptionRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
func (f fakeGoogleSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, platform, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ExistsByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
|
||||
|
||||
@@ -565,6 +565,10 @@ func (r *stubApiKeyRepo) ClearGroupIDByGroupID(ctx context.Context, groupID int6
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID 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")
|
||||
}
|
||||
@@ -646,7 +650,7 @@ func (r *stubUserSubscriptionRepo) ListByGroupID(ctx context.Context, groupID in
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
func (r *stubUserSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, platform, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
|
||||
@@ -198,6 +198,7 @@ func registerDashboardRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
dashboard.GET("/users-ranking", h.Admin.Dashboard.GetUserSpendingRanking)
|
||||
dashboard.POST("/users-usage", h.Admin.Dashboard.GetBatchUsersUsage)
|
||||
dashboard.POST("/api-keys-usage", h.Admin.Dashboard.GetBatchAPIKeysUsage)
|
||||
dashboard.GET("/user-breakdown", h.Admin.Dashboard.GetUserBreakdown)
|
||||
dashboard.POST("/aggregation/backfill", h.Admin.Dashboard.BackfillAggregation)
|
||||
}
|
||||
}
|
||||
@@ -214,6 +215,7 @@ func registerUserManagementRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
users.GET("/:id/api-keys", h.Admin.User.GetUserAPIKeys)
|
||||
users.GET("/:id/usage", h.Admin.User.GetUserUsage)
|
||||
users.GET("/:id/balance-history", h.Admin.User.GetBalanceHistory)
|
||||
users.POST("/:id/replace-group", h.Admin.User.ReplaceGroup)
|
||||
|
||||
// User attribute values
|
||||
users.GET("/:id/attributes", h.Admin.UserAttribute.GetUserAttributes)
|
||||
@@ -226,6 +228,8 @@ func registerGroupRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
{
|
||||
groups.GET("", h.Admin.Group.List)
|
||||
groups.GET("/all", h.Admin.Group.GetAll)
|
||||
groups.GET("/usage-summary", h.Admin.Group.GetUsageSummary)
|
||||
groups.GET("/capacity-summary", h.Admin.Group.GetCapacitySummary)
|
||||
groups.PUT("/sort-order", h.Admin.Group.UpdateSortOrder)
|
||||
groups.GET("/:id", h.Admin.Group.GetByID)
|
||||
groups.POST("", h.Admin.Group.Create)
|
||||
@@ -262,6 +266,7 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
accounts.POST("/:id/clear-rate-limit", h.Admin.Account.ClearRateLimit)
|
||||
accounts.POST("/:id/reset-quota", h.Admin.Account.ResetQuota)
|
||||
accounts.GET("/:id/affinity-clients", h.Admin.Account.GetAffinityClients)
|
||||
accounts.GET("/:id/affinity-details", h.Admin.Account.GetAffinityDetails)
|
||||
accounts.GET("/:id/temp-unschedulable", h.Admin.Account.GetTempUnschedulable)
|
||||
accounts.DELETE("/:id/temp-unschedulable", h.Admin.Account.ClearTempUnschedulable)
|
||||
accounts.POST("/:id/schedulable", h.Admin.Account.SetSchedulable)
|
||||
@@ -400,6 +405,9 @@ func registerSettingsRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
adminSettings.GET("/admin-api-key", h.Admin.Setting.GetAdminAPIKey)
|
||||
adminSettings.POST("/admin-api-key/regenerate", h.Admin.Setting.RegenerateAdminAPIKey)
|
||||
adminSettings.DELETE("/admin-api-key", h.Admin.Setting.DeleteAdminAPIKey)
|
||||
// 529过载冷却配置
|
||||
adminSettings.GET("/overload-cooldown", h.Admin.Setting.GetOverloadCooldownSettings)
|
||||
adminSettings.PUT("/overload-cooldown", h.Admin.Setting.UpdateOverloadCooldownSettings)
|
||||
// 流超时处理配置
|
||||
adminSettings.GET("/stream-timeout", h.Admin.Setting.GetStreamTimeoutSettings)
|
||||
adminSettings.PUT("/stream-timeout", h.Admin.Setting.UpdateStreamTimeoutSettings)
|
||||
|
||||
@@ -1204,16 +1204,21 @@ func (a *Account) IsSessionIDMaskingEnabled() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// IsClientAffinityEnabled 检查是否启用客户端亲和调度
|
||||
// 仅适用于 Anthropic 账号(OAuth/SetupToken/APIKey)
|
||||
// 启用后,新会话会优先调度到之前使用过的账号
|
||||
func (a *Account) IsClientAffinityEnabled() bool {
|
||||
// IsAffinityEnabled 检查是否启用亲和调度(统一入口)
|
||||
// 仅适用于 Anthropic 账号,同时检查新字段 affinity_enabled 和旧字段 client_affinity_enabled(向后兼容)
|
||||
func (a *Account) IsAffinityEnabled() bool {
|
||||
if a.Platform != PlatformAnthropic {
|
||||
return false
|
||||
}
|
||||
if a.Extra == nil {
|
||||
return false
|
||||
}
|
||||
if v, ok := a.Extra["affinity_enabled"]; ok {
|
||||
if enabled, ok := v.(bool); ok {
|
||||
return enabled
|
||||
}
|
||||
return false
|
||||
}
|
||||
if v, ok := a.Extra["client_affinity_enabled"]; ok {
|
||||
if enabled, ok := v.(bool); ok {
|
||||
return enabled
|
||||
@@ -1222,6 +1227,11 @@ func (a *Account) IsClientAffinityEnabled() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// IsClientAffinityEnabled 向后兼容别名,内部调用 IsAffinityEnabled
|
||||
func (a *Account) IsClientAffinityEnabled() bool {
|
||||
return a.IsAffinityEnabled()
|
||||
}
|
||||
|
||||
// AffinityZone 表示账号的客户端亲和分区
|
||||
type AffinityZone int
|
||||
|
||||
@@ -1288,6 +1298,150 @@ func (a *Account) GetAffinityZone(clientCount int64) AffinityZone {
|
||||
return AffinityZoneRed
|
||||
}
|
||||
|
||||
// IsAffinityAllowSwitch 检查亲和账号是否允许在无绿区时切换到其他账号
|
||||
// 默认 true(允许切换),设为 false 时即使全部红区也不切换
|
||||
func (a *Account) IsAffinityAllowSwitch() bool {
|
||||
if a.Extra == nil {
|
||||
return true
|
||||
}
|
||||
if v, ok := a.Extra["affinity_allow_switch"]; ok {
|
||||
if allow, ok := v.(bool); ok {
|
||||
return allow
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// GetPinnedUsers 获取指定亲和用户列表
|
||||
func (a *Account) GetPinnedUsers() []int64 {
|
||||
if a.Extra == nil {
|
||||
return nil
|
||||
}
|
||||
v, ok := a.Extra["pinned_users"]
|
||||
if !ok || v == nil {
|
||||
return nil
|
||||
}
|
||||
arr, ok := v.([]any)
|
||||
if !ok || len(arr) == 0 {
|
||||
return nil
|
||||
}
|
||||
result := make([]int64, 0, len(arr))
|
||||
for _, item := range arr {
|
||||
switch id := item.(type) {
|
||||
case float64:
|
||||
result = append(result, int64(id))
|
||||
case int64:
|
||||
result = append(result, id)
|
||||
case int:
|
||||
result = append(result, int64(id))
|
||||
case json.Number:
|
||||
if i, err := id.Int64(); err == nil {
|
||||
result = append(result, i)
|
||||
}
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// IsPinnedUser 检查 userID 是否在指定亲和用户列表中
|
||||
func (a *Account) IsPinnedUser(userID int64) bool {
|
||||
for _, id := range a.GetPinnedUsers() {
|
||||
if id == userID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// GetAffinityUserBase 获取用户维度亲和基础限制(绿区上限),0 表示未配置
|
||||
func (a *Account) GetAffinityUserBase() int {
|
||||
if a.Extra == nil {
|
||||
return 0
|
||||
}
|
||||
if v, ok := a.Extra["affinity_user_base"]; ok {
|
||||
return parseExtraInt(v)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// GetAffinityUserBuffer 获取用户维度亲和缓冲区大小(黄区范围)
|
||||
// 返回 (value, configured),语义与 GetAffinityBuffer 相同
|
||||
func (a *Account) GetAffinityUserBuffer() (int, bool) {
|
||||
if a.Extra == nil {
|
||||
return 0, false
|
||||
}
|
||||
v, ok := a.Extra["affinity_user_buffer"]
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
if v == nil {
|
||||
return 0, false
|
||||
}
|
||||
return parseExtraInt(v), true
|
||||
}
|
||||
|
||||
// GetPerUserClientLimit 获取每用户客户端限制,0 表示不限制
|
||||
func (a *Account) GetPerUserClientLimit() int {
|
||||
if a.Extra == nil {
|
||||
return 0
|
||||
}
|
||||
if v, ok := a.Extra["per_user_client_limit"]; ok {
|
||||
return parseExtraInt(v)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// getAffinityZoneForDim 根据单一维度的计数和三区参数计算亲和分区
|
||||
func getAffinityZoneForDim(count int64, base int, buffer int, bufferConfigured bool) AffinityZone {
|
||||
if base <= 0 {
|
||||
return AffinityZoneGreen
|
||||
}
|
||||
if count <= int64(base) {
|
||||
return AffinityZoneGreen
|
||||
}
|
||||
if !bufferConfigured {
|
||||
return AffinityZoneYellow
|
||||
}
|
||||
if buffer == 0 {
|
||||
return AffinityZoneRed
|
||||
}
|
||||
if count <= int64(base+buffer) {
|
||||
return AffinityZoneYellow
|
||||
}
|
||||
return AffinityZoneRed
|
||||
}
|
||||
|
||||
// GetMultiDimAffinityZone 根据多维度计数计算亲和分区,取所有维度中最严格的区域。
|
||||
// 未开启亲和的账号永远返回绿区。
|
||||
func (a *Account) GetMultiDimAffinityZone(userCount, clientCount, perUserCount int64) AffinityZone {
|
||||
if !a.IsAffinityEnabled() {
|
||||
return AffinityZoneGreen
|
||||
}
|
||||
worst := AffinityZoneGreen
|
||||
|
||||
// 客户端维度
|
||||
clientBase := a.GetAffinityBase()
|
||||
clientBuffer, clientBufCfg := a.GetAffinityBuffer()
|
||||
if z := getAffinityZoneForDim(clientCount, clientBase, clientBuffer, clientBufCfg); z > worst {
|
||||
worst = z
|
||||
}
|
||||
|
||||
// 用户维度
|
||||
userBase := a.GetAffinityUserBase()
|
||||
userBuffer, userBufCfg := a.GetAffinityUserBuffer()
|
||||
if z := getAffinityZoneForDim(userCount, userBase, userBuffer, userBufCfg); z > worst {
|
||||
worst = z
|
||||
}
|
||||
|
||||
// 每用户客户端维度
|
||||
perUserLimit := a.GetPerUserClientLimit()
|
||||
if perUserLimit > 0 && perUserCount > int64(perUserLimit) {
|
||||
worst = AffinityZoneRed
|
||||
}
|
||||
|
||||
return worst
|
||||
}
|
||||
|
||||
// IsCacheTTLOverrideEnabled 检查是否启用缓存 TTL 强制替换
|
||||
// 仅适用于 Anthropic OAuth/SetupToken 类型账号
|
||||
// 启用后将所有 cache creation tokens 归入指定的 TTL 类型(5m 或 1h)
|
||||
|
||||
@@ -14,6 +14,8 @@ var (
|
||||
ErrAccountNilInput = infraerrors.BadRequest("ACCOUNT_NIL_INPUT", "account input cannot be nil")
|
||||
)
|
||||
|
||||
const AccountListGroupUngrouped int64 = -1
|
||||
|
||||
type AccountRepository interface {
|
||||
Create(ctx context.Context, account *Account) error
|
||||
GetByID(ctx context.Context, id int64) (*Account, error)
|
||||
|
||||
@@ -113,15 +113,18 @@ func (s *AccountTestService) validateUpstreamBaseURL(raw string) (string, error)
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
// generateSessionString generates a Claude Code style session string
|
||||
// generateSessionString generates a Claude Code style session string.
|
||||
// The output format is determined by the UA version in claude.DefaultHeaders,
|
||||
// ensuring consistency between the user_id format and the UA sent to upstream.
|
||||
func generateSessionString() (string, error) {
|
||||
bytes := make([]byte, 32)
|
||||
if _, err := rand.Read(bytes); err != nil {
|
||||
b := make([]byte, 32)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", err
|
||||
}
|
||||
hex64 := hex.EncodeToString(bytes)
|
||||
hex64 := hex.EncodeToString(b)
|
||||
sessionUUID := uuid.New().String()
|
||||
return fmt.Sprintf("user_%s_account__session_%s", hex64, sessionUUID), nil
|
||||
uaVersion := ExtractCLIVersion(claude.DefaultHeaders["User-Agent"])
|
||||
return FormatMetadataUserID(hex64, "", sessionUUID, uaVersion), nil
|
||||
}
|
||||
|
||||
// createTestPayload creates a Claude Code style test request payload
|
||||
@@ -305,7 +308,14 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
|
||||
errMsg := fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body))
|
||||
|
||||
// 403 表示账号被上游封禁,标记为 error 状态
|
||||
if resp.StatusCode == http.StatusForbidden {
|
||||
_ = s.accountRepo.SetError(ctx, account.ID, errMsg)
|
||||
}
|
||||
|
||||
return s.sendErrorAndEnd(c, errMsg)
|
||||
}
|
||||
|
||||
// Process SSE stream
|
||||
|
||||
@@ -48,6 +48,8 @@ type UsageLogRepository interface {
|
||||
GetEndpointStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) ([]usagestats.EndpointStat, error)
|
||||
GetUpstreamEndpointStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) ([]usagestats.EndpointStat, error)
|
||||
GetGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) ([]usagestats.GroupStat, error)
|
||||
GetUserBreakdownStats(ctx context.Context, startTime, endTime time.Time, dim usagestats.UserBreakdownDimension, limit int) ([]usagestats.UserBreakdownItem, error)
|
||||
GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error)
|
||||
GetAPIKeyUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.APIKeyUsageTrendPoint, error)
|
||||
GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.UserUsageTrendPoint, error)
|
||||
GetUserSpendingRanking(ctx context.Context, startTime, endTime time.Time, limit int) (*usagestats.UserSpendingRankingResponse, error)
|
||||
@@ -175,6 +177,7 @@ type AICredit struct {
|
||||
|
||||
// UsageInfo 账号使用量信息
|
||||
type UsageInfo struct {
|
||||
Source string `json:"source,omitempty"` // "passive" or "active"
|
||||
UpdatedAt *time.Time `json:"updated_at,omitempty"` // 更新时间
|
||||
FiveHour *UsageProgress `json:"five_hour"` // 5小时窗口
|
||||
SevenDay *UsageProgress `json:"seven_day,omitempty"` // 7天窗口
|
||||
@@ -391,6 +394,9 @@ func (s *AccountUsageService) GetUsage(ctx context.Context, accountID int64) (*U
|
||||
// 4. 添加窗口统计(有独立缓存,1 分钟)
|
||||
s.addWindowStats(ctx, account, usage)
|
||||
|
||||
// 5. 将主动查询结果同步到被动缓存,下次 passive 加载即为最新值
|
||||
s.syncActiveToPassive(ctx, account.ID, usage)
|
||||
|
||||
s.tryClearRecoverableAccountError(ctx, account)
|
||||
return usage, nil
|
||||
}
|
||||
@@ -407,6 +413,81 @@ func (s *AccountUsageService) GetUsage(ctx context.Context, accountID int64) (*U
|
||||
return nil, fmt.Errorf("account type %s does not support usage query", account.Type)
|
||||
}
|
||||
|
||||
// GetPassiveUsage 从 Account.Extra 中的被动采样数据构建 UsageInfo,不调用外部 API。
|
||||
// 仅适用于 Anthropic OAuth / SetupToken 账号。
|
||||
func (s *AccountUsageService) GetPassiveUsage(ctx context.Context, accountID int64) (*UsageInfo, error) {
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get account failed: %w", err)
|
||||
}
|
||||
|
||||
if !account.IsAnthropicOAuthOrSetupToken() {
|
||||
return nil, fmt.Errorf("passive usage only supported for Anthropic OAuth/SetupToken accounts")
|
||||
}
|
||||
|
||||
// 复用 estimateSetupTokenUsage 构建 5h 窗口(OAuth 和 SetupToken 逻辑一致)
|
||||
info := s.estimateSetupTokenUsage(account)
|
||||
info.Source = "passive"
|
||||
|
||||
// 设置采样时间
|
||||
if raw, ok := account.Extra["passive_usage_sampled_at"]; ok {
|
||||
if str, ok := raw.(string); ok {
|
||||
if t, err := time.Parse(time.RFC3339, str); err == nil {
|
||||
info.UpdatedAt = &t
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 构建 7d 窗口(从被动采样数据)
|
||||
util7d := parseExtraFloat64(account.Extra["passive_usage_7d_utilization"])
|
||||
reset7dRaw := parseExtraFloat64(account.Extra["passive_usage_7d_reset"])
|
||||
if util7d > 0 || reset7dRaw > 0 {
|
||||
var resetAt *time.Time
|
||||
var remaining int
|
||||
if reset7dRaw > 0 {
|
||||
t := time.Unix(int64(reset7dRaw), 0)
|
||||
resetAt = &t
|
||||
remaining = int(time.Until(t).Seconds())
|
||||
if remaining < 0 {
|
||||
remaining = 0
|
||||
}
|
||||
}
|
||||
info.SevenDay = &UsageProgress{
|
||||
Utilization: util7d * 100,
|
||||
ResetsAt: resetAt,
|
||||
RemainingSeconds: remaining,
|
||||
}
|
||||
}
|
||||
|
||||
// 添加窗口统计
|
||||
s.addWindowStats(ctx, account, info)
|
||||
|
||||
return info, nil
|
||||
}
|
||||
|
||||
// syncActiveToPassive 将主动查询的最新数据回写到 Extra 被动缓存,
|
||||
// 这样下次被动加载时能看到最新值。
|
||||
func (s *AccountUsageService) syncActiveToPassive(ctx context.Context, accountID int64, usage *UsageInfo) {
|
||||
extraUpdates := make(map[string]any, 4)
|
||||
|
||||
if usage.FiveHour != nil {
|
||||
extraUpdates["session_window_utilization"] = usage.FiveHour.Utilization / 100
|
||||
}
|
||||
if usage.SevenDay != nil {
|
||||
extraUpdates["passive_usage_7d_utilization"] = usage.SevenDay.Utilization / 100
|
||||
if usage.SevenDay.ResetsAt != nil {
|
||||
extraUpdates["passive_usage_7d_reset"] = usage.SevenDay.ResetsAt.Unix()
|
||||
}
|
||||
}
|
||||
|
||||
if len(extraUpdates) > 0 {
|
||||
extraUpdates["passive_usage_sampled_at"] = time.Now().UTC().Format(time.RFC3339)
|
||||
if err := s.accountRepo.UpdateExtra(ctx, accountID, extraUpdates); err != nil {
|
||||
slog.Warn("sync_active_to_passive_failed", "account_id", accountID, "error", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AccountUsageService) getOpenAIUsage(ctx context.Context, account *Account) (*UsageInfo, error) {
|
||||
now := time.Now()
|
||||
usage := &UsageInfo{UpdatedAt: &now}
|
||||
@@ -1042,6 +1123,11 @@ func buildCodexUsageProgressFromExtra(extra map[string]any, window string, now t
|
||||
}
|
||||
}
|
||||
|
||||
// 窗口已过期(resetAt 在 now 之前)→ 额度已重置,归零
|
||||
if progress.ResetsAt != nil && !now.Before(*progress.ResetsAt) {
|
||||
progress.Utilization = 0
|
||||
}
|
||||
|
||||
return progress
|
||||
}
|
||||
|
||||
|
||||
@@ -148,3 +148,54 @@ func TestAccountUsageService_PersistOpenAICodexProbeSnapshotSetsRateLimit(t *tes
|
||||
t.Fatal("waiting for codex probe rate limit persistence timed out")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildCodexUsageProgressFromExtra_ZerosExpiredWindow(t *testing.T) {
|
||||
t.Parallel()
|
||||
now := time.Date(2026, 3, 16, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
t.Run("expired 5h window zeroes utilization", func(t *testing.T) {
|
||||
extra := map[string]any{
|
||||
"codex_5h_used_percent": 42.0,
|
||||
"codex_5h_reset_at": "2026-03-16T10:00:00Z", // 2h ago
|
||||
}
|
||||
progress := buildCodexUsageProgressFromExtra(extra, "5h", now)
|
||||
if progress == nil {
|
||||
t.Fatal("expected non-nil progress")
|
||||
}
|
||||
if progress.Utilization != 0 {
|
||||
t.Fatalf("expected Utilization=0 for expired window, got %v", progress.Utilization)
|
||||
}
|
||||
if progress.RemainingSeconds != 0 {
|
||||
t.Fatalf("expected RemainingSeconds=0, got %v", progress.RemainingSeconds)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("active 5h window keeps utilization", func(t *testing.T) {
|
||||
resetAt := now.Add(2 * time.Hour).Format(time.RFC3339)
|
||||
extra := map[string]any{
|
||||
"codex_5h_used_percent": 42.0,
|
||||
"codex_5h_reset_at": resetAt,
|
||||
}
|
||||
progress := buildCodexUsageProgressFromExtra(extra, "5h", now)
|
||||
if progress == nil {
|
||||
t.Fatal("expected non-nil progress")
|
||||
}
|
||||
if progress.Utilization != 42.0 {
|
||||
t.Fatalf("expected Utilization=42, got %v", progress.Utilization)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("expired 7d window zeroes utilization", func(t *testing.T) {
|
||||
extra := map[string]any{
|
||||
"codex_7d_used_percent": 88.0,
|
||||
"codex_7d_reset_at": "2026-03-15T00:00:00Z", // yesterday
|
||||
}
|
||||
progress := buildCodexUsageProgressFromExtra(extra, "7d", now)
|
||||
if progress == nil {
|
||||
t.Fatal("expected non-nil progress")
|
||||
}
|
||||
if progress.Utilization != 0 {
|
||||
t.Fatalf("expected Utilization=0 for expired 7d window, got %v", progress.Utilization)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -50,6 +50,9 @@ type AdminService interface {
|
||||
// API Key management (admin)
|
||||
AdminUpdateAPIKeyGroupID(ctx context.Context, keyID int64, groupID *int64) (*AdminUpdateAPIKeyGroupIDResult, error)
|
||||
|
||||
// ReplaceUserGroup 替换用户的专属分组:授予新分组权限、迁移 Key、移除旧分组权限
|
||||
ReplaceUserGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (*ReplaceUserGroupResult, error)
|
||||
|
||||
// Account management
|
||||
ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64) ([]Account, int64, error)
|
||||
GetAccount(ctx context.Context, id int64) (*Account, error)
|
||||
@@ -272,6 +275,11 @@ type AdminUpdateAPIKeyGroupIDResult struct {
|
||||
GrantedGroupName string // the group name that was auto-granted
|
||||
}
|
||||
|
||||
// ReplaceUserGroupResult 分组替换操作的结果
|
||||
type ReplaceUserGroupResult struct {
|
||||
MigratedKeys int64 // 迁移的 Key 数量
|
||||
}
|
||||
|
||||
// BulkUpdateAccountsResult is the aggregated response for bulk updates.
|
||||
type BulkUpdateAccountsResult struct {
|
||||
Success int `json:"success"`
|
||||
@@ -1396,6 +1404,71 @@ func (s *adminServiceImpl) AdminUpdateAPIKeyGroupID(ctx context.Context, keyID i
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// ReplaceUserGroup 替换用户的专属分组
|
||||
func (s *adminServiceImpl) ReplaceUserGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (*ReplaceUserGroupResult, error) {
|
||||
if oldGroupID == newGroupID {
|
||||
return nil, infraerrors.BadRequest("SAME_GROUP", "old and new group must be different")
|
||||
}
|
||||
|
||||
// 验证新分组存在且为活跃的专属标准分组
|
||||
newGroup, err := s.groupRepo.GetByID(ctx, newGroupID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if newGroup.Status != StatusActive {
|
||||
return nil, infraerrors.BadRequest("GROUP_NOT_ACTIVE", "target group is not active")
|
||||
}
|
||||
if !newGroup.IsExclusive {
|
||||
return nil, infraerrors.BadRequest("GROUP_NOT_EXCLUSIVE", "target group is not exclusive")
|
||||
}
|
||||
if newGroup.IsSubscriptionType() {
|
||||
return nil, infraerrors.BadRequest("GROUP_IS_SUBSCRIPTION", "subscription groups are not supported for replacement")
|
||||
}
|
||||
|
||||
// 事务保证原子性
|
||||
if s.entClient == nil {
|
||||
return nil, fmt.Errorf("entClient is nil, cannot perform group replacement")
|
||||
}
|
||||
tx, err := s.entClient.Tx(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("begin transaction: %w", err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
opCtx := dbent.NewTxContext(ctx, tx)
|
||||
|
||||
// 1. 授予新分组权限
|
||||
if err := s.userRepo.AddGroupToAllowedGroups(opCtx, userID, newGroupID); err != nil {
|
||||
return nil, fmt.Errorf("add new group to allowed groups: %w", err)
|
||||
}
|
||||
|
||||
// 2. 迁移绑定旧分组的 Key 到新分组
|
||||
migrated, err := s.apiKeyRepo.UpdateGroupIDByUserAndGroup(opCtx, userID, oldGroupID, newGroupID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("migrate api keys: %w", err)
|
||||
}
|
||||
|
||||
// 3. 移除旧分组权限
|
||||
if err := s.userRepo.RemoveGroupFromUserAllowedGroups(opCtx, userID, oldGroupID); err != nil {
|
||||
return nil, fmt.Errorf("remove old group from allowed groups: %w", err)
|
||||
}
|
||||
|
||||
if err := tx.Commit(); err != nil {
|
||||
return nil, fmt.Errorf("commit transaction: %w", err)
|
||||
}
|
||||
|
||||
// 失效该用户所有 Key 的认证缓存
|
||||
if s.authCacheInvalidator != nil {
|
||||
keys, keyErr := s.apiKeyRepo.ListKeysByUserID(ctx, userID)
|
||||
if keyErr == nil {
|
||||
for _, k := range keys {
|
||||
s.authCacheInvalidator.InvalidateAuthCacheByKey(ctx, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return &ReplaceUserGroupResult{MigratedKeys: migrated}, nil
|
||||
}
|
||||
|
||||
// Account management implementations
|
||||
func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64) ([]Account, int64, error) {
|
||||
params := pagination.PaginationParams{Page: page, PageSize: pageSize}
|
||||
|
||||
@@ -65,6 +65,9 @@ func (s *userRepoStubForGroupUpdate) ExistsByEmail(context.Context, string) (boo
|
||||
func (s *userRepoStubForGroupUpdate) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
|
||||
panic("unexpected")
|
||||
}
|
||||
func (s *userRepoStubForGroupUpdate) RemoveGroupFromUserAllowedGroups(context.Context, int64, int64) error {
|
||||
panic("unexpected")
|
||||
}
|
||||
func (s *userRepoStubForGroupUpdate) UpdateTotpSecret(context.Context, int64, *string) error {
|
||||
panic("unexpected")
|
||||
}
|
||||
@@ -128,6 +131,9 @@ func (s *apiKeyRepoStubForGroupUpdate) SearchAPIKeys(context.Context, int64, str
|
||||
func (s *apiKeyRepoStubForGroupUpdate) ClearGroupIDByGroupID(context.Context, int64) (int64, error) {
|
||||
panic("unexpected")
|
||||
}
|
||||
func (s *apiKeyRepoStubForGroupUpdate) UpdateGroupIDByUserAndGroup(context.Context, int64, int64, int64) (int64, error) {
|
||||
panic("unexpected")
|
||||
}
|
||||
func (s *apiKeyRepoStubForGroupUpdate) CountByGroupID(context.Context, int64) (int64, error) {
|
||||
panic("unexpected")
|
||||
}
|
||||
@@ -194,7 +200,7 @@ func (s *groupRepoStubForGroupUpdate) ListActiveByPlatform(context.Context, stri
|
||||
func (s *groupRepoStubForGroupUpdate) ExistsByName(context.Context, string) (bool, error) {
|
||||
panic("unexpected")
|
||||
}
|
||||
func (s *groupRepoStubForGroupUpdate) GetAccountCount(context.Context, int64) (int64, error) {
|
||||
func (s *groupRepoStubForGroupUpdate) GetAccountCount(context.Context, int64) (int64, int64, error) {
|
||||
panic("unexpected")
|
||||
}
|
||||
func (s *groupRepoStubForGroupUpdate) DeleteAccountGroupsByGroupID(context.Context, int64) (int64, error) {
|
||||
|
||||
@@ -93,6 +93,10 @@ func (s *userRepoStub) RemoveGroupFromAllowedGroups(ctx context.Context, groupID
|
||||
panic("unexpected RemoveGroupFromAllowedGroups call")
|
||||
}
|
||||
|
||||
func (s *userRepoStub) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
panic("unexpected RemoveGroupFromUserAllowedGroups call")
|
||||
}
|
||||
|
||||
func (s *userRepoStub) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
panic("unexpected AddGroupToAllowedGroups call")
|
||||
}
|
||||
@@ -160,7 +164,7 @@ func (s *groupRepoStub) ExistsByName(ctx context.Context, name string) (bool, er
|
||||
panic("unexpected ExistsByName call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStub) GetAccountCount(ctx context.Context, groupID int64) (int64, error) {
|
||||
func (s *groupRepoStub) GetAccountCount(ctx context.Context, groupID int64) (int64, int64, error) {
|
||||
panic("unexpected GetAccountCount call")
|
||||
}
|
||||
|
||||
|
||||
@@ -100,7 +100,7 @@ func (s *groupRepoStubForAdmin) ExistsByName(_ context.Context, _ string) (bool,
|
||||
panic("unexpected ExistsByName call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForAdmin) GetAccountCount(_ context.Context, _ int64) (int64, error) {
|
||||
func (s *groupRepoStubForAdmin) GetAccountCount(_ context.Context, _ int64) (int64, int64, error) {
|
||||
panic("unexpected GetAccountCount call")
|
||||
}
|
||||
|
||||
@@ -383,7 +383,7 @@ func (s *groupRepoStubForFallbackCycle) ExistsByName(_ context.Context, _ string
|
||||
panic("unexpected ExistsByName call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForFallbackCycle) GetAccountCount(_ context.Context, _ int64) (int64, error) {
|
||||
func (s *groupRepoStubForFallbackCycle) GetAccountCount(_ context.Context, _ int64) (int64, int64, error) {
|
||||
panic("unexpected GetAccountCount call")
|
||||
}
|
||||
|
||||
@@ -458,7 +458,7 @@ func (s *groupRepoStubForInvalidRequestFallback) ExistsByName(_ context.Context,
|
||||
panic("unexpected ExistsByName call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForInvalidRequestFallback) GetAccountCount(_ context.Context, _ int64) (int64, error) {
|
||||
func (s *groupRepoStubForInvalidRequestFallback) GetAccountCount(_ context.Context, _ int64) (int64, int64, error) {
|
||||
panic("unexpected GetAccountCount call")
|
||||
}
|
||||
|
||||
|
||||
@@ -45,6 +45,7 @@ var (
|
||||
"minimumcreditamountforusage",
|
||||
"minimum credit amount for usage",
|
||||
"minimum credit",
|
||||
"resource has been exhausted",
|
||||
}
|
||||
)
|
||||
|
||||
@@ -147,9 +148,9 @@ func shouldMarkCreditsExhausted(resp *http.Response, respBody []byte, reqErr err
|
||||
if resp.StatusCode >= 500 || resp.StatusCode == http.StatusRequestTimeout {
|
||||
return false
|
||||
}
|
||||
if isURLLevelRateLimit(respBody) {
|
||||
return false
|
||||
}
|
||||
// 注意:不再检查 isURLLevelRateLimit。此函数仅在积分重试失败后调用,
|
||||
// 如果注入 enabledCreditTypes 后仍返回 "Resource has been exhausted",
|
||||
// 说明积分也已耗尽,应该标记。clearCreditsExhausted 会在后续成功时自动清除。
|
||||
if info := parseAntigravitySmartRetryInfo(respBody); info != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -406,10 +406,16 @@ func TestShouldMarkCreditsExhausted(t *testing.T) {
|
||||
require.False(t, shouldMarkCreditsExhausted(resp, []byte(`{"error":"Insufficient credits"}`), nil))
|
||||
})
|
||||
|
||||
t.Run("URL 级限流不标记", func(t *testing.T) {
|
||||
t.Run("Resource has been exhausted 应标记为积分耗尽", func(t *testing.T) {
|
||||
resp := &http.Response{StatusCode: http.StatusTooManyRequests}
|
||||
body := []byte(`{"error":{"message":"Resource has been exhausted"}}`)
|
||||
require.False(t, shouldMarkCreditsExhausted(resp, body, nil))
|
||||
require.True(t, shouldMarkCreditsExhausted(resp, body, nil))
|
||||
})
|
||||
|
||||
t.Run("Resource has been exhausted (check quota) 完整格式应标记", func(t *testing.T) {
|
||||
resp := &http.Response{StatusCode: http.StatusTooManyRequests}
|
||||
body := []byte(`{"error":{"code":429,"message":"Resource has been exhausted (e.g. check quota).","status":"RESOURCE_EXHAUSTED"}}`)
|
||||
require.True(t, shouldMarkCreditsExhausted(resp, body, nil))
|
||||
})
|
||||
|
||||
t.Run("结构化限流不标记", func(t *testing.T) {
|
||||
|
||||
@@ -930,7 +930,7 @@ func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParam
|
||||
case ErrorPolicyTempUnscheduled:
|
||||
slog.Info("temp_unschedulable_matched",
|
||||
"prefix", p.prefix, "status_code", statusCode, "account_id", p.account.ID)
|
||||
return true, statusCode, &AntigravityAccountSwitchError{OriginalAccountID: p.account.ID, IsStickySession: p.isStickySession}
|
||||
return true, statusCode, &AntigravityAccountSwitchError{OriginalAccountID: p.account.ID, RateLimitedModel: p.requestedModel, IsStickySession: p.isStickySession}
|
||||
}
|
||||
return false, statusCode, nil
|
||||
}
|
||||
@@ -1001,8 +1001,9 @@ type TestConnectionResult struct {
|
||||
MappedModel string // 实际使用的模型
|
||||
}
|
||||
|
||||
// TestConnection 测试 Antigravity 账号连接(非流式,无重试、无计费)
|
||||
// 支持 Claude 和 Gemini 两种协议,根据 modelID 前缀自动选择
|
||||
// TestConnection 测试 Antigravity 账号连接。
|
||||
// 复用 antigravityRetryLoop 的完整重试 / credits overages / 智能重试逻辑,
|
||||
// 与真实调度行为一致。差异:不做账号切换(测试指定账号)、不记录 ops 错误。
|
||||
func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account *Account, modelID string) (*TestConnectionResult, error) {
|
||||
|
||||
// 获取 token
|
||||
@@ -1026,10 +1027,8 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account
|
||||
// 构建请求体
|
||||
var requestBody []byte
|
||||
if strings.HasPrefix(modelID, "gemini-") {
|
||||
// Gemini 模型:直接使用 Gemini 格式
|
||||
requestBody, err = s.buildGeminiTestRequest(projectID, mappedModel)
|
||||
} else {
|
||||
// Claude 模型:使用协议转换
|
||||
requestBody, err = s.buildClaudeTestRequest(projectID, mappedModel)
|
||||
}
|
||||
if err != nil {
|
||||
@@ -1042,64 +1041,63 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
baseURL := resolveAntigravityForwardBaseURL()
|
||||
if baseURL == "" {
|
||||
return nil, errors.New("no antigravity forward base url configured")
|
||||
}
|
||||
availableURLs := []string{baseURL}
|
||||
|
||||
var lastErr error
|
||||
for urlIdx, baseURL := range availableURLs {
|
||||
// 构建 HTTP 请求(总是使用流式 endpoint,与官方客户端一致)
|
||||
req, err := antigravity.NewAPIRequestWithURL(ctx, baseURL, "streamGenerateContent", accessToken, requestBody)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
|
||||
// 调试日志:Test 请求信息
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Test] account=%s request_size=%d url=%s", account.Name, len(requestBody), req.URL.String())
|
||||
|
||||
// 发送请求
|
||||
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("请求失败: %w", err)
|
||||
if shouldAntigravityFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Test] URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1])
|
||||
continue
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
// 读取响应
|
||||
respBody, err := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
_ = resp.Body.Close() // 立即关闭,避免循环内 defer 导致的资源泄漏
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取响应失败: %w", err)
|
||||
}
|
||||
|
||||
// 检查是否需要 URL 降级
|
||||
if shouldAntigravityFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Test] URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1])
|
||||
continue
|
||||
}
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
return nil, fmt.Errorf("API 返回 %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
// 解析流式响应,提取文本
|
||||
text := extractTextFromSSEResponse(respBody)
|
||||
|
||||
// 标记成功的 URL,下次优先使用
|
||||
antigravity.DefaultURLAvailability.MarkSuccess(baseURL)
|
||||
return &TestConnectionResult{
|
||||
Text: text,
|
||||
MappedModel: mappedModel,
|
||||
}, nil
|
||||
// 复用 antigravityRetryLoop:完整的重试 / credits overages / 智能重试
|
||||
prefix := fmt.Sprintf("[antigravity-Test] account=%d(%s)", account.ID, account.Name)
|
||||
p := antigravityRetryLoopParams{
|
||||
ctx: ctx,
|
||||
prefix: prefix,
|
||||
account: account,
|
||||
proxyURL: proxyURL,
|
||||
accessToken: accessToken,
|
||||
action: "streamGenerateContent",
|
||||
body: requestBody,
|
||||
c: nil, // 无 gin.Context → 跳过 ops 追踪
|
||||
httpUpstream: s.httpUpstream,
|
||||
settingService: s.settingService,
|
||||
accountRepo: s.accountRepo,
|
||||
requestedModel: modelID,
|
||||
handleError: testConnectionHandleError,
|
||||
}
|
||||
|
||||
return nil, lastErr
|
||||
result, err := s.antigravityRetryLoop(p)
|
||||
if err != nil {
|
||||
// AccountSwitchError → 测试时不切换账号,返回友好提示
|
||||
var switchErr *AntigravityAccountSwitchError
|
||||
if errors.As(err, &switchErr) {
|
||||
return nil, fmt.Errorf("该账号模型 %s 当前限流中,请稍后重试", switchErr.RateLimitedModel)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if result == nil || result.resp == nil {
|
||||
return nil, errors.New("upstream returned empty response")
|
||||
}
|
||||
defer func() { _ = result.resp.Body.Close() }()
|
||||
|
||||
respBody, err := io.ReadAll(io.LimitReader(result.resp.Body, 2<<20))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取响应失败: %w", err)
|
||||
}
|
||||
|
||||
if result.resp.StatusCode >= 400 {
|
||||
return nil, fmt.Errorf("API 返回 %d: %s", result.resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
text := extractTextFromSSEResponse(respBody)
|
||||
return &TestConnectionResult{Text: text, MappedModel: mappedModel}, nil
|
||||
}
|
||||
|
||||
// testConnectionHandleError 是 TestConnection 使用的轻量 handleError 回调。
|
||||
// 仅记录日志,不做 ops 错误追踪或粘性会话清除。
|
||||
func testConnectionHandleError(
|
||||
_ context.Context, prefix string, account *Account,
|
||||
statusCode int, _ http.Header, body []byte,
|
||||
requestedModel string, _ int64, _ string, _ bool,
|
||||
) *handleModelRateLimitResult {
|
||||
logger.LegacyPrintf("service.antigravity_gateway",
|
||||
"%s test_handle_error status=%d model=%s account=%d body=%s",
|
||||
prefix, statusCode, requestedModel, account.ID, truncateForLog(body, 200))
|
||||
return nil
|
||||
}
|
||||
|
||||
// buildGeminiTestRequest 构建 Gemini 格式测试请求
|
||||
@@ -1361,7 +1359,10 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, s.writeClaudeError(c, http.StatusBadGateway, "authentication_error", "Failed to get upstream access token")
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
ResponseBody: []byte(`{"error":{"type":"authentication_error","message":"Failed to get upstream access token"},"type":"error"}`),
|
||||
}
|
||||
}
|
||||
|
||||
// 获取 project_id(部分账户类型可能没有)
|
||||
@@ -2106,7 +2107,10 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, s.writeGoogleError(c, http.StatusBadGateway, "Failed to get upstream access token")
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
ResponseBody: []byte(`{"error":{"message":"Failed to get upstream access token","status":"UNAVAILABLE"}}`),
|
||||
}
|
||||
}
|
||||
|
||||
// 获取 project_id(部分账户类型可能没有)
|
||||
@@ -3082,6 +3086,22 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
|
||||
intervalCh = intervalTicker.C
|
||||
}
|
||||
|
||||
// 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开
|
||||
keepaliveInterval := time.Duration(0)
|
||||
if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 {
|
||||
keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second
|
||||
}
|
||||
var keepaliveTicker *time.Ticker
|
||||
if keepaliveInterval > 0 {
|
||||
keepaliveTicker = time.NewTicker(keepaliveInterval)
|
||||
defer keepaliveTicker.Stop()
|
||||
}
|
||||
var keepaliveCh <-chan time.Time
|
||||
if keepaliveTicker != nil {
|
||||
keepaliveCh = keepaliveTicker.C
|
||||
}
|
||||
lastDataAt := time.Now()
|
||||
|
||||
cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity gemini")
|
||||
|
||||
// 仅发送一次错误事件,避免多次写入导致协议混乱
|
||||
@@ -3114,6 +3134,8 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
|
||||
return nil, ev.err
|
||||
}
|
||||
|
||||
lastDataAt = time.Now()
|
||||
|
||||
line := ev.line
|
||||
trimmed := strings.TrimRight(line, "\r\n")
|
||||
if strings.HasPrefix(trimmed, "data:") {
|
||||
@@ -3173,6 +3195,19 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)")
|
||||
sendErrorEvent("stream_timeout")
|
||||
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")
|
||||
|
||||
case <-keepaliveCh:
|
||||
if cw.Disconnected() {
|
||||
continue
|
||||
}
|
||||
if time.Since(lastDataAt) < keepaliveInterval {
|
||||
continue
|
||||
}
|
||||
// SSE ping/keepalive:保持连接活跃防止 Cloudflare Tunnel 等代理断开
|
||||
if !cw.Fprintf(":\n\n") {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity gemini), continuing to drain upstream for billing")
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3903,6 +3938,22 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
|
||||
intervalCh = intervalTicker.C
|
||||
}
|
||||
|
||||
// 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开
|
||||
keepaliveInterval := time.Duration(0)
|
||||
if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 {
|
||||
keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second
|
||||
}
|
||||
var keepaliveTicker *time.Ticker
|
||||
if keepaliveInterval > 0 {
|
||||
keepaliveTicker = time.NewTicker(keepaliveInterval)
|
||||
defer keepaliveTicker.Stop()
|
||||
}
|
||||
var keepaliveCh <-chan time.Time
|
||||
if keepaliveTicker != nil {
|
||||
keepaliveCh = keepaliveTicker.C
|
||||
}
|
||||
lastDataAt := time.Now()
|
||||
|
||||
cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity claude")
|
||||
|
||||
// 仅发送一次错误事件,避免多次写入导致协议混乱
|
||||
@@ -3955,6 +4006,8 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
|
||||
return nil, fmt.Errorf("stream read error: %w", ev.err)
|
||||
}
|
||||
|
||||
lastDataAt = time.Now()
|
||||
|
||||
// 处理 SSE 行,转换为 Claude 格式
|
||||
claudeEvents := processor.ProcessLine(strings.TrimRight(ev.line, "\r\n"))
|
||||
if len(claudeEvents) > 0 {
|
||||
@@ -3977,6 +4030,20 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)")
|
||||
sendErrorEvent("stream_timeout")
|
||||
return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")
|
||||
|
||||
case <-keepaliveCh:
|
||||
if cw.Disconnected() {
|
||||
continue
|
||||
}
|
||||
if time.Since(lastDataAt) < keepaliveInterval {
|
||||
continue
|
||||
}
|
||||
// SSE ping 事件:Anthropic 原生格式,客户端会正确处理,
|
||||
// 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开
|
||||
if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity claude), continuing to drain upstream for billing")
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -4307,6 +4374,22 @@ func (s *AntigravityGatewayService) streamUpstreamResponse(c *gin.Context, resp
|
||||
intervalCh = intervalTicker.C
|
||||
}
|
||||
|
||||
// 下游 keepalive:防止代理/Cloudflare Tunnel 因连接空闲而断开
|
||||
keepaliveInterval := time.Duration(0)
|
||||
if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 {
|
||||
keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second
|
||||
}
|
||||
var keepaliveTicker *time.Ticker
|
||||
if keepaliveInterval > 0 {
|
||||
keepaliveTicker = time.NewTicker(keepaliveInterval)
|
||||
defer keepaliveTicker.Stop()
|
||||
}
|
||||
var keepaliveCh <-chan time.Time
|
||||
if keepaliveTicker != nil {
|
||||
keepaliveCh = keepaliveTicker.C
|
||||
}
|
||||
lastDataAt := time.Now()
|
||||
|
||||
flusher, _ := c.Writer.(http.Flusher)
|
||||
cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity upstream")
|
||||
|
||||
@@ -4324,6 +4407,8 @@ func (s *AntigravityGatewayService) streamUpstreamResponse(c *gin.Context, resp
|
||||
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}
|
||||
}
|
||||
|
||||
lastDataAt = time.Now()
|
||||
|
||||
line := ev.line
|
||||
|
||||
// 记录首 token 时间
|
||||
@@ -4349,6 +4434,20 @@ func (s *AntigravityGatewayService) streamUpstreamResponse(c *gin.Context, resp
|
||||
}
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity upstream)")
|
||||
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}
|
||||
|
||||
case <-keepaliveCh:
|
||||
if cw.Disconnected() {
|
||||
continue
|
||||
}
|
||||
if time.Since(lastDataAt) < keepaliveInterval {
|
||||
continue
|
||||
}
|
||||
// SSE ping 事件:Anthropic 原生格式,客户端会正确处理,
|
||||
// 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开
|
||||
if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity upstream), continuing to drain upstream for billing")
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -57,16 +57,16 @@ func TestAntigravityGatewayService_GetMappedModel(t *testing.T) {
|
||||
expected: "claude-opus-4-6-thinking",
|
||||
},
|
||||
{
|
||||
name: "默认映射 - claude-haiku-4-5 → claude-sonnet-4-5",
|
||||
name: "默认映射 - claude-haiku-4-5 → claude-sonnet-4-6",
|
||||
requestedModel: "claude-haiku-4-5",
|
||||
accountMapping: nil,
|
||||
expected: "claude-sonnet-4-5",
|
||||
expected: "claude-sonnet-4-6",
|
||||
},
|
||||
{
|
||||
name: "默认映射 - claude-haiku-4-5-20251001 → claude-sonnet-4-5",
|
||||
name: "默认映射 - claude-haiku-4-5-20251001 → claude-sonnet-4-6",
|
||||
requestedModel: "claude-haiku-4-5-20251001",
|
||||
accountMapping: nil,
|
||||
expected: "claude-sonnet-4-5",
|
||||
expected: "claude-sonnet-4-6",
|
||||
},
|
||||
{
|
||||
name: "默认映射 - claude-sonnet-4-5-20250929 → claude-sonnet-4-5",
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user