feat: OpenAI OAuth账号显示Codex使用量
从响应头提取x-codex-*使用量信息并保存到账号Extra字段, 前端账号列表展示5h/7d窗口的使用进度条。
This commit is contained in:
@@ -23,14 +23,18 @@ func (r *AccountRepository) Create(ctx context.Context, account *model.Account)
|
||||
|
||||
func (r *AccountRepository) GetByID(ctx context.Context, id int64) (*model.Account, error) {
|
||||
var account model.Account
|
||||
err := r.db.WithContext(ctx).Preload("Proxy").Preload("AccountGroups").First(&account, id).Error
|
||||
err := r.db.WithContext(ctx).Preload("Proxy").Preload("AccountGroups.Group").First(&account, id).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 填充 GroupIDs 虚拟字段
|
||||
// 填充 GroupIDs 和 Groups 虚拟字段
|
||||
account.GroupIDs = make([]int64, 0, len(account.AccountGroups))
|
||||
account.Groups = make([]*model.Group, 0, len(account.AccountGroups))
|
||||
for _, ag := range account.AccountGroups {
|
||||
account.GroupIDs = append(account.GroupIDs, ag.GroupID)
|
||||
if ag.Group != nil {
|
||||
account.Groups = append(account.Groups, ag.Group)
|
||||
}
|
||||
}
|
||||
return &account, nil
|
||||
}
|
||||
@@ -303,3 +307,31 @@ func (r *AccountRepository) SetSchedulable(ctx context.Context, id int64, schedu
|
||||
return r.db.WithContext(ctx).Model(&model.Account{}).Where("id = ?", id).
|
||||
Update("schedulable", schedulable).Error
|
||||
}
|
||||
|
||||
// UpdateExtra updates specific fields in account's Extra JSONB field
|
||||
// It merges the updates into existing Extra data without overwriting other fields
|
||||
func (r *AccountRepository) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
|
||||
if len(updates) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get current account to preserve existing Extra data
|
||||
var account model.Account
|
||||
if err := r.db.WithContext(ctx).Select("extra").Where("id = ?", id).First(&account).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Initialize Extra if nil
|
||||
if account.Extra == nil {
|
||||
account.Extra = make(model.JSONB)
|
||||
}
|
||||
|
||||
// Merge updates into existing Extra
|
||||
for k, v := range updates {
|
||||
account.Extra[k] = v
|
||||
}
|
||||
|
||||
// Save updated Extra
|
||||
return r.db.WithContext(ctx).Model(&model.Account{}).Where("id = ?", id).
|
||||
Update("extra", account.Extra).Error
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -38,6 +39,18 @@ var openaiAllowedHeaders = map[string]bool{
|
||||
"session_id": true,
|
||||
}
|
||||
|
||||
// OpenAICodexUsageSnapshot represents Codex API usage limits from response headers
|
||||
type OpenAICodexUsageSnapshot struct {
|
||||
PrimaryUsedPercent *float64 `json:"primary_used_percent,omitempty"`
|
||||
PrimaryResetAfterSeconds *int `json:"primary_reset_after_seconds,omitempty"`
|
||||
PrimaryWindowMinutes *int `json:"primary_window_minutes,omitempty"`
|
||||
SecondaryUsedPercent *float64 `json:"secondary_used_percent,omitempty"`
|
||||
SecondaryResetAfterSeconds *int `json:"secondary_reset_after_seconds,omitempty"`
|
||||
SecondaryWindowMinutes *int `json:"secondary_window_minutes,omitempty"`
|
||||
PrimaryOverSecondaryPercent *float64 `json:"primary_over_secondary_percent,omitempty"`
|
||||
UpdatedAt string `json:"updated_at,omitempty"`
|
||||
}
|
||||
|
||||
// OpenAIUsage represents OpenAI API response usage
|
||||
type OpenAIUsage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
@@ -284,6 +297,13 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
}
|
||||
}
|
||||
|
||||
// Extract and save Codex usage snapshot from response headers (for OAuth accounts)
|
||||
if account.Type == model.AccountTypeOAuth {
|
||||
if snapshot := extractCodexUsageHeaders(resp.Header); snapshot != nil {
|
||||
s.updateCodexUsageSnapshot(ctx, account.ID, snapshot)
|
||||
}
|
||||
}
|
||||
|
||||
return &OpenAIForwardResult{
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
Usage: *usage,
|
||||
@@ -708,3 +728,109 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// extractCodexUsageHeaders extracts Codex usage limits from response headers
|
||||
func extractCodexUsageHeaders(headers http.Header) *OpenAICodexUsageSnapshot {
|
||||
snapshot := &OpenAICodexUsageSnapshot{}
|
||||
hasData := false
|
||||
|
||||
// Helper to parse float64 from header
|
||||
parseFloat := func(key string) *float64 {
|
||||
if v := headers.Get(key); v != "" {
|
||||
if f, err := strconv.ParseFloat(v, 64); err == nil {
|
||||
return &f
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Helper to parse int from header
|
||||
parseInt := func(key string) *int {
|
||||
if v := headers.Get(key); v != "" {
|
||||
if i, err := strconv.Atoi(v); err == nil {
|
||||
return &i
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Primary (weekly) limits
|
||||
if v := parseFloat("x-codex-primary-used-percent"); v != nil {
|
||||
snapshot.PrimaryUsedPercent = v
|
||||
hasData = true
|
||||
}
|
||||
if v := parseInt("x-codex-primary-reset-after-seconds"); v != nil {
|
||||
snapshot.PrimaryResetAfterSeconds = v
|
||||
hasData = true
|
||||
}
|
||||
if v := parseInt("x-codex-primary-window-minutes"); v != nil {
|
||||
snapshot.PrimaryWindowMinutes = v
|
||||
hasData = true
|
||||
}
|
||||
|
||||
// Secondary (5h) limits
|
||||
if v := parseFloat("x-codex-secondary-used-percent"); v != nil {
|
||||
snapshot.SecondaryUsedPercent = v
|
||||
hasData = true
|
||||
}
|
||||
if v := parseInt("x-codex-secondary-reset-after-seconds"); v != nil {
|
||||
snapshot.SecondaryResetAfterSeconds = v
|
||||
hasData = true
|
||||
}
|
||||
if v := parseInt("x-codex-secondary-window-minutes"); v != nil {
|
||||
snapshot.SecondaryWindowMinutes = v
|
||||
hasData = true
|
||||
}
|
||||
|
||||
// Overflow ratio
|
||||
if v := parseFloat("x-codex-primary-over-secondary-limit-percent"); v != nil {
|
||||
snapshot.PrimaryOverSecondaryPercent = v
|
||||
hasData = true
|
||||
}
|
||||
|
||||
if !hasData {
|
||||
return nil
|
||||
}
|
||||
|
||||
snapshot.UpdatedAt = time.Now().Format(time.RFC3339)
|
||||
return snapshot
|
||||
}
|
||||
|
||||
// updateCodexUsageSnapshot saves the Codex usage snapshot to account's Extra field
|
||||
func (s *OpenAIGatewayService) updateCodexUsageSnapshot(ctx context.Context, accountID int64, snapshot *OpenAICodexUsageSnapshot) {
|
||||
if snapshot == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert snapshot to map for merging into Extra
|
||||
updates := make(map[string]any)
|
||||
if snapshot.PrimaryUsedPercent != nil {
|
||||
updates["codex_primary_used_percent"] = *snapshot.PrimaryUsedPercent
|
||||
}
|
||||
if snapshot.PrimaryResetAfterSeconds != nil {
|
||||
updates["codex_primary_reset_after_seconds"] = *snapshot.PrimaryResetAfterSeconds
|
||||
}
|
||||
if snapshot.PrimaryWindowMinutes != nil {
|
||||
updates["codex_primary_window_minutes"] = *snapshot.PrimaryWindowMinutes
|
||||
}
|
||||
if snapshot.SecondaryUsedPercent != nil {
|
||||
updates["codex_secondary_used_percent"] = *snapshot.SecondaryUsedPercent
|
||||
}
|
||||
if snapshot.SecondaryResetAfterSeconds != nil {
|
||||
updates["codex_secondary_reset_after_seconds"] = *snapshot.SecondaryResetAfterSeconds
|
||||
}
|
||||
if snapshot.SecondaryWindowMinutes != nil {
|
||||
updates["codex_secondary_window_minutes"] = *snapshot.SecondaryWindowMinutes
|
||||
}
|
||||
if snapshot.PrimaryOverSecondaryPercent != nil {
|
||||
updates["codex_primary_over_secondary_percent"] = *snapshot.PrimaryOverSecondaryPercent
|
||||
}
|
||||
updates["codex_usage_updated_at"] = snapshot.UpdatedAt
|
||||
|
||||
// Update account's Extra field asynchronously
|
||||
go func() {
|
||||
updateCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = s.accountRepo.UpdateExtra(updateCtx, accountID, updates)
|
||||
}()
|
||||
}
|
||||
|
||||
@@ -34,4 +34,5 @@ type AccountRepository interface {
|
||||
SetOverloaded(ctx context.Context, id int64, until time.Time) error
|
||||
ClearRateLimit(ctx context.Context, id int64) error
|
||||
UpdateSessionWindow(ctx context.Context, id int64, start, end *time.Time, status string) error
|
||||
UpdateExtra(ctx context.Context, id int64, updates map[string]any) error
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user