Skip to content

Commit 8097558

Browse files
james-6-23james-6-23
committed
feat: borrow image proxy and usage ideas from pr100
1 parent 80236d1 commit 8097558

21 files changed

Lines changed: 1359 additions & 66 deletions

‎admin/handler.go‎

Lines changed: 109 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,8 @@ func (h *Handler) SetPoolSizes(pgMaxConns, redisPoolSize int) {
9797

9898
// RegisterRoutes 注册管理 API 路由
9999
func (h *Handler) RegisterRoutes(r *gin.Engine) {
100+
r.GET("/p/img/:id", h.GetSignedImageAssetFile)
101+
100102
api := r.Group("/api/admin")
101103
api.Use(h.adminAuthMiddleware())
102104
api.GET("/stats", h.GetStats)
@@ -283,6 +285,8 @@ type accountResponse struct {
283285
ErrorRequests int64 `json:"error_requests"`
284286
UsagePercent7d *float64 `json:"usage_percent_7d"`
285287
UsagePercent5h *float64 `json:"usage_percent_5h"`
288+
Usage5hDetail *accountUsageWindow `json:"usage_5h_detail,omitempty"`
289+
Usage7dDetail *accountUsageWindow `json:"usage_7d_detail,omitempty"`
286290
Reset5hAt string `json:"reset_5h_at,omitempty"`
287291
Reset7dAt string `json:"reset_7d_at,omitempty"`
288292
ScoreBreakdown schedulerBreakdownResponse `json:"scheduler_breakdown"`
@@ -294,6 +298,11 @@ type accountResponse struct {
294298
AllowedAPIKeyIDs []int64 `json:"allowed_api_key_ids"`
295299
}
296300

301+
type accountUsageWindow struct {
302+
Requests int64 `json:"requests"`
303+
Tokens int64 `json:"tokens"`
304+
}
305+
297306
type schedulerBreakdownResponse struct {
298307
UnauthorizedPenalty float64 `json:"unauthorized_penalty"`
299308
RateLimitPenalty float64 `json:"rate_limit_penalty"`
@@ -328,6 +337,7 @@ func (h *Handler) ListAccounts(c *gin.Context) {
328337

329338
// 获取每账号近 7 天请求统计（带 30 秒内存缓存）
330339
reqCounts := h.getCachedRequestCounts()
340+
usage5h, usage7d := h.getAccountUsageWindows(ctx)
331341

332342
accounts := make([]accountResponse, 0, len(rows))
333343
for _, row := range rows {
@@ -412,6 +422,12 @@ func (h *Handler) ListAccounts(c *gin.Context) {
412422
resp.SuccessRequests = rc.SuccessCount
413423
resp.ErrorRequests = rc.ErrorCount
414424
}
425+
if usage, ok := usage5h[row.ID]; ok {
426+
resp.Usage5hDetail = &accountUsageWindow{Requests: usage.Requests, Tokens: usage.Tokens}
427+
}
428+
if usage, ok := usage7d[row.ID]; ok {
429+
resp.Usage7dDetail = &accountUsageWindow{Requests: usage.Requests, Tokens: usage.Tokens}
430+
}
415431
accounts = append(accounts, resp)
416432
}
417433

@@ -689,6 +705,21 @@ func (h *Handler) getCachedRequestCounts() map[int64]*database.AccountRequestCou
689705
return counts
690706
}
691707

708+
func (h *Handler) getAccountUsageWindows(ctx context.Context) (map[int64]*database.AccountTimeRangeUsage, map[int64]*database.AccountTimeRangeUsage) {
709+
now := time.Now()
710+
usage5h, err := h.db.GetAccountTimeRangeUsage(ctx, now.Add(-5*time.Hour))
711+
if err != nil {
712+
log.Printf("获取账号 5h 用量统计失败: %v", err)
713+
usage5h = make(map[int64]*database.AccountTimeRangeUsage)
714+
}
715+
usage7d, err := h.db.GetAccountTimeRangeUsage(ctx, now.AddDate(0, 0, -7))
716+
if err != nil {
717+
log.Printf("获取账号 7d 用量统计失败: %v", err)
718+
usage7d = make(map[int64]*database.AccountTimeRangeUsage)
719+
}
720+
return usage5h, usage7d
721+
}
722+
692723
type addAccountReq struct {
693724
Name string `json:"name"`
694725
RefreshToken string `json:"refresh_token"`
@@ -944,6 +975,7 @@ func (h *Handler) AddATAccount(c *gin.Context) {
944975
// importToken 导入时的统一 token 载体
945976
type importToken struct {
946977
refreshToken string
978+
sessionToken string
947979
accessToken string // AT-only 兼容路径
948980
name string
949981
email string
@@ -955,15 +987,17 @@ type importToken struct {
955987

956988
// jsonAccountEntry CLIProxyAPI 凭证 JSON 条目
957989
type jsonAccountEntry struct {
958-
RefreshToken string `json:"refresh_token"`
959-
AccessToken string `json:"access_token"`
960-
IDToken string `json:"id_token"`
961-
AccountID string `json:"account_id"`
962-
Email string `json:"email"`
963-
Name string `json:"name"`
964-
PlanType string `json:"plan_type"`
965-
Expired string `json:"expired"`
966-
ExpiresAt string `json:"expires_at"`
990+
RefreshToken string `json:"refresh_token"`
991+
SessionToken string `json:"session_token"`
992+
SessionTokenCamel string `json:"sessionToken"`
993+
AccessToken string `json:"access_token"`
994+
IDToken string `json:"id_token"`
995+
AccountID string `json:"account_id"`
996+
Email string `json:"email"`
997+
Name string `json:"name"`
998+
PlanType string `json:"plan_type"`
999+
Expired string `json:"expired"`
1000+
ExpiresAt string `json:"expires_at"`
9671001
}
9681002

9691003
type sub2apiImportPayload struct {
@@ -976,14 +1010,16 @@ type sub2apiAccountEntry struct {
9761010
}
9771011

9781012
type sub2apiAccountCredentials struct {
979-
RefreshToken string `json:"refresh_token"`
980-
AccessToken string `json:"access_token"`
981-
IDToken string `json:"id_token"`
982-
AccountID string `json:"account_id"`
983-
Email string `json:"email"`
984-
PlanType string `json:"plan_type"`
985-
ExpiresAt string `json:"expires_at"`
986-
Expired string `json:"expired"`
1013+
RefreshToken string `json:"refresh_token"`
1014+
SessionToken string `json:"session_token"`
1015+
SessionTokenCamel string `json:"sessionToken"`
1016+
AccessToken string `json:"access_token"`
1017+
IDToken string `json:"id_token"`
1018+
AccountID string `json:"account_id"`
1019+
Email string `json:"email"`
1020+
PlanType string `json:"plan_type"`
1021+
ExpiresAt string `json:"expires_at"`
1022+
Expired string `json:"expired"`
9871023
}
9881024

9891025
var utf8BOM = []byte{0xef, 0xbb, 0xbf}
@@ -1028,13 +1064,15 @@ func jsonAccountEntriesToTokens(entries []jsonAccountEntry) []importToken {
10281064
tokens := make([]importToken, 0, len(entries))
10291065
for _, entry := range entries {
10301066
rt := strings.TrimSpace(entry.RefreshToken)
1067+
st := firstNonEmpty(entry.SessionToken, entry.SessionTokenCamel)
10311068
at := strings.TrimSpace(entry.AccessToken)
10321069
email := strings.TrimSpace(entry.Email)
10331070
name := firstNonEmpty(entry.Name, email)
10341071

1035-
if rt != "" || at != "" {
1072+
if rt != "" || st != "" || at != "" {
10361073
tokens = append(tokens, importToken{
10371074
refreshToken: rt,
1075+
sessionToken: st,
10381076
accessToken: at,
10391077
name: name,
10401078
email: email,
@@ -1057,6 +1095,7 @@ func parseSub2APIJSONImportTokens(data []byte) []importToken {
10571095
tokens := make([]importToken, 0, len(payload.Accounts))
10581096
for _, account := range payload.Accounts {
10591097
rt := strings.TrimSpace(account.Credentials.RefreshToken)
1098+
st := firstNonEmpty(account.Credentials.SessionToken, account.Credentials.SessionTokenCamel)
10601099
at := strings.TrimSpace(account.Credentials.AccessToken)
10611100
name := strings.TrimSpace(account.Name)
10621101
email := strings.TrimSpace(account.Credentials.Email)
@@ -1065,9 +1104,10 @@ func parseSub2APIJSONImportTokens(data []byte) []importToken {
10651104
name = email
10661105
}
10671106

1068-
if rt != "" || at != "" {
1107+
if rt != "" || st != "" || at != "" {
10691108
tokens = append(tokens, importToken{
10701109
refreshToken: rt,
1110+
sessionToken: st,
10711111
accessToken: at,
10721112
name: name,
10731113
email: email,
@@ -1247,8 +1287,9 @@ func setupSSE(c *gin.Context) {
12471287

12481288
// importAccountsCommon 公共的去重、并发插入、SSE 进度推送逻辑（支持 RT 和 AT-only 混合导入）
12491289
func (h *Handler) importAccountsCommon(c *gin.Context, tokens []importToken, proxyURL string) {
1250-
// 文件内去重（RT 和 AT 分别去重）
1290+
// 文件内去重（RT、ST 和 AT 分别去重）
12511291
seenRT := make(map[string]bool)
1292+
seenST := make(map[string]bool)
12521293
seenAT := make(map[string]bool)
12531294
var unique []importToken
12541295
for _, t := range tokens {
@@ -1257,6 +1298,17 @@ func (h *Handler) importAccountsCommon(c *gin.Context, tokens []importToken, pro
12571298
seenRT[t.refreshToken] = true
12581299
unique = append(unique, t)
12591300
}
1301+
if t.sessionToken != "" {
1302+
seenST[t.sessionToken] = true
1303+
}
1304+
if t.accessToken != "" {
1305+
seenAT[t.accessToken] = true
1306+
}
1307+
} else if t.sessionToken != "" {
1308+
if !seenST[t.sessionToken] {
1309+
seenST[t.sessionToken] = true
1310+
unique = append(unique, t)
1311+
}
12601312
if t.accessToken != "" {
12611313
seenAT[t.accessToken] = true
12621314
}
@@ -1290,6 +1342,15 @@ func (h *Handler) importAccountsCommon(c *gin.Context, tokens []importToken, pro
12901342
existingATs = make(map[string]bool)
12911343
}
12921344
}
1345+
hasST := len(seenST) > 0
1346+
var existingSTs map[string]bool
1347+
if hasST {
1348+
existingSTs, err = h.db.GetAllSessionTokens(dedupeCtx)
1349+
if err != nil {
1350+
log.Printf("查询已有 ST 失败: %v", err)
1351+
existingSTs = make(map[string]bool)
1352+
}
1353+
}
12931354

12941355
var newTokens []importToken
12951356
duplicateCount := 0
@@ -1298,6 +1359,16 @@ func (h *Handler) importAccountsCommon(c *gin.Context, tokens []importToken, pro
12981359
case t.refreshToken != "":
12991360
if existingRTs[t.refreshToken] {
13001361
duplicateCount++
1362+
} else if t.sessionToken != "" && existingSTs[t.sessionToken] {
1363+
duplicateCount++
1364+
} else if t.accessToken != "" && existingATs[t.accessToken] {
1365+
duplicateCount++
1366+
} else {
1367+
newTokens = append(newTokens, t)
1368+
}
1369+
case t.sessionToken != "":
1370+
if existingSTs[t.sessionToken] {
1371+
duplicateCount++
13011372
} else if t.accessToken != "" && existingATs[t.accessToken] {
13021373
duplicateCount++
13031374
} else {
@@ -1388,6 +1459,7 @@ func (h *Handler) importAccountsCommon(c *gin.Context, tokens []importToken, pro
13881459
h.db.InsertAccountEventAsync(id, "added", "import_at")
13891460

13901461
seed := normalizeTokenCredentialSeed(tokenCredentialSeed{
1462+
sessionToken: tok.sessionToken,
13911463
accessToken: tok.accessToken,
13921464
idToken: tok.idToken,
13931465
accountID: tok.accountID,
@@ -1409,7 +1481,22 @@ func (h *Handler) importAccountsCommon(c *gin.Context, tokens []importToken, pro
14091481
}
14101482

14111483
insertCtx, insertCancel := context.WithTimeout(context.Background(), 5*time.Second)
1412-
id, err := h.db.InsertAccount(insertCtx, name, tok.refreshToken, proxyURL)
1484+
var id int64
1485+
var err error
1486+
if tok.refreshToken != "" {
1487+
id, err = h.db.InsertAccount(insertCtx, name, tok.refreshToken, proxyURL)
1488+
} else {
1489+
seed := normalizeTokenCredentialSeed(tokenCredentialSeed{
1490+
sessionToken: tok.sessionToken,
1491+
accessToken: tok.accessToken,
1492+
idToken: tok.idToken,
1493+
accountID: tok.accountID,
1494+
email: tok.email,
1495+
planType: tok.planType,
1496+
expiresAtRaw: tok.expiresAt,
1497+
})
1498+
id, err = h.db.InsertAccountWithCredentials(insertCtx, name, tokenCredentialMap(seed), proxyURL)
1499+
}
14131500
insertCancel()
14141501

14151502
if err != nil {
@@ -1425,6 +1512,7 @@ func (h *Handler) importAccountsCommon(c *gin.Context, tokens []importToken, pro
14251512

14261513
seed := normalizeTokenCredentialSeed(tokenCredentialSeed{
14271514
refreshToken: tok.refreshToken,
1515+
sessionToken: tok.sessionToken,
14281516
accessToken: tok.accessToken,
14291517
idToken: tok.idToken,
14301518
accountID: tok.accountID,

0 commit comments

Comments
 (0)