Skip to content

Commit be51167

Browse files
flantofunflantofun
authored
fix: separate codex quotas by account identity (#418)
1 parent a2242ad commit be51167

File tree

9 files changed

+247
-28
lines changed

9 files changed

+247
-28
lines changed

‎internal/domain/model.go‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -654,7 +654,7 @@ type CodexRateLimitInfo struct {
654654
SecondaryWindow *CodexQuotaWindow `json:"secondaryWindow,omitempty"`
655655
}
656656

657-
// Codex 账户配额（基于邮箱存储）
657+
// Codex 账户配额（优先按 account_id 区分，回退到 email）
658658
type CodexQuota struct {
659659
ID uint64 `json:"id"`
660660
CreatedAt time.Time `json:"createdAt"`
@@ -666,7 +666,10 @@ type CodexQuota struct {
666666
// 所属租户
667667
TenantID uint64 `json:"tenantID"`
668668

669-
// 邮箱作为唯一标识
669+
// 配额身份键：优先 account:<account_id>，否则 email:<email>
670+
IdentityKey string `json:"identityKey"`
671+
672+
// 邮箱（展示和回退匹配用）
670673
Email string `json:"email"`
671674

672675
// 账户 ID
@@ -688,6 +691,18 @@ type CodexQuota struct {
688691
CodeReviewWindow *CodexQuotaWindow `json:"codeReviewWindow,omitempty"`
689692
}
690693

694+
func CodexQuotaIdentityKey(email, accountID string) string {
695+
accountID = strings.TrimSpace(accountID)
696+
if accountID != "" {
697+
return "account:" + accountID
698+
}
699+
email = strings.TrimSpace(email)
700+
if email != "" {
701+
return "email:" + email
702+
}
703+
return ""
704+
}
705+
691706
// Provider 统计信息
692707
type ProviderStats struct {
693708
ProviderID uint64 `json:"providerID"`

‎internal/handler/codex.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -603,10 +603,11 @@ func (h *CodexHandler) GetBatchQuotas(ctx context.Context) (*CodexBatchQuotaResu
603603

604604
config := provider.Config.Codex
605605
email := config.Email
606+
identityKey := domain.CodexQuotaIdentityKey(config.Email, config.AccountID)
606607

607608
// 优先从数据库获取缓存的配额（无论是否过期）
608-
if email != "" && h.quotaRepo != nil {
609-
cachedQuota, err := h.quotaRepo.GetByEmail(tenantID, email)
609+
if identityKey != "" && h.quotaRepo != nil {
610+
cachedQuota, err := h.quotaRepo.GetByIdentityKey(tenantID, identityKey)
610611
if err == nil && cachedQuota != nil {
611612
result.Quotas[provider.ID] = h.domainQuotaToResponse(cachedQuota)
612613
continue
@@ -680,11 +681,12 @@ func (h *CodexHandler) isTokenExpired(expiresAt string) bool {
680681

681682
// saveQuotaToDB saves Codex quota to database
682683
func (h *CodexHandler) saveQuotaToDB(email, accountID, planType string, usage *codex.CodexUsageResponse, isForbidden bool) {
683-
if h.quotaRepo == nil || email == "" {
684+
if h.quotaRepo == nil || domain.CodexQuotaIdentityKey(email, accountID) == "" {
684685
return
685686
}
686687

687688
quota := &domain.CodexQuota{
689+
IdentityKey: domain.CodexQuotaIdentityKey(email, accountID),
688690
Email: email,
689691
AccountID: accountID,
690692
PlanType: planType,

‎internal/repository/interfaces.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,8 +191,10 @@ type AntigravityQuotaRepository interface {
191191
}
192192

193193
type CodexQuotaRepository interface {
194-
// Upsert 更新或插入配额（基于邮箱）
194+
// Upsert 更新或插入配额（优先基于 identityKey，其次回退邮箱）
195195
Upsert(quota *domain.CodexQuota) error
196+
// GetByIdentityKey 根据身份键获取配额
197+
GetByIdentityKey(tenantID uint64, identityKey string) (*domain.CodexQuota, error)
196198
// GetByEmail 根据邮箱获取配额
197199
GetByEmail(tenantID uint64, email string) (*domain.CodexQuota, error)
198200
// List 获取所有配额

‎internal/repository/sqlite/codex_quota.go‎

Lines changed: 37 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,27 @@ func NewCodexQuotaRepository(d *DB) *CodexQuotaRepository {
1717

1818
func (r *CodexQuotaRepository) Upsert(quota *domain.CodexQuota) error {
1919
now := time.Now()
20+
identityKey := domain.CodexQuotaIdentityKey(quota.Email, quota.AccountID)
21+
quota.IdentityKey = identityKey
22+
23+
query := tenantScope(r.db.gorm.Model(&CodexQuota{}), quota.TenantID)
24+
if identityKey != "" {
25+
query = query.Where("identity_key = ? AND deleted_at = 0", identityKey)
26+
} else {
27+
query = query.Where("email = ? AND deleted_at = 0", quota.Email)
28+
}
2029

21-
// Try to update first
22-
result := tenantScope(r.db.gorm.Model(&CodexQuota{}), quota.TenantID).
23-
Where("email = ? AND deleted_at = 0", quota.Email).
24-
Updates(map[string]any{
25-
"updated_at": toTimestamp(now),
26-
"account_id": quota.AccountID,
27-
"plan_type": quota.PlanType,
28-
"is_forbidden": quota.IsForbidden,
29-
"primary_window": toJSON(quota.PrimaryWindow),
30-
"secondary_window": toJSON(quota.SecondaryWindow),
31-
"code_review_window": toJSON(quota.CodeReviewWindow),
32-
})
30+
result := query.Updates(map[string]any{
31+
"updated_at": toTimestamp(now),
32+
"identity_key": identityKey,
33+
"email": quota.Email,
34+
"account_id": quota.AccountID,
35+
"plan_type": quota.PlanType,
36+
"is_forbidden": quota.IsForbidden,
37+
"primary_window": toJSON(quota.PrimaryWindow),
38+
"secondary_window": toJSON(quota.SecondaryWindow),
39+
"code_review_window": toJSON(quota.CodeReviewWindow),
40+
})
3341

3442
if result.Error != nil {
3543
return result.Error
@@ -53,6 +61,21 @@ func (r *CodexQuotaRepository) Upsert(quota *domain.CodexQuota) error {
5361
return nil
5462
}
5563

64+
func (r *CodexQuotaRepository) GetByIdentityKey(tenantID uint64, identityKey string) (*domain.CodexQuota, error) {
65+
if identityKey == "" {
66+
return nil, nil
67+
}
68+
var model CodexQuota
69+
err := tenantScope(r.db.gorm, tenantID).Where("identity_key = ? AND deleted_at = 0", identityKey).First(&model).Error
70+
if err != nil {
71+
if err == gorm.ErrRecordNotFound {
72+
return nil, nil
73+
}
74+
return nil, err
75+
}
76+
return r.toDomain(&model), nil
77+
}
78+
5679
func (r *CodexQuotaRepository) GetByEmail(tenantID uint64, email string) (*domain.CodexQuota, error) {
5780
var model CodexQuota
5881
err := tenantScope(r.db.gorm, tenantID).Where("email = ? AND deleted_at = 0", email).First(&model).Error
@@ -94,6 +117,7 @@ func (r *CodexQuotaRepository) toModel(q *domain.CodexQuota) *CodexQuota {
94117
DeletedAt: toTimestampPtr(q.DeletedAt),
95118
},
96119
TenantID: q.TenantID,
120+
IdentityKey: q.IdentityKey,
97121
Email: q.Email,
98122
AccountID: q.AccountID,
99123
PlanType: q.PlanType,
@@ -111,6 +135,7 @@ func (r *CodexQuotaRepository) toDomain(m *CodexQuota) *domain.CodexQuota {
111135
UpdatedAt: fromTimestamp(m.UpdatedAt),
112136
DeletedAt: fromTimestampPtr(m.DeletedAt),
113137
TenantID: m.TenantID,
138+
IdentityKey: m.IdentityKey,
114139
Email: m.Email,
115140
AccountID: m.AccountID,
116141
PlanType: m.PlanType,

‎internal/repository/sqlite/codex_quota_test.go‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
package sqlite
2+
3+
import (
4+
"testing"
5+
6+
"github.com/awsl-project/maxx/internal/domain"
7+
)
8+
9+
func TestCodexQuotaRepository_UpsertUsesIdentityKey(t *testing.T) {
10+
db, err := NewDBWithDSN("sqlite://:memory:")
11+
if err != nil {
12+
t.Fatalf("Failed to create DB: %v", err)
13+
}
14+
defer db.Close()
15+
16+
repo := NewCodexQuotaRepository(db)
17+
18+
quota1 := &domain.CodexQuota{
19+
TenantID: 1,
20+
Email: "same@example.com",
21+
AccountID: "acct-1",
22+
PlanType: "team",
23+
}
24+
if err := repo.Upsert(quota1); err != nil {
25+
t.Fatalf("Upsert quota1 failed: %v", err)
26+
}
27+
28+
quota2 := &domain.CodexQuota{
29+
TenantID: 1,
30+
Email: "same@example.com",
31+
AccountID: "acct-2",
32+
PlanType: "team",
33+
}
34+
if err := repo.Upsert(quota2); err != nil {
35+
t.Fatalf("Upsert quota2 failed: %v", err)
36+
}
37+
38+
quota1.PlanType = "team-updated"
39+
if err := repo.Upsert(quota1); err != nil {
40+
t.Fatalf("Upsert quota1 update failed: %v", err)
41+
}
42+
43+
list, err := repo.List(1)
44+
if err != nil {
45+
t.Fatalf("List failed: %v", err)
46+
}
47+
if len(list) != 2 {
48+
t.Fatalf("expected 2 quota rows, got %d", len(list))
49+
}
50+
51+
got1, err := repo.GetByIdentityKey(1, domain.CodexQuotaIdentityKey(quota1.Email, quota1.AccountID))
52+
if err != nil {
53+
t.Fatalf("GetByIdentityKey quota1 failed: %v", err)
54+
}
55+
if got1 == nil || got1.AccountID != "acct-1" || got1.PlanType != "team-updated" {
56+
t.Fatalf("unexpected quota1 row: %#v", got1)
57+
}
58+
59+
got2, err := repo.GetByIdentityKey(1, domain.CodexQuotaIdentityKey(quota2.Email, quota2.AccountID))
60+
if err != nil {
61+
t.Fatalf("GetByIdentityKey quota2 failed: %v", err)
62+
}
63+
if got2 == nil || got2.AccountID != "acct-2" {
64+
t.Fatalf("unexpected quota2 row: %#v", got2)
65+
}
66+
}

‎internal/repository/sqlite/migrations.go‎

Lines changed: 94 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ var migrations = []Migration{
119119

120120
// 2. Update all existing rows to belong to default tenant
121121
for _, table := range tenantScopedTables {
122-
result := db.Exec("UPDATE "+table+" SET tenant_id = 1 WHERE tenant_id = 0 OR tenant_id IS NULL")
122+
result := db.Exec("UPDATE " + table + " SET tenant_id = 1 WHERE tenant_id = 0 OR tenant_id IS NULL")
123123
if result.Error != nil {
124124
log.Printf("[Migration] Warning: Failed to update tenant_id for %s: %v", table, result.Error)
125125
// Continue with other tables
@@ -217,6 +217,87 @@ var migrations = []Migration{
217217
return nil
218218
},
219219
},
220+
{
221+
Version: 7,
222+
Description: "Make Codex quota identity account-aware to avoid same-email quota collisions",
223+
Up: func(db *gorm.DB) error {
224+
var backfillSQL string
225+
switch db.Dialector.Name() {
226+
case "mysql":
227+
backfillSQL = `
228+
UPDATE codex_quotas
229+
SET identity_key = CASE
230+
WHEN account_id IS NOT NULL AND TRIM(account_id) != '' THEN CONCAT('account:', TRIM(account_id))
231+
WHEN email IS NOT NULL AND TRIM(email) != '' THEN CONCAT('email:', TRIM(email))
232+
ELSE NULL
233+
END
234+
WHERE identity_key IS NULL OR TRIM(identity_key) = ''
235+
`
236+
default:
237+
backfillSQL = `
238+
UPDATE codex_quotas
239+
SET identity_key = CASE
240+
WHEN account_id IS NOT NULL AND TRIM(account_id) != '' THEN 'account:' || TRIM(account_id)
241+
WHEN email IS NOT NULL AND TRIM(email) != '' THEN 'email:' || TRIM(email)
242+
ELSE NULL
243+
END
244+
WHERE identity_key IS NULL OR TRIM(identity_key) = ''
245+
`
246+
}
247+
if err := db.Exec(backfillSQL).Error; err != nil {
248+
return err
249+
}
250+
251+
switch db.Dialector.Name() {
252+
case "mysql":
253+
if err := db.Exec("DROP INDEX idx_codex_quotas_tenant_email ON codex_quotas").Error; err != nil && !isMySQLMissingIndexError(err) {
254+
return err
255+
}
256+
if err := db.Exec("CREATE INDEX idx_codex_quotas_email ON codex_quotas(email)").Error; err != nil && !isMySQLDuplicateIndexError(err) {
257+
return err
258+
}
259+
if err := db.Exec("CREATE UNIQUE INDEX idx_codex_quotas_tenant_identity ON codex_quotas(tenant_id, identity_key)").Error; err != nil && !isMySQLDuplicateIndexError(err) {
260+
return err
261+
}
262+
default:
263+
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_tenant_email").Error; err != nil {
264+
return err
265+
}
266+
if err := db.Exec("CREATE INDEX IF NOT EXISTS idx_codex_quotas_email ON codex_quotas(email)").Error; err != nil {
267+
return err
268+
}
269+
if err := db.Exec("CREATE UNIQUE INDEX IF NOT EXISTS idx_codex_quotas_tenant_identity ON codex_quotas(tenant_id, identity_key)").Error; err != nil {
270+
return err
271+
}
272+
}
273+
return nil
274+
},
275+
Down: func(db *gorm.DB) error {
276+
switch db.Dialector.Name() {
277+
case "mysql":
278+
if err := db.Exec("DROP INDEX idx_codex_quotas_tenant_identity ON codex_quotas").Error; err != nil && !isMySQLMissingIndexError(err) {
279+
return err
280+
}
281+
if err := db.Exec("DROP INDEX idx_codex_quotas_email ON codex_quotas").Error; err != nil && !isMySQLMissingIndexError(err) {
282+
return err
283+
}
284+
if err := db.Exec("CREATE UNIQUE INDEX idx_codex_quotas_tenant_email ON codex_quotas(tenant_id, email)").Error; err != nil && !isMySQLDuplicateIndexError(err) {
285+
return err
286+
}
287+
default:
288+
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_tenant_identity").Error; err != nil {
289+
return err
290+
}
291+
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_email").Error; err != nil {
292+
return err
293+
}
294+
if err := db.Exec("CREATE UNIQUE INDEX IF NOT EXISTS idx_codex_quotas_tenant_email ON codex_quotas(tenant_id, email)").Error; err != nil {
295+
return err
296+
}
297+
}
298+
return nil
299+
},
300+
},
220301
}
221302

222303
func isMySQLDuplicateIndexError(err error) bool {
@@ -232,6 +313,18 @@ func isMySQLDuplicateIndexError(err error) bool {
232313
return strings.Contains(lower, "duplicate key name") || strings.Contains(lower, "error 1061")
233314
}
234315

316+
func isMySQLMissingIndexError(err error) bool {
317+
if err == nil {
318+
return false
319+
}
320+
var mysqlErr *mysqlDriver.MySQLError
321+
if errors.As(err, &mysqlErr) {
322+
return mysqlErr.Number == 1091 // ER_CANT_DROP_FIELD_OR_KEY
323+
}
324+
lower := strings.ToLower(err.Error())
325+
return strings.Contains(lower, "can't drop") && strings.Contains(lower, "check that column/key exists")
326+
}
327+
235328
// RunMigrations 运行所有待执行的迁移
236329
func (d *DB) RunMigrations() error {
237330
// 确保迁移表存在（由 GORM AutoMigrate 处理）

‎internal/repository/sqlite/migrations_test.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,15 @@ func TestIsMySQLDuplicateIndexError(t *testing.T) {
2121
t.Fatalf("expected false for unrelated error")
2222
}
2323
}
24+
25+
func TestIsMySQLMissingIndexError(t *testing.T) {
26+
if !isMySQLMissingIndexError(&mysqlDriver.MySQLError{Number: 1091, Message: "Can't DROP"}) {
27+
t.Fatalf("expected true for ER_CANT_DROP_FIELD_OR_KEY(1091)")
28+
}
29+
if !isMySQLMissingIndexError(errors.New("Error 1091: Can't DROP 'idx_x'; check that column/key exists")) {
30+
t.Fatalf("expected true for missing index string match fallback")
31+
}
32+
if isMySQLMissingIndexError(errors.New("some other error")) {
33+
t.Fatalf("expected false for unrelated error")
34+
}
35+
}

‎internal/repository/sqlite/models.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -252,8 +252,9 @@ func (AntigravityQuota) TableName() string { return "antigravity_quotas" }
252252
// CodexQuota model
253253
type CodexQuota struct {
254254
SoftDeleteModel
255-
TenantID uint64 `gorm:"uniqueIndex:idx_codex_quotas_tenant_email"`
256-
Email string `gorm:"size:255;uniqueIndex:idx_codex_quotas_tenant_email"`
255+
TenantID uint64 `gorm:"uniqueIndex:idx_codex_quotas_tenant_identity"`
256+
IdentityKey string `gorm:"size:255;column:identity_key;uniqueIndex:idx_codex_quotas_tenant_identity"`
257+
Email string `gorm:"size:255;index:idx_codex_quotas_email"`
257258
AccountID string `gorm:"size:128;column:account_id"`
258259
PlanType string `gorm:"size:64"`
259260
IsForbidden int

0 commit comments

Comments
 (0)