Merge origin/main into codex/fix-native-responses-namespace

Resolve conflict in openai_gateway_passthrough.go streaming path: keep
main's normalizeCompletedImageGenerationStatus normalization ahead of
this branch's namespace restore block, mirroring the established order
in openai_gateway_response_handling.go.
This commit is contained in:
shaw
2026-07-14 14:45:04 +08:00
310 changed files with 18489 additions and 1087 deletions
+9
View File
@@ -8,6 +8,15 @@ permissions:
contents: read
jobs:
shell:
runs-on: macos-15
steps:
- uses: actions/checkout@v6
- name: Check deployment scripts
run: |
/bin/bash -n deploy/apple-container.sh
/bin/bash deploy/tests/apple-container-test.sh
test:
runs-on: ubuntu-latest
steps:
+2
View File
@@ -116,6 +116,8 @@ backend/.installed
# 其他
# ===================
tests
!deploy/tests/
!deploy/tests/**
CLAUDE.md
.claude
scripts
+40 -2
View File
@@ -329,6 +329,7 @@ cd sub2api/deploy
# 2. Copy environment configuration
cp .env.example .env
chmod 600 .env
# 3. Edit configuration (generate secure passwords)
nano .env
@@ -448,7 +449,23 @@ rm -rf data/ postgres_data/ redis_data/
---
### Method 3: Build from Source
### Method 3: Apple container (macOS)
Apple-silicon Macs running macOS 26 can run the full Sub2API, PostgreSQL, and Redis stack with Apple `container` 1.1.0 or newer:
```bash
git clone https://github.com/Wei-Shaw/sub2api.git
cd sub2api/deploy
./apple-container.sh init
./apple-container.sh up
./apple-container.sh status
```
This is an operator-managed local workflow; Docker Compose remains the recommended production path. See [deploy/APPLE_CONTAINER.md](deploy/APPLE_CONTAINER.md) for lifecycle commands, persistence, upgrades, and runtime limitations.
---
### Method 4: Build from Source
Build and run from source code for development or customization.
@@ -579,6 +596,27 @@ If you disable URL validation or response header filtering, harden your network
- Enforce TLS-only outbound traffic
- Strip sensitive upstream response headers at the proxy
#### OpenAI Responses WebSocket ingress limits
`gateway.openai_ws` bounds the lifetime and aggregate count of client-facing
Responses WebSocket sessions. These safeguards apply independently from
per-turn user and account concurrency slots, which are released between turns.
```yaml
gateway:
openai_ws:
# Close a client socket idle between completed turns; 0 disables this safeguard.
ingress_inter_turn_idle_timeout_seconds: 300
# Distributed API-key limit for live client ingress sessions; 0 disables it.
max_ingress_connections_per_api_key: 64
```
The connection cap is coordinated through Redis using a 60-second lease that
is refreshed every 20 seconds. A process that cannot confirm a lease for a
full lease lifetime closes its local WebSocket rather than continuing outside
the global cap. Use `http_bridge` for client-WebSocket/upstream-HTTP operation
when rolling out or mitigating upstream WebSocket issues.
#### ⚠️ Important: Creating the Admin Account
The initial admin account is **only created via the setup wizard** (served at `http://<host>:8080` on first run). The `default.admin_email` / `default.admin_password` fields in `config.yaml` are **not used** to create it — they exist in the template for historical reasons.
@@ -650,7 +688,7 @@ Sub2API supports both Grok subscription accounts through xAI OAuth and standard
- Public Chat Completions targets: `/v1/chat/completions` and `/chat/completions`, forwarded to the account-type-specific xAI upstream
- Codex CLI style Responses WebSocket ingress is accepted on the Responses targets and bridged to xAI HTTP/SSE Responses upstream
- Text models: `grok-4.5`, `grok-4.3`, `grok-build-0.1`, `grok-composer-2.5-fast`, `grok-4.20-0309-reasoning`, `grok-4.20-0309-non-reasoning`, and `grok-4.20-multi-agent-0309`
- Media targets for Grok groups: `/v1/images/generations`, `/images/generations`, `/v1/images/edits`, `/images/edits`, `/v1/videos/generations`, `/videos/generations`, `/v1/videos/{request_id}`, and `/videos/{request_id}`. Generation requests require the group image-generation permission.
- Media targets for Grok groups: `/v1/images/generations`, `/images/generations`, `/v1/images/edits`, `/images/edits`, `/v1/videos/generations`, `/videos/generations`, `/v1/videos/edits`, `/videos/edits`, `/v1/videos/extensions`, `/videos/extensions`, `/v1/videos/{request_id}`, and `/videos/{request_id}`. Generation, editing, and extension requests require the group image-generation permission.
- Media models: `grok-imagine`, `grok-imagine-image-quality`, `grok-imagine-image`, `grok-imagine-edit`, `grok-imagine-video`, and `grok-imagine-video-1.5`
- Out of scope for this provider: TTS, transcription, browser automation, cookies, and Grok web scraping
+18 -1
View File
@@ -333,6 +333,7 @@ cd sub2api/deploy
# 2. 复制环境配置文件
cp .env.example .env
chmod 600 .env
# 3. 编辑配置(生成安全密码)
nano .env
@@ -464,7 +465,23 @@ rm -rf data/ postgres_data/ redis_data/
---
### 方式三:源码编译
### 方式三:Apple containermacOS
Apple 芯片 Mac 在 macOS 26 上可使用 Apple `container` 1.1.0 或更高版本运行完整的 Sub2API、PostgreSQL 和 Redis
```bash
git clone https://github.com/Wei-Shaw/sub2api.git
cd sub2api/deploy
./apple-container.sh init
./apple-container.sh up
./apple-container.sh status
```
该方式面向本地开发和人工运维,不提供持续重启监管;生产部署仍推荐 Docker Compose。生命周期命令、持久化、升级和运行时限制见 [deploy/APPLE_CONTAINER.md](deploy/APPLE_CONTAINER.md)。
---
### 方式四:源码编译
从源码编译安装,适合开发或定制需求。
+18 -1
View File
@@ -327,6 +327,7 @@ cd sub2api/deploy
# 2. 環境設定ファイルをコピー
cp .env.example .env
chmod 600 .env
# 3. 設定を編集(セキュアなパスワードを生成)
nano .env
@@ -446,7 +447,23 @@ rm -rf data/ postgres_data/ redis_data/
---
### 方法3: ソースからビルド
### 方法3: Apple containermacOS
Apple シリコン搭載 Mac と macOS 26 では、Apple `container` 1.1.0 以降を使用して Sub2API、PostgreSQL、Redis の完全なスタックを実行できます:
```bash
git clone https://github.com/Wei-Shaw/sub2api.git
cd sub2api/deploy
./apple-container.sh init
./apple-container.sh up
./apple-container.sh status
```
これはローカル開発および手動運用向けです。本番環境では引き続き Docker Compose を推奨します。ライフサイクル、永続化、アップグレード、制限については [deploy/APPLE_CONTAINER.md](deploy/APPLE_CONTAINER.md) を参照してください。
---
### 方法4: ソースからビルド
開発やカスタマイズのためにソースコードからビルドして実行します。
+1 -1
View File
@@ -1 +1 @@
0.1.152
0.1.153
+4 -4
View File
@@ -190,12 +190,13 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream)
antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository)
grokQuotaFetcher := service.NewGrokQuotaFetcher()
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, usageLogRepository)
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory)
usageCache := service.NewUsageCache()
accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService)
accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService)
accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService)
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator)
accountHandler := admin.ProvideAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, grokQuotaService)
adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService)
dataManagementService := service.NewDataManagementService()
dataManagementHandler := admin.NewDataManagementHandler(dataManagementService)
@@ -207,7 +208,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
openAIOAuthHandler := admin.NewOpenAIOAuthHandler(openAIOAuthService, adminService, openAIQuotaService)
geminiOAuthHandler := admin.NewGeminiOAuthHandler(geminiOAuthService)
antigravityOAuthHandler := admin.NewAntigravityOAuthHandler(antigravityOAuthService)
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream)
grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService, grokQuotaService)
proxyHandler := admin.NewProxyHandler(adminService)
adminRedeemHandler := admin.NewRedeemHandler(adminService, redeemService)
@@ -264,7 +264,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, configConfig)
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
totpHandler := handler.NewTotpHandler(totpService)
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService, channelService)
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService)
paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry)
availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService)
batchImageHandler := handler.NewBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService)
+2 -1
View File
@@ -167,6 +167,7 @@ const (
ProviderOpenai Provider = "openai"
ProviderAnthropic Provider = "anthropic"
ProviderGemini Provider = "gemini"
ProviderGrok Provider = "grok"
)
func (pr Provider) String() string {
@@ -176,7 +177,7 @@ func (pr Provider) String() string {
// ProviderValidator is a validator for the "provider" field enum values. It is called by the builders before save.
func ProviderValidator(pr Provider) error {
switch pr {
case ProviderOpenai, ProviderAnthropic, ProviderGemini:
case ProviderOpenai, ProviderAnthropic, ProviderGemini, ProviderGrok:
return nil
default:
return fmt.Errorf("channelmonitor: invalid enum value for provider field: %q", pr)
@@ -103,6 +103,7 @@ const (
ProviderOpenai Provider = "openai"
ProviderAnthropic Provider = "anthropic"
ProviderGemini Provider = "gemini"
ProviderGrok Provider = "grok"
)
func (pr Provider) String() string {
@@ -112,7 +113,7 @@ func (pr Provider) String() string {
// ProviderValidator is a validator for the "provider" field enum values. It is called by the builders before save.
func ProviderValidator(pr Provider) error {
switch pr {
case ProviderOpenai, ProviderAnthropic, ProviderGemini:
case ProviderOpenai, ProviderAnthropic, ProviderGemini, ProviderGrok:
return nil
default:
return fmt.Errorf("channelmonitorrequesttemplate: invalid enum value for provider field: %q", pr)
+17 -16
View File
@@ -623,7 +623,7 @@ var (
{Name: "created_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}},
{Name: "updated_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}},
{Name: "name", Type: field.TypeString, Size: 100},
{Name: "provider", Type: field.TypeEnum, Enums: []string{"openai", "anthropic", "gemini"}},
{Name: "provider", Type: field.TypeEnum, Enums: []string{"openai", "anthropic", "gemini", "grok"}},
{Name: "api_mode", Type: field.TypeString, Size: 32, Default: "chat_completions"},
{Name: "endpoint", Type: field.TypeString, Size: 500},
{Name: "api_key_encrypted", Type: field.TypeString},
@@ -768,7 +768,7 @@ var (
{Name: "created_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}},
{Name: "updated_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}},
{Name: "name", Type: field.TypeString, Size: 100},
{Name: "provider", Type: field.TypeEnum, Enums: []string{"openai", "anthropic", "gemini"}},
{Name: "provider", Type: field.TypeEnum, Enums: []string{"openai", "anthropic", "gemini", "grok"}},
{Name: "api_mode", Type: field.TypeString, Size: 32, Default: "chat_completions"},
{Name: "description", Type: field.TypeString, Nullable: true, Size: 500, Default: ""},
{Name: "extra_headers", Type: field.TypeJSON},
@@ -1560,6 +1560,7 @@ var (
{Name: "total_cost", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,10)"}},
{Name: "actual_cost", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,10)"}},
{Name: "rate_multiplier", Type: field.TypeFloat64, Default: 1, SchemaType: map[string]string{"postgres": "decimal(10,4)"}},
{Name: "long_context_billing_applied", Type: field.TypeBool, Default: false},
{Name: "account_rate_multiplier", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(10,4)"}},
{Name: "billing_type", Type: field.TypeInt8, Default: 0},
{Name: "stream", Type: field.TypeBool, Default: false},
@@ -1592,31 +1593,31 @@ var (
ForeignKeys: []*schema.ForeignKey{
{
Symbol: "usage_logs_api_keys_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[40]},
Columns: []*schema.Column{UsageLogsColumns[41]},
RefColumns: []*schema.Column{APIKeysColumns[0]},
OnDelete: schema.NoAction,
},
{
Symbol: "usage_logs_accounts_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[41]},
Columns: []*schema.Column{UsageLogsColumns[42]},
RefColumns: []*schema.Column{AccountsColumns[0]},
OnDelete: schema.NoAction,
},
{
Symbol: "usage_logs_groups_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[42]},
Columns: []*schema.Column{UsageLogsColumns[43]},
RefColumns: []*schema.Column{GroupsColumns[0]},
OnDelete: schema.SetNull,
},
{
Symbol: "usage_logs_users_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[43]},
Columns: []*schema.Column{UsageLogsColumns[44]},
RefColumns: []*schema.Column{UsersColumns[0]},
OnDelete: schema.NoAction,
},
{
Symbol: "usage_logs_user_subscriptions_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[44]},
Columns: []*schema.Column{UsageLogsColumns[45]},
RefColumns: []*schema.Column{UserSubscriptionsColumns[0]},
OnDelete: schema.SetNull,
},
@@ -1625,32 +1626,32 @@ var (
{
Name: "usagelog_user_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[43]},
Columns: []*schema.Column{UsageLogsColumns[44]},
},
{
Name: "usagelog_api_key_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[40]},
Columns: []*schema.Column{UsageLogsColumns[41]},
},
{
Name: "usagelog_account_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[41]},
Columns: []*schema.Column{UsageLogsColumns[42]},
},
{
Name: "usagelog_group_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[42]},
Columns: []*schema.Column{UsageLogsColumns[43]},
},
{
Name: "usagelog_subscription_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[44]},
Columns: []*schema.Column{UsageLogsColumns[45]},
},
{
Name: "usagelog_created_at",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[39]},
Columns: []*schema.Column{UsageLogsColumns[40]},
},
{
Name: "usagelog_model",
@@ -1670,17 +1671,17 @@ var (
{
Name: "usagelog_user_id_created_at",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[43], UsageLogsColumns[39]},
Columns: []*schema.Column{UsageLogsColumns[44], UsageLogsColumns[40]},
},
{
Name: "usagelog_api_key_id_created_at",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[40], UsageLogsColumns[39]},
Columns: []*schema.Column{UsageLogsColumns[41], UsageLogsColumns[40]},
},
{
Name: "usagelog_group_id_created_at",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[42], UsageLogsColumns[39]},
Columns: []*schema.Column{UsageLogsColumns[43], UsageLogsColumns[40]},
},
},
}
+132 -78
View File
@@ -41763,83 +41763,84 @@ func (m *UsageCleanupTaskMutation) ResetEdge(name string) error {
// UsageLogMutation represents an operation that mutates the UsageLog nodes in the graph.
type UsageLogMutation struct {
config
op Op
typ string
id *int64
request_id *string
model *string
requested_model *string
upstream_model *string
channel_id *int64
addchannel_id *int64
model_mapping_chain *string
billing_tier *string
billing_mode *string
input_tokens *int
addinput_tokens *int
output_tokens *int
addoutput_tokens *int
cache_creation_tokens *int
addcache_creation_tokens *int
cache_read_tokens *int
addcache_read_tokens *int
cache_creation_5m_tokens *int
addcache_creation_5m_tokens *int
cache_creation_1h_tokens *int
addcache_creation_1h_tokens *int
input_cost *float64
addinput_cost *float64
output_cost *float64
addoutput_cost *float64
cache_creation_cost *float64
addcache_creation_cost *float64
cache_read_cost *float64
addcache_read_cost *float64
total_cost *float64
addtotal_cost *float64
actual_cost *float64
addactual_cost *float64
rate_multiplier *float64
addrate_multiplier *float64
account_rate_multiplier *float64
addaccount_rate_multiplier *float64
billing_type *int8
addbilling_type *int8
stream *bool
duration_ms *int
addduration_ms *int
first_token_ms *int
addfirst_token_ms *int
user_agent *string
ip_address *string
image_count *int
addimage_count *int
image_size *string
image_input_size *string
image_output_size *string
image_size_source *string
image_size_breakdown *map[string]int
video_count *int
addvideo_count *int
video_resolution *string
video_duration_seconds *int
addvideo_duration_seconds *int
cache_ttl_overridden *bool
created_at *time.Time
clearedFields map[string]struct{}
user *int64
cleareduser bool
api_key *int64
clearedapi_key bool
account *int64
clearedaccount bool
group *int64
clearedgroup bool
subscription *int64
clearedsubscription bool
done bool
oldValue func(context.Context) (*UsageLog, error)
predicates []predicate.UsageLog
op Op
typ string
id *int64
request_id *string
model *string
requested_model *string
upstream_model *string
channel_id *int64
addchannel_id *int64
model_mapping_chain *string
billing_tier *string
billing_mode *string
input_tokens *int
addinput_tokens *int
output_tokens *int
addoutput_tokens *int
cache_creation_tokens *int
addcache_creation_tokens *int
cache_read_tokens *int
addcache_read_tokens *int
cache_creation_5m_tokens *int
addcache_creation_5m_tokens *int
cache_creation_1h_tokens *int
addcache_creation_1h_tokens *int
input_cost *float64
addinput_cost *float64
output_cost *float64
addoutput_cost *float64
cache_creation_cost *float64
addcache_creation_cost *float64
cache_read_cost *float64
addcache_read_cost *float64
total_cost *float64
addtotal_cost *float64
actual_cost *float64
addactual_cost *float64
rate_multiplier *float64
addrate_multiplier *float64
long_context_billing_applied *bool
account_rate_multiplier *float64
addaccount_rate_multiplier *float64
billing_type *int8
addbilling_type *int8
stream *bool
duration_ms *int
addduration_ms *int
first_token_ms *int
addfirst_token_ms *int
user_agent *string
ip_address *string
image_count *int
addimage_count *int
image_size *string
image_input_size *string
image_output_size *string
image_size_source *string
image_size_breakdown *map[string]int
video_count *int
addvideo_count *int
video_resolution *string
video_duration_seconds *int
addvideo_duration_seconds *int
cache_ttl_overridden *bool
created_at *time.Time
clearedFields map[string]struct{}
user *int64
cleareduser bool
api_key *int64
clearedapi_key bool
account *int64
clearedaccount bool
group *int64
clearedgroup bool
subscription *int64
clearedsubscription bool
done bool
oldValue func(context.Context) (*UsageLog, error)
predicates []predicate.UsageLog
}
var _ ent.Mutation = (*UsageLogMutation)(nil)
@@ -43261,6 +43262,42 @@ func (m *UsageLogMutation) ResetRateMultiplier() {
m.addrate_multiplier = nil
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (m *UsageLogMutation) SetLongContextBillingApplied(b bool) {
m.long_context_billing_applied = &b
}
// LongContextBillingApplied returns the value of the "long_context_billing_applied" field in the mutation.
func (m *UsageLogMutation) LongContextBillingApplied() (r bool, exists bool) {
v := m.long_context_billing_applied
if v == nil {
return
}
return *v, true
}
// OldLongContextBillingApplied returns the old "long_context_billing_applied" 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) OldLongContextBillingApplied(ctx context.Context) (v bool, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldLongContextBillingApplied is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldLongContextBillingApplied requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldLongContextBillingApplied: %w", err)
}
return oldValue.LongContextBillingApplied, nil
}
// ResetLongContextBillingApplied resets all changes to the "long_context_billing_applied" field.
func (m *UsageLogMutation) ResetLongContextBillingApplied() {
m.long_context_billing_applied = nil
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (m *UsageLogMutation) SetAccountRateMultiplier(f float64) {
m.account_rate_multiplier = &f
@@ -44378,7 +44415,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, 44)
fields := make([]string, 0, 45)
if m.user != nil {
fields = append(fields, usagelog.FieldUserID)
}
@@ -44457,6 +44494,9 @@ func (m *UsageLogMutation) Fields() []string {
if m.rate_multiplier != nil {
fields = append(fields, usagelog.FieldRateMultiplier)
}
if m.long_context_billing_applied != nil {
fields = append(fields, usagelog.FieldLongContextBillingApplied)
}
if m.account_rate_multiplier != nil {
fields = append(fields, usagelog.FieldAccountRateMultiplier)
}
@@ -44571,6 +44611,8 @@ func (m *UsageLogMutation) Field(name string) (ent.Value, bool) {
return m.ActualCost()
case usagelog.FieldRateMultiplier:
return m.RateMultiplier()
case usagelog.FieldLongContextBillingApplied:
return m.LongContextBillingApplied()
case usagelog.FieldAccountRateMultiplier:
return m.AccountRateMultiplier()
case usagelog.FieldBillingType:
@@ -44668,6 +44710,8 @@ func (m *UsageLogMutation) OldField(ctx context.Context, name string) (ent.Value
return m.OldActualCost(ctx)
case usagelog.FieldRateMultiplier:
return m.OldRateMultiplier(ctx)
case usagelog.FieldLongContextBillingApplied:
return m.OldLongContextBillingApplied(ctx)
case usagelog.FieldAccountRateMultiplier:
return m.OldAccountRateMultiplier(ctx)
case usagelog.FieldBillingType:
@@ -44895,6 +44939,13 @@ func (m *UsageLogMutation) SetField(name string, value ent.Value) error {
}
m.SetRateMultiplier(v)
return nil
case usagelog.FieldLongContextBillingApplied:
v, ok := value.(bool)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetLongContextBillingApplied(v)
return nil
case usagelog.FieldAccountRateMultiplier:
v, ok := value.(float64)
if !ok {
@@ -45526,6 +45577,9 @@ func (m *UsageLogMutation) ResetField(name string) error {
case usagelog.FieldRateMultiplier:
m.ResetRateMultiplier()
return nil
case usagelog.FieldLongContextBillingApplied:
m.ResetLongContextBillingApplied()
return nil
case usagelog.FieldAccountRateMultiplier:
m.ResetAccountRateMultiplier()
return nil
+17 -13
View File
@@ -1940,56 +1940,60 @@ func init() {
usagelogDescRateMultiplier := usagelogFields[25].Descriptor()
// usagelog.DefaultRateMultiplier holds the default value on creation for the rate_multiplier field.
usagelog.DefaultRateMultiplier = usagelogDescRateMultiplier.Default.(float64)
// usagelogDescLongContextBillingApplied is the schema descriptor for long_context_billing_applied field.
usagelogDescLongContextBillingApplied := usagelogFields[26].Descriptor()
// usagelog.DefaultLongContextBillingApplied holds the default value on creation for the long_context_billing_applied field.
usagelog.DefaultLongContextBillingApplied = usagelogDescLongContextBillingApplied.Default.(bool)
// usagelogDescBillingType is the schema descriptor for billing_type field.
usagelogDescBillingType := usagelogFields[27].Descriptor()
usagelogDescBillingType := usagelogFields[28].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[28].Descriptor()
usagelogDescStream := usagelogFields[29].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[31].Descriptor()
usagelogDescUserAgent := usagelogFields[32].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[32].Descriptor()
usagelogDescIPAddress := usagelogFields[33].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[33].Descriptor()
usagelogDescImageCount := usagelogFields[34].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[34].Descriptor()
usagelogDescImageSize := usagelogFields[35].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)
// usagelogDescImageInputSize is the schema descriptor for image_input_size field.
usagelogDescImageInputSize := usagelogFields[35].Descriptor()
usagelogDescImageInputSize := usagelogFields[36].Descriptor()
// usagelog.ImageInputSizeValidator is a validator for the "image_input_size" field. It is called by the builders before save.
usagelog.ImageInputSizeValidator = usagelogDescImageInputSize.Validators[0].(func(string) error)
// usagelogDescImageOutputSize is the schema descriptor for image_output_size field.
usagelogDescImageOutputSize := usagelogFields[36].Descriptor()
usagelogDescImageOutputSize := usagelogFields[37].Descriptor()
// usagelog.ImageOutputSizeValidator is a validator for the "image_output_size" field. It is called by the builders before save.
usagelog.ImageOutputSizeValidator = usagelogDescImageOutputSize.Validators[0].(func(string) error)
// usagelogDescImageSizeSource is the schema descriptor for image_size_source field.
usagelogDescImageSizeSource := usagelogFields[37].Descriptor()
usagelogDescImageSizeSource := usagelogFields[38].Descriptor()
// usagelog.ImageSizeSourceValidator is a validator for the "image_size_source" field. It is called by the builders before save.
usagelog.ImageSizeSourceValidator = usagelogDescImageSizeSource.Validators[0].(func(string) error)
// usagelogDescVideoCount is the schema descriptor for video_count field.
usagelogDescVideoCount := usagelogFields[39].Descriptor()
usagelogDescVideoCount := usagelogFields[40].Descriptor()
// usagelog.DefaultVideoCount holds the default value on creation for the video_count field.
usagelog.DefaultVideoCount = usagelogDescVideoCount.Default.(int)
// usagelogDescVideoResolution is the schema descriptor for video_resolution field.
usagelogDescVideoResolution := usagelogFields[40].Descriptor()
usagelogDescVideoResolution := usagelogFields[41].Descriptor()
// usagelog.VideoResolutionValidator is a validator for the "video_resolution" field. It is called by the builders before save.
usagelog.VideoResolutionValidator = usagelogDescVideoResolution.Validators[0].(func(string) error)
// usagelogDescCacheTTLOverridden is the schema descriptor for cache_ttl_overridden field.
usagelogDescCacheTTLOverridden := usagelogFields[42].Descriptor()
usagelogDescCacheTTLOverridden := usagelogFields[43].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[43].Descriptor()
usagelogDescCreatedAt := usagelogFields[44].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()
+1 -1
View File
@@ -35,7 +35,7 @@ func (ChannelMonitor) Fields() []ent.Field {
NotEmpty().
MaxLen(100),
field.Enum("provider").
Values("openai", "anthropic", "gemini"),
Values("openai", "anthropic", "gemini", "grok"),
field.String("api_mode").
Default("chat_completions").
MaxLen(32).
@@ -39,7 +39,7 @@ func (ChannelMonitorRequestTemplate) Fields() []ent.Field {
NotEmpty().
MaxLen(100),
field.Enum("provider").
Values("openai", "anthropic", "gemini"),
Values("openai", "anthropic", "gemini", "grok"),
field.String("api_mode").
Default("chat_completions").
MaxLen(32).
+3
View File
@@ -100,6 +100,9 @@ func (UsageLog) Fields() []ent.Field {
field.Float("rate_multiplier").
Default(1).
SchemaType(map[string]string{dialect.Postgres: "decimal(10,4)"}),
field.Bool("long_context_billing_applied").
Default(false).
Comment("Whether long-context pricing changed token prices for this request"),
// account_rate_multiplier: 账号计费倍率快照(NULL 表示按 1.0 处理)
field.Float("account_rate_multiplier").
+12 -1
View File
@@ -75,6 +75,8 @@ type UsageLog struct {
ActualCost float64 `json:"actual_cost,omitempty"`
// RateMultiplier holds the value of the "rate_multiplier" field.
RateMultiplier float64 `json:"rate_multiplier,omitempty"`
// Whether long-context pricing changed token prices for this request
LongContextBillingApplied bool `json:"long_context_billing_applied,omitempty"`
// AccountRateMultiplier holds the value of the "account_rate_multiplier" field.
AccountRateMultiplier *float64 `json:"account_rate_multiplier,omitempty"`
// BillingType holds the value of the "billing_type" field.
@@ -196,7 +198,7 @@ func (*UsageLog) scanValues(columns []string) ([]any, error) {
switch columns[i] {
case usagelog.FieldImageSizeBreakdown:
values[i] = new([]byte)
case usagelog.FieldStream, usagelog.FieldCacheTTLOverridden:
case usagelog.FieldLongContextBillingApplied, usagelog.FieldStream, usagelog.FieldCacheTTLOverridden:
values[i] = new(sql.NullBool)
case usagelog.FieldInputCost, usagelog.FieldOutputCost, usagelog.FieldCacheCreationCost, usagelog.FieldCacheReadCost, usagelog.FieldTotalCost, usagelog.FieldActualCost, usagelog.FieldRateMultiplier, usagelog.FieldAccountRateMultiplier:
values[i] = new(sql.NullFloat64)
@@ -391,6 +393,12 @@ func (_m *UsageLog) assignValues(columns []string, values []any) error {
} else if value.Valid {
_m.RateMultiplier = value.Float64
}
case usagelog.FieldLongContextBillingApplied:
if value, ok := values[i].(*sql.NullBool); !ok {
return fmt.Errorf("unexpected type %T for field long_context_billing_applied", values[i])
} else if value.Valid {
_m.LongContextBillingApplied = value.Bool
}
case usagelog.FieldAccountRateMultiplier:
if value, ok := values[i].(*sql.NullFloat64); !ok {
return fmt.Errorf("unexpected type %T for field account_rate_multiplier", values[i])
@@ -667,6 +675,9 @@ func (_m *UsageLog) String() string {
builder.WriteString("rate_multiplier=")
builder.WriteString(fmt.Sprintf("%v", _m.RateMultiplier))
builder.WriteString(", ")
builder.WriteString("long_context_billing_applied=")
builder.WriteString(fmt.Sprintf("%v", _m.LongContextBillingApplied))
builder.WriteString(", ")
if v := _m.AccountRateMultiplier; v != nil {
builder.WriteString("account_rate_multiplier=")
builder.WriteString(fmt.Sprintf("%v", *v))
+10
View File
@@ -66,6 +66,8 @@ const (
FieldActualCost = "actual_cost"
// FieldRateMultiplier holds the string denoting the rate_multiplier field in the database.
FieldRateMultiplier = "rate_multiplier"
// FieldLongContextBillingApplied holds the string denoting the long_context_billing_applied field in the database.
FieldLongContextBillingApplied = "long_context_billing_applied"
// FieldAccountRateMultiplier holds the string denoting the account_rate_multiplier field in the database.
FieldAccountRateMultiplier = "account_rate_multiplier"
// FieldBillingType holds the string denoting the billing_type field in the database.
@@ -180,6 +182,7 @@ var Columns = []string{
FieldTotalCost,
FieldActualCost,
FieldRateMultiplier,
FieldLongContextBillingApplied,
FieldAccountRateMultiplier,
FieldBillingType,
FieldStream,
@@ -251,6 +254,8 @@ var (
DefaultActualCost float64
// DefaultRateMultiplier holds the default value on creation for the "rate_multiplier" field.
DefaultRateMultiplier float64
// DefaultLongContextBillingApplied holds the default value on creation for the "long_context_billing_applied" field.
DefaultLongContextBillingApplied bool
// DefaultBillingType holds the default value on creation for the "billing_type" field.
DefaultBillingType int8
// DefaultStream holds the default value on creation for the "stream" field.
@@ -417,6 +422,11 @@ func ByRateMultiplier(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldRateMultiplier, opts...).ToFunc()
}
// ByLongContextBillingApplied orders the results by the long_context_billing_applied field.
func ByLongContextBillingApplied(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldLongContextBillingApplied, opts...).ToFunc()
}
// ByAccountRateMultiplier orders the results by the account_rate_multiplier field.
func ByAccountRateMultiplier(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldAccountRateMultiplier, opts...).ToFunc()
+15
View File
@@ -185,6 +185,11 @@ func RateMultiplier(v float64) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldRateMultiplier, v))
}
// LongContextBillingApplied applies equality check predicate on the "long_context_billing_applied" field. It's identical to LongContextBillingAppliedEQ.
func LongContextBillingApplied(v bool) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldLongContextBillingApplied, v))
}
// AccountRateMultiplier applies equality check predicate on the "account_rate_multiplier" field. It's identical to AccountRateMultiplierEQ.
func AccountRateMultiplier(v float64) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldAccountRateMultiplier, v))
@@ -1465,6 +1470,16 @@ func RateMultiplierLTE(v float64) predicate.UsageLog {
return predicate.UsageLog(sql.FieldLTE(FieldRateMultiplier, v))
}
// LongContextBillingAppliedEQ applies the EQ predicate on the "long_context_billing_applied" field.
func LongContextBillingAppliedEQ(v bool) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldLongContextBillingApplied, v))
}
// LongContextBillingAppliedNEQ applies the NEQ predicate on the "long_context_billing_applied" field.
func LongContextBillingAppliedNEQ(v bool) predicate.UsageLog {
return predicate.UsageLog(sql.FieldNEQ(FieldLongContextBillingApplied, v))
}
// AccountRateMultiplierEQ applies the EQ predicate on the "account_rate_multiplier" field.
func AccountRateMultiplierEQ(v float64) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldAccountRateMultiplier, v))
+65
View File
@@ -351,6 +351,20 @@ func (_c *UsageLogCreate) SetNillableRateMultiplier(v *float64) *UsageLogCreate
return _c
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (_c *UsageLogCreate) SetLongContextBillingApplied(v bool) *UsageLogCreate {
_c.mutation.SetLongContextBillingApplied(v)
return _c
}
// SetNillableLongContextBillingApplied sets the "long_context_billing_applied" field if the given value is not nil.
func (_c *UsageLogCreate) SetNillableLongContextBillingApplied(v *bool) *UsageLogCreate {
if v != nil {
_c.SetLongContextBillingApplied(*v)
}
return _c
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (_c *UsageLogCreate) SetAccountRateMultiplier(v float64) *UsageLogCreate {
_c.mutation.SetAccountRateMultiplier(v)
@@ -707,6 +721,10 @@ func (_c *UsageLogCreate) defaults() {
v := usagelog.DefaultRateMultiplier
_c.mutation.SetRateMultiplier(v)
}
if _, ok := _c.mutation.LongContextBillingApplied(); !ok {
v := usagelog.DefaultLongContextBillingApplied
_c.mutation.SetLongContextBillingApplied(v)
}
if _, ok := _c.mutation.BillingType(); !ok {
v := usagelog.DefaultBillingType
_c.mutation.SetBillingType(v)
@@ -824,6 +842,9 @@ func (_c *UsageLogCreate) check() error {
if _, ok := _c.mutation.RateMultiplier(); !ok {
return &ValidationError{Name: "rate_multiplier", err: errors.New(`ent: missing required field "UsageLog.rate_multiplier"`)}
}
if _, ok := _c.mutation.LongContextBillingApplied(); !ok {
return &ValidationError{Name: "long_context_billing_applied", err: errors.New(`ent: missing required field "UsageLog.long_context_billing_applied"`)}
}
if _, ok := _c.mutation.BillingType(); !ok {
return &ValidationError{Name: "billing_type", err: errors.New(`ent: missing required field "UsageLog.billing_type"`)}
}
@@ -997,6 +1018,10 @@ func (_c *UsageLogCreate) createSpec() (*UsageLog, *sqlgraph.CreateSpec) {
_spec.SetField(usagelog.FieldRateMultiplier, field.TypeFloat64, value)
_node.RateMultiplier = value
}
if value, ok := _c.mutation.LongContextBillingApplied(); ok {
_spec.SetField(usagelog.FieldLongContextBillingApplied, field.TypeBool, value)
_node.LongContextBillingApplied = value
}
if value, ok := _c.mutation.AccountRateMultiplier(); ok {
_spec.SetField(usagelog.FieldAccountRateMultiplier, field.TypeFloat64, value)
_node.AccountRateMultiplier = &value
@@ -1650,6 +1675,18 @@ func (u *UsageLogUpsert) AddRateMultiplier(v float64) *UsageLogUpsert {
return u
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (u *UsageLogUpsert) SetLongContextBillingApplied(v bool) *UsageLogUpsert {
u.Set(usagelog.FieldLongContextBillingApplied, v)
return u
}
// UpdateLongContextBillingApplied sets the "long_context_billing_applied" field to the value that was provided on create.
func (u *UsageLogUpsert) UpdateLongContextBillingApplied() *UsageLogUpsert {
u.SetExcluded(usagelog.FieldLongContextBillingApplied)
return u
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (u *UsageLogUpsert) SetAccountRateMultiplier(v float64) *UsageLogUpsert {
u.Set(usagelog.FieldAccountRateMultiplier, v)
@@ -2531,6 +2568,20 @@ func (u *UsageLogUpsertOne) UpdateRateMultiplier() *UsageLogUpsertOne {
})
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (u *UsageLogUpsertOne) SetLongContextBillingApplied(v bool) *UsageLogUpsertOne {
return u.Update(func(s *UsageLogUpsert) {
s.SetLongContextBillingApplied(v)
})
}
// UpdateLongContextBillingApplied sets the "long_context_billing_applied" field to the value that was provided on create.
func (u *UsageLogUpsertOne) UpdateLongContextBillingApplied() *UsageLogUpsertOne {
return u.Update(func(s *UsageLogUpsert) {
s.UpdateLongContextBillingApplied()
})
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (u *UsageLogUpsertOne) SetAccountRateMultiplier(v float64) *UsageLogUpsertOne {
return u.Update(func(s *UsageLogUpsert) {
@@ -3631,6 +3682,20 @@ func (u *UsageLogUpsertBulk) UpdateRateMultiplier() *UsageLogUpsertBulk {
})
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (u *UsageLogUpsertBulk) SetLongContextBillingApplied(v bool) *UsageLogUpsertBulk {
return u.Update(func(s *UsageLogUpsert) {
s.SetLongContextBillingApplied(v)
})
}
// UpdateLongContextBillingApplied sets the "long_context_billing_applied" field to the value that was provided on create.
func (u *UsageLogUpsertBulk) UpdateLongContextBillingApplied() *UsageLogUpsertBulk {
return u.Update(func(s *UsageLogUpsert) {
s.UpdateLongContextBillingApplied()
})
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (u *UsageLogUpsertBulk) SetAccountRateMultiplier(v float64) *UsageLogUpsertBulk {
return u.Update(func(s *UsageLogUpsert) {
+34
View File
@@ -542,6 +542,20 @@ func (_u *UsageLogUpdate) AddRateMultiplier(v float64) *UsageLogUpdate {
return _u
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (_u *UsageLogUpdate) SetLongContextBillingApplied(v bool) *UsageLogUpdate {
_u.mutation.SetLongContextBillingApplied(v)
return _u
}
// SetNillableLongContextBillingApplied sets the "long_context_billing_applied" field if the given value is not nil.
func (_u *UsageLogUpdate) SetNillableLongContextBillingApplied(v *bool) *UsageLogUpdate {
if v != nil {
_u.SetLongContextBillingApplied(*v)
}
return _u
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (_u *UsageLogUpdate) SetAccountRateMultiplier(v float64) *UsageLogUpdate {
_u.mutation.ResetAccountRateMultiplier()
@@ -1199,6 +1213,9 @@ func (_u *UsageLogUpdate) sqlSave(ctx context.Context) (_node int, err error) {
if value, ok := _u.mutation.AddedRateMultiplier(); ok {
_spec.AddField(usagelog.FieldRateMultiplier, field.TypeFloat64, value)
}
if value, ok := _u.mutation.LongContextBillingApplied(); ok {
_spec.SetField(usagelog.FieldLongContextBillingApplied, field.TypeBool, value)
}
if value, ok := _u.mutation.AccountRateMultiplier(); ok {
_spec.SetField(usagelog.FieldAccountRateMultiplier, field.TypeFloat64, value)
}
@@ -1982,6 +1999,20 @@ func (_u *UsageLogUpdateOne) AddRateMultiplier(v float64) *UsageLogUpdateOne {
return _u
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (_u *UsageLogUpdateOne) SetLongContextBillingApplied(v bool) *UsageLogUpdateOne {
_u.mutation.SetLongContextBillingApplied(v)
return _u
}
// SetNillableLongContextBillingApplied sets the "long_context_billing_applied" field if the given value is not nil.
func (_u *UsageLogUpdateOne) SetNillableLongContextBillingApplied(v *bool) *UsageLogUpdateOne {
if v != nil {
_u.SetLongContextBillingApplied(*v)
}
return _u
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (_u *UsageLogUpdateOne) SetAccountRateMultiplier(v float64) *UsageLogUpdateOne {
_u.mutation.ResetAccountRateMultiplier()
@@ -2669,6 +2700,9 @@ func (_u *UsageLogUpdateOne) sqlSave(ctx context.Context) (_node *UsageLog, err
if value, ok := _u.mutation.AddedRateMultiplier(); ok {
_spec.AddField(usagelog.FieldRateMultiplier, field.TypeFloat64, value)
}
if value, ok := _u.mutation.LongContextBillingApplied(); ok {
_spec.SetField(usagelog.FieldLongContextBillingApplied, field.TypeBool, value)
}
if value, ok := _u.mutation.AccountRateMultiplier(); ok {
_spec.SetField(usagelog.FieldAccountRateMultiplier, field.TypeFloat64, value)
}
+29
View File
@@ -601,6 +601,7 @@ type ServerConfig struct {
Host string `mapstructure:"host"`
Port int `mapstructure:"port"`
Mode string `mapstructure:"mode"` // debug/release
EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header
FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接
ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒)
IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒)
@@ -818,6 +819,8 @@ type GatewayConfig struct {
ImageStreamDataIntervalTimeout int `mapstructure:"image_stream_data_interval_timeout"`
// ImageStreamKeepaliveInterval: 图片流式 keepalive 间隔(秒),0表示禁用
ImageStreamKeepaliveInterval int `mapstructure:"image_stream_keepalive_interval"`
// ImageNonstreamKeepaliveInterval: 图片非流式 JSON keepalive 间隔(秒),0表示禁用
ImageNonstreamKeepaliveInterval int `mapstructure:"image_nonstream_keepalive_interval"`
// MaxLineSize: 上游 SSE 单行最大字节数(0使用默认值)
MaxLineSize int `mapstructure:"max_line_size"`
@@ -923,6 +926,12 @@ type GatewayOpenAIWSConfig struct {
ModeRouterV2Enabled bool `mapstructure:"mode_router_v2_enabled"`
// IngressModeDefault: ingress 默认模式(off/ctx_pool/passthrough/http_bridge
IngressModeDefault string `mapstructure:"ingress_mode_default"`
// IngressInterTurnIdleTimeoutSeconds bounds the time a client may remain idle
// between completed ingress turns. Zero disables this protection.
IngressInterTurnIdleTimeoutSeconds int `mapstructure:"ingress_inter_turn_idle_timeout_seconds"`
// MaxIngressConnectionsPerAPIKey bounds live client WebSocket ingress sessions
// per API key across all instances. Zero disables this protection.
MaxIngressConnectionsPerAPIKey int `mapstructure:"max_ingress_connections_per_api_key"`
// Enabled: 全局总开关(默认 true
Enabled bool `mapstructure:"enabled"`
// OAuthEnabled: 是否允许 OpenAI OAuth 账号使用 WS
@@ -1453,6 +1462,9 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
// 环境变量支持
viper.AutomaticEnv()
viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
if err := viper.BindEnv("server.enable_server_timing", "ENABLE_SERVER_TIMING"); err != nil {
return nil, fmt.Errorf("bind ENABLE_SERVER_TIMING: %w", err)
}
// 默认值
setDefaults()
@@ -1608,6 +1620,7 @@ func setDefaults() {
viper.SetDefault("server.host", "0.0.0.0")
viper.SetDefault("server.port", 8080)
viper.SetDefault("server.mode", "release")
viper.SetDefault("server.enable_server_timing", false)
viper.SetDefault("server.frontend_url", "")
viper.SetDefault("server.read_header_timeout", 30) // 30秒读取请求头
viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时
@@ -1945,6 +1958,8 @@ func setDefaults() {
viper.SetDefault("gateway.openai_ws.enabled", true)
viper.SetDefault("gateway.openai_ws.mode_router_v2_enabled", false)
viper.SetDefault("gateway.openai_ws.ingress_mode_default", "ctx_pool")
viper.SetDefault("gateway.openai_ws.ingress_inter_turn_idle_timeout_seconds", 300)
viper.SetDefault("gateway.openai_ws.max_ingress_connections_per_api_key", 64)
viper.SetDefault("gateway.openai_ws.oauth_enabled", true)
viper.SetDefault("gateway.openai_ws.apikey_enabled", true)
viper.SetDefault("gateway.openai_ws.force_http", false)
@@ -2024,6 +2039,7 @@ func setDefaults() {
viper.SetDefault("gateway.stream_keepalive_interval", 10)
viper.SetDefault("gateway.image_stream_data_interval_timeout", 900)
viper.SetDefault("gateway.image_stream_keepalive_interval", 10)
viper.SetDefault("gateway.image_nonstream_keepalive_interval", 0)
viper.SetDefault("gateway.max_line_size", 500*1024*1024)
viper.SetDefault("gateway.scheduling.sticky_session_max_waiting", 3)
viper.SetDefault("gateway.scheduling.sticky_session_wait_timeout", 120*time.Second)
@@ -2710,6 +2726,13 @@ func (c *Config) Validate() error {
(c.Gateway.ImageStreamKeepaliveInterval < 5 || c.Gateway.ImageStreamKeepaliveInterval > 60) {
return fmt.Errorf("gateway.image_stream_keepalive_interval must be 0 or between 5-60 seconds")
}
if c.Gateway.ImageNonstreamKeepaliveInterval < 0 {
return fmt.Errorf("gateway.image_nonstream_keepalive_interval must be non-negative")
}
if c.Gateway.ImageNonstreamKeepaliveInterval != 0 &&
(c.Gateway.ImageNonstreamKeepaliveInterval < 5 || c.Gateway.ImageNonstreamKeepaliveInterval > 60) {
return fmt.Errorf("gateway.image_nonstream_keepalive_interval must be 0 or between 5-60 seconds")
}
// 兼容旧键 sticky_previous_response_ttl_seconds
if c.Gateway.OpenAIWS.StickyResponseIDTTLSeconds <= 0 && c.Gateway.OpenAIWS.StickyPreviousResponseTTLSeconds > 0 {
c.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = c.Gateway.OpenAIWS.StickyPreviousResponseTTLSeconds
@@ -2717,6 +2740,12 @@ func (c *Config) Validate() error {
if c.Gateway.OpenAIWS.MaxConnsPerAccount <= 0 {
return fmt.Errorf("gateway.openai_ws.max_conns_per_account must be positive")
}
if c.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds < 0 {
return fmt.Errorf("gateway.openai_ws.ingress_inter_turn_idle_timeout_seconds must be non-negative")
}
if c.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey < 0 {
return fmt.Errorf("gateway.openai_ws.max_ingress_connections_per_api_key must be non-negative")
}
if c.Gateway.OpenAIWS.MinIdlePerAccount < 0 {
return fmt.Errorf("gateway.openai_ws.min_idle_per_account must be non-negative")
}
+55
View File
@@ -17,6 +17,23 @@ func resetViperWithJWTSecret(t *testing.T) {
t.Setenv("JWT_SECRET", strings.Repeat("x", 32))
}
func TestLoadServerTimingConfig(t *testing.T) {
t.Run("disabled by default", func(t *testing.T) {
resetViperWithJWTSecret(t)
cfg, err := Load()
require.NoError(t, err)
require.False(t, cfg.Server.EnableServerTiming)
})
t.Run("enabled by exact environment variable", func(t *testing.T) {
resetViperWithJWTSecret(t)
t.Setenv("ENABLE_SERVER_TIMING", "true")
cfg, err := Load()
require.NoError(t, err)
require.True(t, cfg.Server.EnableServerTiming)
})
}
func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) {
viper.Reset()
t.Setenv("JWT_SECRET", "")
@@ -182,6 +199,12 @@ func TestLoadDefaultOpenAIWSConfig(t *testing.T) {
if cfg.Gateway.OpenAIWS.IngressModeDefault != "ctx_pool" {
t.Fatalf("Gateway.OpenAIWS.IngressModeDefault = %q, want %q", cfg.Gateway.OpenAIWS.IngressModeDefault, "ctx_pool")
}
if cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds != 300 {
t.Fatalf("Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = %d, want 300", cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds)
}
if cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey != 64 {
t.Fatalf("Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = %d, want 64", cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey)
}
}
func TestLoadDefaultOpenAICompactModel(t *testing.T) {
@@ -236,6 +259,15 @@ func TestLoadOpenAIResponseHeaderTimeoutFromEnv(t *testing.T) {
require.Equal(t, 1800, cfg.Gateway.OpenAIResponseHeaderTimeout)
}
func TestLoadImageNonstreamKeepaliveFromEnv(t *testing.T) {
resetViperWithJWTSecret(t)
t.Setenv("GATEWAY_IMAGE_NONSTREAM_KEEPALIVE_INTERVAL", "15")
cfg, err := Load()
require.NoError(t, err)
require.Equal(t, 15, cfg.Gateway.ImageNonstreamKeepaliveInterval)
}
func TestLoadOpenAIWSStickyTTLCompatibility(t *testing.T) {
resetViperWithJWTSecret(t)
t.Setenv("GATEWAY_OPENAI_WS_STICKY_RESPONSE_ID_TTL_SECONDS", "0")
@@ -1406,6 +1438,16 @@ func TestValidateConfigErrors(t *testing.T) {
mutate: func(c *Config) { c.Gateway.ImageStreamKeepaliveInterval = -1 },
wantErr: "gateway.image_stream_keepalive_interval must be non-negative",
},
{
name: "gateway image nonstream keepalive range",
mutate: func(c *Config) { c.Gateway.ImageNonstreamKeepaliveInterval = 4 },
wantErr: "gateway.image_nonstream_keepalive_interval",
},
{
name: "gateway image nonstream keepalive negative",
mutate: func(c *Config) { c.Gateway.ImageNonstreamKeepaliveInterval = -1 },
wantErr: "gateway.image_nonstream_keepalive_interval must be non-negative",
},
{
name: "gateway image stream data interval range",
mutate: func(c *Config) { c.Gateway.ImageStreamDataIntervalTimeout = 30 },
@@ -1640,6 +1682,16 @@ func TestValidateConfig_OpenAIWSRules(t *testing.T) {
mutate: func(c *Config) { c.Gateway.OpenAIWS.MaxConnsPerAccount = 0 },
wantErr: "gateway.openai_ws.max_conns_per_account",
},
{
name: "ingress_inter_turn_idle_timeout_seconds 不能为负数",
mutate: func(c *Config) { c.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = -1 },
wantErr: "gateway.openai_ws.ingress_inter_turn_idle_timeout_seconds",
},
{
name: "max_ingress_connections_per_api_key 不能为负数",
mutate: func(c *Config) { c.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = -1 },
wantErr: "gateway.openai_ws.max_ingress_connections_per_api_key",
},
{
name: "min_idle_per_account 不能为负数",
mutate: func(c *Config) { c.Gateway.OpenAIWS.MinIdlePerAccount = -1 },
@@ -1964,6 +2016,9 @@ func TestLoad_DefaultGatewayImageStreamConfig(t *testing.T) {
if cfg.Gateway.ImageStreamKeepaliveInterval != 10 {
t.Fatalf("image_stream_keepalive_interval = %d, want 10", cfg.Gateway.ImageStreamKeepaliveInterval)
}
if cfg.Gateway.ImageNonstreamKeepaliveInterval != 0 {
t.Fatalf("image_nonstream_keepalive_interval = %d, want 0", cfg.Gateway.ImageNonstreamKeepaliveInterval)
}
if cfg.Gateway.ImageConcurrency.Enabled {
t.Fatalf("image_concurrency.enabled = true, want false")
}
@@ -115,6 +115,10 @@ func (h *AccountHandler) ImportCodexSession(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
if err := service.ValidateOpenAILongContextBillingExtra(service.PlatformOpenAI, req.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
if req.Concurrency != nil && *req.Concurrency < 0 {
response.BadRequest(c, "concurrency must be >= 0")
return
@@ -630,6 +630,7 @@ func TestImportCodexSessionsAccessTokenOnlySameUserUpdatesExisting(t *testing.T)
"chatgpt_user_id": "user-1",
"access_token": existingToken,
},
Extra: map[string]any{"openai_long_context_billing_enabled": false},
}})
handler := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
req := CodexSessionImportRequest{SkipDefaultGroupBind: boolPtr(true)}
@@ -650,6 +651,9 @@ func TestImportCodexSessionsAccessTokenOnlySameUserUpdatesExisting(t *testing.T)
if len(svc.updatedAccounts) != 1 || svc.updatedAccounts[0].id != 10 {
t.Fatalf("updated accounts = %+v, want account 10", svc.updatedAccounts)
}
if got := svc.updatedAccounts[0].input.Extra["openai_long_context_billing_enabled"]; got != false {
t.Fatalf("openai_long_context_billing_enabled = %v, want false", got)
}
}
func TestImportCodexSessionsUpgradesAccessTokenOnlyAccountWithRefreshToken(t *testing.T) {
@@ -460,6 +460,7 @@ func (h *AccountHandler) importData(ctx context.Context, req DataImportRequest)
if created.Platform == service.PlatformAntigravity && created.Type == service.AccountTypeOAuth {
privacyAccounts = append(privacyAccounts, created)
}
h.scheduleGrokImportProbe(created)
result.AccountCreated++
}
@@ -60,6 +60,7 @@ type AccountHandler struct {
sessionLimitCache service.SessionLimitCache
rpmCache service.RPMCache
tokenCacheInvalidator service.TokenCacheInvalidator
grokImportProber grokUsageProber
}
// NewAccountHandler creates a new admin account handler
@@ -784,6 +785,10 @@ func (h *AccountHandler) Create(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
if err := service.ValidateOpenAILongContextBillingExtra(req.Platform, req.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
if req.RateMultiplier != nil && *req.RateMultiplier < 0 {
response.BadRequest(c, "rate_multiplier must be >= 0")
return
@@ -851,6 +856,7 @@ func (h *AccountHandler) Create(c *gin.Context) {
// OpenAI APIKey 账号创建后异步探测上游 /v1/responses 能力。
// 探测失败不影响账号创建响应。
h.scheduleOpenAIResponsesProbe(createdAccount)
h.scheduleGrokImportProbe(createdAccount)
response.Success(c, result.Data)
}
@@ -1299,6 +1305,10 @@ func (h *AccountHandler) ApplyOAuthCredentials(c *gin.Context) {
response.ErrorFrom(c, infraerrors.BadRequest("NOT_OAUTH", "cannot apply oauth credentials to non-OAuth account"))
return
}
if err := service.ValidateOpenAILongContextBillingExtra(existing.Platform, req.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
updatedAccount, err := h.adminService.UpdateAccount(ctx, accountID, &service.UpdateAccountInput{
Type: req.Type,
@@ -1592,6 +1602,12 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
for _, item := range req.Accounts {
if err := service.ValidateOpenAILongContextBillingExtra(item.Platform, item.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
}
executeAdminIdempotentJSON(c, "admin.accounts.batch_create", req, service.DefaultWriteIdempotencyTTL(), func(ctx context.Context) (any, error) {
success := 0
@@ -1653,6 +1669,7 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) {
}
// OpenAI APIKey 账号异步探测 /v1/responses 能力。
h.scheduleOpenAIResponsesProbe(account)
h.scheduleGrokImportProbe(account)
success++
results = append(results, gin.H{
"name": item.Name,
@@ -0,0 +1,165 @@
package admin
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestAccountAdminBoundariesRejectMalformedOpenAILongContextBillingValue(t *testing.T) {
const malformedExtra = `"extra":{"openai_long_context_billing_enabled":"true"}`
tests := []struct {
name string
method string
path string
body string
mount func(*gin.Engine, *AccountHandler)
setup func(*stubAdminService)
}{
{
name: "create",
method: http.MethodPost,
path: "/accounts",
body: `{"name":"account","platform":"openai","type":"apikey","credentials":{"api_key":"test"},` + malformedExtra + `}`,
mount: func(router *gin.Engine, handler *AccountHandler) { router.POST("/accounts", handler.Create) },
},
{
name: "update",
method: http.MethodPut,
path: "/accounts/1",
body: `{` + malformedExtra + `}`,
mount: func(router *gin.Engine, handler *AccountHandler) { router.PUT("/accounts/:id", handler.Update) },
setup: func(stub *stubAdminService) {
stub.updateAccountErr = infraerrors.BadRequest("OPENAI_LONG_CONTEXT_BILLING_INVALID", "invalid")
},
},
{
name: "bulk update",
method: http.MethodPost,
path: "/accounts/bulk-update",
body: `{"account_ids":[1],` + malformedExtra + `}`,
mount: func(router *gin.Engine, handler *AccountHandler) {
router.POST("/accounts/bulk-update", handler.BulkUpdate)
},
setup: func(stub *stubAdminService) {
stub.bulkUpdateAccountErr = infraerrors.BadRequest("OPENAI_LONG_CONTEXT_BILLING_INVALID", "invalid")
},
},
{
name: "batch create",
method: http.MethodPost,
path: "/accounts/batch",
body: `{"accounts":[{"name":"account","platform":"openai","type":"apikey","credentials":{"api_key":"test"},` + malformedExtra + `}]}`,
mount: func(router *gin.Engine, handler *AccountHandler) { router.POST("/accounts/batch", handler.BatchCreate) },
},
{
name: "Codex session import",
method: http.MethodPost,
path: "/accounts/import-codex-session",
body: `{"content":"token",` + malformedExtra + `}`,
mount: func(router *gin.Engine, handler *AccountHandler) {
router.POST("/accounts/import-codex-session", handler.ImportCodexSession)
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
stub := newStubAdminService()
if tt.setup != nil {
tt.setup(stub)
}
handler := NewAccountHandler(stub, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
router := gin.New()
tt.mount(router, handler)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(tt.method, tt.path, bytes.NewBufferString(tt.body))
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
var responseBody struct {
Reason string `json:"reason"`
}
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &responseBody))
require.Equal(t, "OPENAI_LONG_CONTEXT_BILLING_INVALID", responseBody.Reason)
})
}
}
func TestAccountCreateBoundaryDoesNotApplyOpenAIValidationToOtherPlatforms(t *testing.T) {
gin.SetMode(gin.TestMode)
handler := NewAccountHandler(newStubAdminService(), nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
router := gin.New()
router.POST("/accounts", handler.Create)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/accounts", bytes.NewBufferString(
`{"name":"account","platform":"anthropic","type":"apikey","credentials":{"api_key":"test"},"extra":{"openai_long_context_billing_enabled":"provider-owned"}}`,
))
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
}
func TestApplyOAuthCredentialsRejectsMalformedOpenAILongContextBillingBeforeMutation(t *testing.T) {
gin.SetMode(gin.TestMode)
stub := newStubAdminService()
stub.getAccountResult = &service.Account{
ID: 1,
Platform: service.PlatformOpenAI,
Type: service.AccountTypeOAuth,
}
handler := NewAccountHandler(stub, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
router := gin.New()
router.POST("/accounts/:id/apply-oauth-credentials", handler.ApplyOAuthCredentials)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/accounts/1/apply-oauth-credentials", bytes.NewBufferString(
`{"type":"oauth","credentials":{"access_token":"new-token"},"extra":{"openai_long_context_billing_enabled":"true"}}`,
))
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
var responseBody struct {
Reason string `json:"reason"`
}
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &responseBody))
require.Equal(t, "OPENAI_LONG_CONTEXT_BILLING_INVALID", responseBody.Reason)
require.Zero(t, stub.updateAccountCalls)
require.Zero(t, stub.updateAccountExtraCalls)
}
func TestOpenAIOAuthCodexPATBoundaryRejectsMalformedOpenAILongContextBillingValueBeforeTokenValidation(t *testing.T) {
gin.SetMode(gin.TestMode)
handler := NewOpenAIOAuthHandler(nil, newStubAdminService(), nil)
router := gin.New()
router.Use(gin.Recovery())
router.POST("/openai/create-from-codex-pat", handler.CreateAccountFromCodexPAT)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/openai/create-from-codex-pat", bytes.NewBufferString(
`{"access_token":"token","extra":{"openai_long_context_billing_enabled":1}}`,
))
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
var responseBody struct {
Reason string `json:"reason"`
}
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &responseBody))
require.Equal(t, "OPENAI_LONG_CONTEXT_BILLING_INVALID", responseBody.Reason)
}
@@ -33,6 +33,9 @@ type stubAdminService struct {
createSparkShadowErr error
updateAccountErr error
bulkUpdateAccountErr error
getAccountResult *service.Account
updateAccountCalls int
updateAccountExtraCalls int
checkMixedErr error
lastMixedCheck struct {
accountID int64
@@ -388,6 +391,9 @@ func (s *stubAdminService) ListOpenAISchedulableAccountsForSchedulerScore(_ cont
}
func (s *stubAdminService) GetAccount(ctx context.Context, id int64) (*service.Account, error) {
if s.getAccountResult != nil {
return s.getAccountResult, nil
}
account := service.Account{ID: id, Name: "account", Status: service.StatusActive}
return &account, nil
}
@@ -413,6 +419,7 @@ func (s *stubAdminService) CreateAccount(ctx context.Context, input *service.Cre
}
func (s *stubAdminService) UpdateAccount(ctx context.Context, id int64, input *service.UpdateAccountInput) (*service.Account, error) {
s.updateAccountCalls++
if s.updateAccountErr != nil {
return nil, s.updateAccountErr
}
@@ -421,6 +428,7 @@ func (s *stubAdminService) UpdateAccount(ctx context.Context, id int64, input *s
}
func (s *stubAdminService) UpdateAccountExtra(ctx context.Context, id int64, updates map[string]any) error {
s.updateAccountExtraCalls++
return nil
}
@@ -37,11 +37,11 @@ func NewChannelMonitorHandler(monitorService *service.ChannelMonitorService) *Ch
type channelMonitorCreateRequest struct {
Name string `json:"name" binding:"required,max=100"`
Provider string `json:"provider" binding:"required,oneof=openai anthropic gemini"`
Provider string `json:"provider" binding:"required,oneof=openai anthropic gemini grok"`
APIMode string `json:"api_mode" binding:"omitempty,oneof=chat_completions responses"`
Endpoint string `json:"endpoint" binding:"required,max=500"`
APIKey string `json:"api_key" binding:"required,max=2000"`
PrimaryModel string `json:"primary_model" binding:"required,max=200"`
PrimaryModel string `json:"primary_model" binding:"max=200"`
ExtraModels []string `json:"extra_models"`
GroupName string `json:"group_name" binding:"max=100"`
Enabled *bool `json:"enabled"`
@@ -55,7 +55,7 @@ type channelMonitorCreateRequest struct {
type channelMonitorUpdateRequest struct {
Name *string `json:"name" binding:"omitempty,max=100"`
Provider *string `json:"provider" binding:"omitempty,oneof=openai anthropic gemini"`
Provider *string `json:"provider" binding:"omitempty,oneof=openai anthropic gemini grok"`
APIMode *string `json:"api_mode" binding:"omitempty,oneof=chat_completions responses"`
Endpoint *string `json:"endpoint" binding:"omitempty,max=500"`
APIKey *string `json:"api_key" binding:"omitempty,max=2000"`
@@ -26,7 +26,7 @@ func NewChannelMonitorRequestTemplateHandler(templateService *service.ChannelMon
type channelMonitorTemplateCreateRequest struct {
Name string `json:"name" binding:"required,max=100"`
Provider string `json:"provider" binding:"required,oneof=openai anthropic gemini"`
Provider string `json:"provider" binding:"required,oneof=openai anthropic gemini grok"`
APIMode string `json:"api_mode" binding:"omitempty,oneof=chat_completions responses"`
Description string `json:"description" binding:"max=500"`
ExtraHeaders map[string]string `json:"extra_headers"`
@@ -0,0 +1,203 @@
package admin
import (
"context"
"log/slog"
"sync"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/service"
)
const (
grokImportProbeConcurrency = 3
grokImportProbeTimeout = 25 * time.Second
)
type grokUsageProber interface {
ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
}
type grokImportProbeTask struct {
prober grokUsageProber
accountID int64
}
type grokImportProbeScheduler struct {
mu sync.Mutex
queue []grokImportProbeTask
concurrency int
workers int
maxWorkers int
timeout time.Duration
}
var defaultGrokImportProbeScheduler = newGrokImportProbeScheduler(
grokImportProbeConcurrency,
grokImportProbeTimeout,
)
func newGrokImportProbeScheduler(concurrency int, timeout time.Duration) *grokImportProbeScheduler {
if concurrency <= 0 {
concurrency = 1
}
if timeout <= 0 {
timeout = grokImportProbeTimeout
}
return &grokImportProbeScheduler{
concurrency: concurrency,
timeout: timeout,
}
}
func (s *grokImportProbeScheduler) schedule(prober grokUsageProber, account *service.Account) {
if s == nil || prober == nil || account == nil || account.ID <= 0 {
return
}
if account.Platform != service.PlatformGrok || account.Type != service.AccountTypeOAuth {
return
}
s.mu.Lock()
s.queue = append(s.queue, grokImportProbeTask{prober: prober, accountID: account.ID})
if s.workers < s.concurrency {
s.workers++
if s.workers > s.maxWorkers {
s.maxWorkers = s.workers
}
go s.worker()
}
s.mu.Unlock()
}
func (s *grokImportProbeScheduler) worker() {
for {
task, ok := s.nextTask()
if !ok {
return
}
s.run(task.prober, task.accountID)
}
}
func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) {
s.mu.Lock()
defer s.mu.Unlock()
if len(s.queue) == 0 {
s.workers--
return grokImportProbeTask{}, false
}
task := s.queue[0]
s.queue[0] = grokImportProbeTask{}
s.queue = s.queue[1:]
if len(s.queue) == 0 {
s.queue = nil
}
return task, true
}
func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) {
defer func() {
if recovered := recover(); recovered != nil {
slog.Error(
"grok_import_active_probe_panic",
"account_id", accountID,
"recovery_type", panicType(recovered),
)
}
}()
// Queue time is intentionally excluded: every imported account is probed,
// while this timeout only bounds the actual upstream probe execution.
ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
defer cancel()
result, err := prober.ProbeUsage(ctx, accountID)
if err != nil {
slog.Warn(
"grok_import_active_probe_failed",
"account_id", accountID,
"status", infraerrors.Code(err),
"reason", infraerrors.Reason(err),
)
return
}
if result == nil {
slog.Warn(
"grok_import_active_probe_failed",
"account_id", accountID,
"reason", "empty_result",
)
return
}
slog.Info(
"grok_import_active_probe_completed",
"account_id", accountID,
"model", result.Model,
"status", result.StatusCode,
"headers_observed", result.HeadersObserved,
)
}
func panicType(value any) string {
switch value.(type) {
case string:
return "string"
case error:
return "error"
default:
return "unknown"
}
}
func (h *AccountHandler) scheduleGrokImportProbe(account *service.Account) {
if h == nil {
return
}
defaultGrokImportProbeScheduler.schedule(h.grokImportProber, account)
}
func (h *GrokOAuthHandler) scheduleGrokImportProbe(account *service.Account) {
if h == nil {
return
}
defaultGrokImportProbeScheduler.schedule(h.importProber, account)
}
// ProvideAccountHandler injects the Grok active prober for production while
// keeping NewAccountHandler convenient for focused unit tests.
func ProvideAccountHandler(
adminService service.AdminService,
oauthService *service.OAuthService,
openaiOAuthService *service.OpenAIOAuthService,
geminiOAuthService *service.GeminiOAuthService,
antigravityOAuthService *service.AntigravityOAuthService,
rateLimitService *service.RateLimitService,
accountUsageService *service.AccountUsageService,
accountTestService *service.AccountTestService,
concurrencyService *service.ConcurrencyService,
crsSyncService *service.CRSSyncService,
sessionLimitCache service.SessionLimitCache,
rpmCache service.RPMCache,
tokenCacheInvalidator service.TokenCacheInvalidator,
grokQuotaService *service.GrokQuotaService,
) *AccountHandler {
handler := NewAccountHandler(
adminService,
oauthService,
openaiOAuthService,
geminiOAuthService,
antigravityOAuthService,
rateLimitService,
accountUsageService,
accountTestService,
concurrencyService,
crsSyncService,
sessionLimitCache,
rpmCache,
tokenCacheInvalidator,
)
handler.grokImportProber = grokQuotaService
return handler
}
@@ -0,0 +1,119 @@
//go:build unit
package admin
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type grokImportAdminService struct {
*stubAdminService
mu sync.Mutex
nextID int64
}
func newGrokImportAdminService() *grokImportAdminService {
return &grokImportAdminService{
stubAdminService: newStubAdminService(),
nextID: 500,
}
}
func (s *grokImportAdminService) CreateAccount(_ context.Context, input *service.CreateAccountInput) (*service.Account, error) {
s.mu.Lock()
s.nextID++
id := s.nextID
s.mu.Unlock()
return &service.Account{
ID: id,
Name: input.Name,
Platform: input.Platform,
Type: input.Type,
Credentials: input.Credentials,
Extra: input.Extra,
ProxyID: input.ProxyID,
Concurrency: input.Concurrency,
Status: service.StatusActive,
Schedulable: true,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}, nil
}
type grokImportOAuthClientStub struct{}
func (grokImportOAuthClientStub) ExchangeCode(context.Context, string, string, string, string, string) (*xai.TokenResponse, error) {
return &xai.TokenResponse{AccessToken: "access-token", RefreshToken: "refresh-token", ExpiresIn: 3600}, nil
}
func (grokImportOAuthClientStub) RefreshToken(context.Context, string, string, string) (*xai.TokenResponse, error) {
return &xai.TokenResponse{AccessToken: "access-token", RefreshToken: "refresh-token", ExpiresIn: 3600}, nil
}
func (grokImportOAuthClientStub) ConvertSSOToBuild(context.Context, string, string) (*xai.TokenResponse, error) {
return &xai.TokenResponse{AccessToken: "access-token", RefreshToken: "refresh-token", ExpiresIn: 3600}, nil
}
func TestGrokSSOBatchImportKeepsCreatedAccountsWhenOneAutomaticProbeFails(t *testing.T) {
gin.SetMode(gin.TestMode)
adminService := newGrokImportAdminService()
oauthService := service.NewGrokOAuthService(nil, grokImportOAuthClientStub{})
defer oauthService.Stop()
prober := newGrokImportProbeStub(3)
prober.failures[502] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "sensitive-upstream-body")
handler := NewGrokOAuthHandler(oauthService, adminService, nil)
handler.importProber = prober
router := gin.New()
router.POST("/api/v1/admin/grok/sso-to-oauth", handler.CreateAccountsFromSSO)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(
http.MethodPost,
"/api/v1/admin/grok/sso-to-oauth",
strings.NewReader(`{"sso_tokens":["sso-one","sso-two","sso-three"]}`),
)
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Contains(t, recorder.Body.String(), `"created"`)
require.NotContains(t, recorder.Body.String(), `GROK_TEST_PROBE_FAILED`)
for i := 0; i < 3; i++ {
awaitGrokProbeSignal(t, prober.done)
}
calls, _, _ := prober.snapshot()
require.Equal(t, map[int64]int{501: 1, 502: 1, 503: 1}, calls)
}
func TestAccountCreateWithoutAutomaticGrokProbeServiceStillSucceeds(t *testing.T) {
gin.SetMode(gin.TestMode)
handler := NewAccountHandler(
newGrokImportAdminService(),
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
)
router := gin.New()
router.POST("/api/v1/admin/accounts", handler.Create)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(
http.MethodPost,
"/api/v1/admin/accounts",
strings.NewReader(`{"name":"grok-rt","platform":"grok","type":"oauth","credentials":{"refresh_token":"secret"}}`),
)
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
}
@@ -0,0 +1,232 @@
//go:build unit
package admin
import (
"bytes"
"context"
"log/slog"
"sync"
"testing"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
type grokImportProbeStub struct {
mu sync.Mutex
calls map[int64]int
failures map[int64]error
active int
maxActive int
deadlineSeen bool
block <-chan struct{}
started chan int64
done chan int64
}
func newGrokImportProbeStub(buffer int) *grokImportProbeStub {
return &grokImportProbeStub{
calls: make(map[int64]int),
failures: make(map[int64]error),
started: make(chan int64, buffer),
done: make(chan int64, buffer),
}
}
func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
_, deadlineSeen := ctx.Deadline()
s.mu.Lock()
s.calls[accountID]++
s.active++
if s.active > s.maxActive {
s.maxActive = s.active
}
s.deadlineSeen = s.deadlineSeen || deadlineSeen
s.mu.Unlock()
s.started <- accountID
var ctxErr error
if s.block != nil {
select {
case <-s.block:
case <-ctx.Done():
ctxErr = ctx.Err()
}
}
s.mu.Lock()
s.active--
failure := s.failures[accountID]
s.mu.Unlock()
s.done <- accountID
if ctxErr != nil {
return nil, ctxErr
}
if failure != nil {
return nil, failure
}
return &service.GrokQuotaProbeResult{
Source: "active_probe",
Model: "grok-4.5",
StatusCode: 200,
ResetSupported: false,
}, nil
}
func (s *grokImportProbeStub) snapshot() (map[int64]int, int, bool) {
s.mu.Lock()
defer s.mu.Unlock()
calls := make(map[int64]int, len(s.calls))
for id, count := range s.calls {
calls[id] = count
}
return calls, s.maxActive, s.deadlineSeen
}
type grokImportProbeSchedulerTestSnapshot struct {
queued int
workers int
maxWorkers int
}
func snapshotGrokImportProbeScheduler(s *grokImportProbeScheduler) grokImportProbeSchedulerTestSnapshot {
if s == nil {
return grokImportProbeSchedulerTestSnapshot{}
}
s.mu.Lock()
defer s.mu.Unlock()
return grokImportProbeSchedulerTestSnapshot{
queued: len(s.queue),
workers: s.workers,
maxWorkers: s.maxWorkers,
}
}
func newGrokOAuthImportAccount(id int64) *service.Account {
return &service.Account{
ID: id,
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
}
}
func awaitGrokProbeSignal(t *testing.T, signals <-chan int64) int64 {
t.Helper()
select {
case id := <-signals:
return id
case <-time.After(time.Second):
t.Fatal("timed out waiting for Grok import probe")
return 0
}
}
func TestGrokImportProbeSchedulerProbesSingleAccountOnce(t *testing.T) {
scheduler := newGrokImportProbeScheduler(1, time.Second)
prober := newGrokImportProbeStub(1)
scheduler.schedule(prober, newGrokOAuthImportAccount(101))
require.Equal(t, int64(101), awaitGrokProbeSignal(t, prober.done))
calls, maxActive, deadlineSeen := prober.snapshot()
require.Equal(t, map[int64]int{101: 1}, calls)
require.Equal(t, 1, maxActive)
require.True(t, deadlineSeen)
require.Eventually(t, func() bool {
snapshot := snapshotGrokImportProbeScheduler(scheduler)
return snapshot.queued == 0 && snapshot.workers == 0
}, time.Second, 10*time.Millisecond)
}
func TestGrokImportProbeSchedulerQueuesBatchWithoutPerTaskGoroutines(t *testing.T) {
const taskCount = 100
release := make(chan struct{})
scheduler := newGrokImportProbeScheduler(3, time.Second)
prober := newGrokImportProbeStub(taskCount)
prober.block = release
prober.failures[150] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "sensitive-upstream-body")
for id := int64(101); id < 101+taskCount; id++ {
scheduler.schedule(prober, newGrokOAuthImportAccount(id))
}
for i := 0; i < 3; i++ {
awaitGrokProbeSignal(t, prober.started)
}
snapshot := snapshotGrokImportProbeScheduler(scheduler)
require.Equal(t, 97, snapshot.queued)
require.Equal(t, 3, snapshot.workers)
require.Equal(t, 3, snapshot.maxWorkers)
select {
case id := <-prober.started:
t.Fatalf("probe %d started before a concurrency slot was released", id)
case <-time.After(75 * time.Millisecond):
}
close(release)
for i := 0; i < taskCount; i++ {
awaitGrokProbeSignal(t, prober.done)
}
calls, maxActive, _ := prober.snapshot()
require.Len(t, calls, taskCount)
for id := int64(101); id < 101+taskCount; id++ {
require.Equal(t, 1, calls[id])
}
require.Equal(t, 3, maxActive)
require.Eventually(t, func() bool {
snapshot = snapshotGrokImportProbeScheduler(scheduler)
return snapshot.queued == 0 && snapshot.workers == 0
}, time.Second, 10*time.Millisecond)
require.Equal(t, 3, snapshot.maxWorkers)
}
func TestGrokImportProbeSchedulerTimeoutCancelsProbe(t *testing.T) {
neverRelease := make(chan struct{})
scheduler := newGrokImportProbeScheduler(1, 20*time.Millisecond)
prober := newGrokImportProbeStub(1)
prober.block = neverRelease
scheduler.schedule(prober, newGrokOAuthImportAccount(201))
require.Equal(t, int64(201), awaitGrokProbeSignal(t, prober.done))
calls, _, _ := prober.snapshot()
require.Equal(t, 1, calls[201])
}
func TestGrokImportProbeSchedulerSkipsMissingServiceAndNonGrokAccounts(t *testing.T) {
scheduler := newGrokImportProbeScheduler(1, time.Second)
prober := newGrokImportProbeStub(1)
scheduler.schedule(nil, newGrokOAuthImportAccount(301))
scheduler.schedule(prober, &service.Account{ID: 302, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth})
scheduler.schedule(prober, &service.Account{ID: 303, Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey})
select {
case id := <-prober.started:
t.Fatalf("unexpected probe for account %d", id)
case <-time.After(50 * time.Millisecond):
}
calls, _, _ := prober.snapshot()
require.Empty(t, calls)
}
func TestGrokImportProbeFailureLogDoesNotIncludeErrorMessage(t *testing.T) {
var logs bytes.Buffer
previousLogger := slog.Default()
slog.SetDefault(slog.New(slog.NewTextHandler(&logs, nil)))
defer slog.SetDefault(previousLogger)
scheduler := newGrokImportProbeScheduler(1, time.Second)
prober := newGrokImportProbeStub(1)
prober.failures[401] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "refresh-token-secret")
scheduler.schedule(prober, newGrokOAuthImportAccount(401))
awaitGrokProbeSignal(t, prober.done)
require.Eventually(t, func() bool {
return bytes.Contains(logs.Bytes(), []byte("grok_import_active_probe_failed"))
}, time.Second, 10*time.Millisecond)
require.Contains(t, logs.String(), "GROK_TEST_PROBE_FAILED")
require.NotContains(t, logs.String(), "refresh-token-secret")
}
@@ -1,20 +1,28 @@
package admin
import (
"context"
"fmt"
"log/slog"
"strconv"
"strings"
"sync"
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
const grokSSOImportConcurrency = 3
type GrokOAuthHandler struct {
grokOAuthService *service.GrokOAuthService
adminService service.AdminService
quotaService *service.GrokQuotaService
importProber grokUsageProber
}
func NewGrokOAuthHandler(
@@ -26,6 +34,7 @@ func NewGrokOAuthHandler(
grokOAuthService: grokOAuthService,
adminService: adminService,
quotaService: quotaService,
importProber: quotaService,
}
}
@@ -202,9 +211,243 @@ func (h *GrokOAuthHandler) CreateAccountFromOAuth(c *gin.Context) {
response.ErrorFrom(c, err)
return
}
h.scheduleGrokImportProbe(account)
response.Success(c, dto.AccountFromService(account))
}
type GrokSSOToOAuthRequest struct {
SSOTokens []string `json:"sso_tokens"`
SSOToken string `json:"sso_token"`
Name string `json:"name"`
Notes *string `json:"notes"`
ProxyID *int64 `json:"proxy_id"`
GroupIDs []int64 `json:"group_ids"`
Credentials map[string]any `json:"credentials"`
Extra map[string]any `json:"extra"`
Concurrency int `json:"concurrency"`
LoadFactor *int `json:"load_factor"`
Priority int `json:"priority"`
RateMultiplier *float64 `json:"rate_multiplier"`
ExpiresAt *int64 `json:"expires_at"`
AutoPauseOnExpired *bool `json:"auto_pause_on_expired"`
}
type GrokSSOToOAuthItemResult struct {
Index int `json:"index"`
Name string `json:"name,omitempty"`
Email string `json:"email,omitempty"`
Account *dto.Account `json:"account,omitempty"`
Error string `json:"error,omitempty"`
}
type GrokSSOToOAuthResponse struct {
Created []GrokSSOToOAuthItemResult `json:"created"`
Failed []GrokSSOToOAuthItemResult `json:"failed"`
}
type grokSSOImportJob struct {
index int
token string
}
type grokSSOImportWorkerResult struct {
created bool
item GrokSSOToOAuthItemResult
}
func (h *GrokOAuthHandler) CreateAccountsFromSSO(c *gin.Context) {
var req GrokSSOToOAuthRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
tokens := normalizeSSOImportTokens(req.SSOTokens, req.SSOToken)
if len(tokens) == 0 {
response.BadRequest(c, "sso_tokens is required")
return
}
ctx := c.Request.Context()
workerCount := grokSSOImportConcurrency
if len(tokens) < workerCount {
workerCount = len(tokens)
}
jobs := make(chan grokSSOImportJob)
items := make([]grokSSOImportWorkerResult, len(tokens))
var wg sync.WaitGroup
for i := 0; i < workerCount; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for job := range jobs {
items[job.index] = h.safeCreateAccountFromSSOToken(ctx, req, job.token, job.index+1, len(tokens))
}
}()
}
for i, token := range tokens {
jobs <- grokSSOImportJob{index: i, token: token}
}
close(jobs)
wg.Wait()
result := GrokSSOToOAuthResponse{
Created: make([]GrokSSOToOAuthItemResult, 0, len(tokens)),
Failed: make([]GrokSSOToOAuthItemResult, 0),
}
for _, item := range items {
if item.created {
result.Created = append(result.Created, item.item)
} else {
result.Failed = append(result.Failed, item.item)
}
}
response.Success(c, result)
}
func (h *GrokOAuthHandler) safeCreateAccountFromSSOToken(ctx context.Context, req GrokSSOToOAuthRequest, token string, index, total int) (result grokSSOImportWorkerResult) {
defer func() {
if recovered := recover(); recovered != nil {
slog.Error("grok_sso_import_worker_panic", "index", index, "recover", recovered)
result = grokSSOImportWorkerResult{
item: GrokSSOToOAuthItemResult{
Index: index,
Error: fmt.Sprintf("internal worker panic: %v", recovered),
},
}
}
}()
return h.createAccountFromSSOToken(ctx, req, token, index, total)
}
func (h *GrokOAuthHandler) createAccountFromSSOToken(ctx context.Context, req GrokSSOToOAuthRequest, token string, index, total int) grokSSOImportWorkerResult {
tokenInfo, err := h.grokOAuthService.ConvertFromSSO(ctx, token, req.ProxyID)
if err != nil {
return grokSSOImportWorkerResult{item: GrokSSOToOAuthItemResult{Index: index, Error: grokSSOImportErrorMessage(err)}}
}
credentials := h.grokOAuthService.BuildAccountCredentials(tokenInfo)
credentials = service.MergeCredentials(cloneGrokSSOMap(req.Credentials), credentials)
name := grokSSOImportAccountName(req.Name, tokenInfo, index, total)
expiresAt, autoPauseOnExpired := grokSSOImportExpiry(req.ExpiresAt, req.AutoPauseOnExpired, tokenInfo)
account, err := h.adminService.CreateAccount(ctx, &service.CreateAccountInput{
Name: name,
Notes: req.Notes,
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Credentials: credentials,
Extra: cloneGrokSSOMap(req.Extra),
ProxyID: req.ProxyID,
Concurrency: req.Concurrency,
LoadFactor: req.LoadFactor,
Priority: req.Priority,
RateMultiplier: req.RateMultiplier,
GroupIDs: append([]int64(nil), req.GroupIDs...),
ExpiresAt: expiresAt,
AutoPauseOnExpired: autoPauseOnExpired,
})
if err != nil {
return grokSSOImportWorkerResult{item: GrokSSOToOAuthItemResult{Index: index, Name: name, Email: tokenInfo.Email, Error: grokSSOImportErrorMessage(err)}}
}
h.scheduleGrokImportProbe(account)
return grokSSOImportWorkerResult{
created: true,
item: GrokSSOToOAuthItemResult{
Index: index,
Name: name,
Email: tokenInfo.Email,
Account: dto.AccountFromService(account),
},
}
}
func grokSSOImportExpiry(requestExpiresAt *int64, requestAutoPause *bool, tokenInfo *service.GrokTokenInfo) (*int64, *bool) {
if tokenInfo == nil || strings.TrimSpace(tokenInfo.RefreshToken) != "" || tokenInfo.ExpiresAt <= 0 {
return requestExpiresAt, requestAutoPause
}
expiresAt := tokenInfo.ExpiresAt
if requestExpiresAt != nil && *requestExpiresAt > 0 && *requestExpiresAt < expiresAt {
expiresAt = *requestExpiresAt
}
autoPause := true
return &expiresAt, &autoPause
}
func cloneGrokSSOMap(source map[string]any) map[string]any {
if source == nil {
return nil
}
clone := make(map[string]any, len(source))
for key, value := range source {
clone[key] = cloneGrokSSOValue(value)
}
return clone
}
func cloneGrokSSOValue(value any) any {
switch v := value.(type) {
case map[string]any:
return cloneGrokSSOMap(v)
case []any:
clone := make([]any, len(v))
for i, item := range v {
clone[i] = cloneGrokSSOValue(item)
}
return clone
default:
return value
}
}
func normalizeSSOImportTokens(tokens []string, single string) []string {
items := make([]string, 0, len(tokens)+1)
if strings.TrimSpace(single) != "" {
items = append(items, single)
}
items = append(items, tokens...)
seen := make(map[string]struct{}, len(items))
result := make([]string, 0, len(items))
for _, item := range items {
parts := strings.Split(strings.NewReplacer(",", "\n", "\r", "\n").Replace(item), "\n")
for _, token := range parts {
if token = xai.NormalizeSSOToken(token); token == "" {
continue
}
if _, ok := seen[token]; ok {
continue
}
seen[token] = struct{}{}
result = append(result, token)
}
}
return result
}
func grokSSOImportAccountName(base string, tokenInfo *service.GrokTokenInfo, index, total int) string {
base = strings.TrimSpace(base)
if base == "" && tokenInfo != nil {
base = strings.TrimSpace(tokenInfo.Email)
}
if base == "" {
base = "Grok OAuth Account"
}
if total > 1 {
return base + " #" + strconv.Itoa(index)
}
return base
}
func grokSSOImportErrorMessage(err error) string {
status := infraerrors.FromError(err)
if status == nil {
return ""
}
if status.Reason != "" {
return status.Reason + ": " + status.Message
}
return status.Message
}
func (h *GrokOAuthHandler) QueryQuota(c *gin.Context) {
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
@@ -215,7 +458,7 @@ func (h *GrokOAuthHandler) QueryQuota(c *gin.Context) {
response.BadRequest(c, "grok quota service is not enabled")
return
}
result, err := h.quotaService.ProbeUsage(c.Request.Context(), accountID)
result, err := h.quotaService.QueryQuota(c.Request.Context(), accountID)
if err != nil {
response.ErrorFrom(c, err)
return
@@ -8,6 +8,7 @@ import (
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
@@ -41,17 +42,35 @@ func (r *grokQuotaHandlerAccountRepo) UpdateExtra(_ context.Context, id int64, u
}
type grokQuotaHandlerUpstream struct {
resp *http.Response
lastReq *http.Request
lastBody []byte
mu sync.Mutex
requests []*http.Request
bodies [][]byte
}
func (u *grokQuotaHandlerUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
u.lastReq = req
var body []byte
if req.Body != nil {
u.lastBody, _ = io.ReadAll(req.Body)
body, _ = io.ReadAll(req.Body)
}
return u.resp, nil
u.mu.Lock()
u.requests = append(u.requests, req)
u.bodies = append(u.bodies, body)
u.mu.Unlock()
if req.URL.Path == "/v1/responses" {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"8"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
}, nil
}
payload := `{"config":{"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"}}`
if req.URL.RawQuery == "format=credits" {
payload = `{"config":{"currentPeriod":{"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"}}}`
}
return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(payload))}, nil
}
func (u *grokQuotaHandlerUpstream) DoWithTLS(
@@ -77,14 +96,7 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
},
}}
upstream := &grokQuotaHandlerUpstream{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"8"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
}}
upstream := &grokQuotaHandlerUpstream{}
quotaService := service.NewGrokQuotaService(repo, nil, service.NewGrokTokenProvider(repo, nil), upstream)
handler := NewGrokOAuthHandler(nil, nil, quotaService)
@@ -95,12 +107,23 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Contains(t, rec.Body.String(), `"source":"active_probe"`)
require.Contains(t, rec.Body.String(), `"source":"hybrid_probe"`)
require.Contains(t, rec.Body.String(), `"billing":`)
require.Contains(t, rec.Body.String(), `"snapshot":`)
require.Contains(t, rec.Body.String(), `"headers_observed":true`)
require.NotContains(t, rec.Body.String(), "access-token")
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Contains(t, string(upstream.lastBody), `"store":false`)
upstream.mu.Lock()
requests := append([]*http.Request(nil), upstream.requests...)
bodies := append([][]byte(nil), upstream.bodies...)
upstream.mu.Unlock()
require.Len(t, requests, 3)
for i, upstreamReq := range requests {
require.Equal(t, "Bearer access-token", upstreamReq.Header.Get("Authorization"))
if upstreamReq.URL.String() == xai.DefaultCLIBaseURL+"/responses" {
require.Contains(t, string(bodies[i]), `"model":"grok-4.5"`)
require.Contains(t, string(bodies[i]), `"store":false`)
}
}
require.NotNil(t, repo.updates[42])
}
@@ -145,3 +168,51 @@ func TestGrokOAuthHandlerRuntimeSanityDoesNotExposeSecrets(t *testing.T) {
require.NotContains(t, rec.Body.String(), "secret")
require.NotContains(t, rec.Body.String(), "client-secret-like-value")
}
func TestGrokSSOImportExpiryUsesTokenExpiryWithoutRefreshToken(t *testing.T) {
tokenExpiry := time.Now().Add(6 * time.Hour).Unix()
expiresAt, autoPause := grokSSOImportExpiry(nil, nil, &service.GrokTokenInfo{
ExpiresAt: tokenExpiry,
})
require.NotNil(t, expiresAt)
require.Equal(t, tokenExpiry, *expiresAt)
require.NotNil(t, autoPause)
require.True(t, *autoPause)
}
func TestGrokSSOImportExpiryUsesEarlierRequestedExpiryWithoutRefreshToken(t *testing.T) {
requestedExpiry := time.Now().Add(2 * time.Hour).Unix()
tokenExpiry := time.Now().Add(6 * time.Hour).Unix()
requestedAutoPause := false
expiresAt, autoPause := grokSSOImportExpiry(&requestedExpiry, &requestedAutoPause, &service.GrokTokenInfo{
ExpiresAt: tokenExpiry,
})
require.NotNil(t, expiresAt)
require.Equal(t, requestedExpiry, *expiresAt)
require.NotNil(t, autoPause)
require.True(t, *autoPause)
}
func TestGrokSSOImportExpiryPreservesRequestSettingsWithRefreshToken(t *testing.T) {
requestedExpiry := time.Now().Add(2 * time.Hour).Unix()
requestedAutoPause := false
expiresAt, autoPause := grokSSOImportExpiry(&requestedExpiry, &requestedAutoPause, &service.GrokTokenInfo{
RefreshToken: "refresh-token",
ExpiresAt: time.Now().Add(6 * time.Hour).Unix(),
})
require.Same(t, &requestedExpiry, expiresAt)
require.Same(t, &requestedAutoPause, autoPause)
}
func TestGrokSSOImportWorkerRecoversPanic(t *testing.T) {
h := &GrokOAuthHandler{}
result := h.safeCreateAccountFromSSOToken(context.Background(), GrokSSOToOAuthRequest{}, "token", 2, 3)
// Without a service, createAccountFromSSOToken would panic on nil service access.
// Recovery must convert that into a failed item and keep the worker alive.
require.False(t, result.created)
require.Equal(t, 2, result.item.Index)
require.Contains(t, result.item.Error, "internal worker panic")
}
@@ -304,6 +304,10 @@ func (h *OpenAIOAuthHandler) CreateAccountFromCodexPAT(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
if err := service.ValidateOpenAILongContextBillingExtra(service.PlatformOpenAI, req.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
if req.Concurrency != nil && *req.Concurrency < 0 {
response.BadRequest(c, "concurrency must be >= 0")
return
@@ -15,6 +15,7 @@ import (
type opsSystemLogCleanupRequest struct {
StartTime string `json:"start_time"`
EndTime string `json:"end_time"`
Host string `json:"host"`
Level string `json:"level"`
Component string `json:"component"`
@@ -56,6 +57,7 @@ func (h *OpsHandler) ListSystemLogs(c *gin.Context) {
PageSize: pageSize,
StartTime: &start,
EndTime: &end,
Host: strings.TrimSpace(c.Query("host")),
Level: strings.TrimSpace(c.Query("level")),
Component: strings.TrimSpace(c.Query("component")),
RequestID: strings.TrimSpace(c.Query("request_id")),
@@ -153,6 +155,7 @@ func (h *OpsHandler) CleanupSystemLogs(c *gin.Context) {
filter := &service.OpsSystemLogCleanupFilter{
StartTime: start,
EndTime: end,
Host: strings.TrimSpace(req.Host),
Level: strings.TrimSpace(req.Level),
Component: strings.TrimSpace(req.Component),
RequestID: strings.TrimSpace(req.RequestID),
@@ -2,6 +2,7 @@ package admin
import (
"bytes"
"context"
"encoding/json"
"net/http"
"net/http/httptest"
@@ -19,6 +20,26 @@ type responseEnvelope struct {
Data json.RawMessage `json:"data"`
}
type opsSystemLogCaptureRepo struct {
service.OpsRepository
listFilter *service.OpsSystemLogFilter
cleanupFilter *service.OpsSystemLogCleanupFilter
}
func (r *opsSystemLogCaptureRepo) ListSystemLogs(_ context.Context, filter *service.OpsSystemLogFilter) (*service.OpsSystemLogList, error) {
r.listFilter = filter
return &service.OpsSystemLogList{Logs: []*service.OpsSystemLog{}, Page: filter.Page, PageSize: filter.PageSize}, nil
}
func (r *opsSystemLogCaptureRepo) DeleteSystemLogs(_ context.Context, filter *service.OpsSystemLogCleanupFilter) (int64, error) {
r.cleanupFilter = filter
return 1, nil
}
func (r *opsSystemLogCaptureRepo) InsertSystemLogCleanupAudit(_ context.Context, _ *service.OpsSystemLogCleanupAudit) error {
return nil
}
func newOpsSystemLogTestRouter(handler *OpsHandler, withUser bool) *gin.Engine {
gin.SetMode(gin.TestMode)
r := gin.New()
@@ -121,6 +142,23 @@ func TestOpsSystemLogHandler_ListSuccess(t *testing.T) {
}
}
func TestOpsSystemLogHandler_ListAcceptsHost(t *testing.T) {
repo := &opsSystemLogCaptureRepo{}
svc := service.NewOpsService(repo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
h := NewOpsHandler(svc)
r := newOpsSystemLogTestRouter(h, false)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/logs?host=api-node-1", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d, want 200", w.Code)
}
if repo.listFilter == nil || repo.listFilter.Host != "api-node-1" {
t.Fatalf("host filter = %+v, want api-node-1", repo.listFilter)
}
}
func TestOpsSystemLogHandler_CleanupUnauthorized(t *testing.T) {
svc := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
h := NewOpsHandler(svc)
@@ -205,6 +243,24 @@ func TestOpsSystemLogHandler_CleanupAcceptsAPIKeyID(t *testing.T) {
}
}
func TestOpsSystemLogHandler_CleanupAcceptsHost(t *testing.T) {
repo := &opsSystemLogCaptureRepo{}
svc := service.NewOpsService(repo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
h := NewOpsHandler(svc)
r := newOpsSystemLogTestRouter(h, true)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/logs/cleanup", bytes.NewBufferString(`{"host":"api-node-1"}`))
req.Header.Set("Content-Type", "application/json")
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d, want 200", w.Code)
}
if repo.cleanupFilter == nil || repo.cleanupFilter.Host != "api-node-1" {
t.Fatalf("host filter = %+v, want api-node-1", repo.cleanupFilter)
}
}
func TestOpsSystemLogHandler_CleanupInvalidAPIKeyID(t *testing.T) {
svc := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
h := NewOpsHandler(svc)
@@ -16,6 +16,7 @@ import (
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
@@ -323,7 +324,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) {
// If realtime monitoring is disabled, prefer a successful WS upgrade followed by a clean close
// with a deterministic close code. This prevents clients from spinning on 404/1006 reconnect loops.
if !h.opsService.IsRealtimeMonitoringEnabled(c.Request.Context()) {
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c))
if err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "ops realtime monitoring is disabled"})
return
@@ -358,7 +359,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) {
defer releaseOpsWSIPSlot(clientIP)
}
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c))
if err != nil {
logger.LegacyPrintf("handler.admin.ops_ws", "[OpsWS] upgrade failed: %v", err)
return
+49 -48
View File
@@ -599,54 +599,55 @@ func usageLogFromServiceUser(l *service.UsageLog) UsageLog {
requestedModel = l.Model
}
return UsageLog{
ID: l.ID,
UserID: l.UserID,
APIKeyID: l.APIKeyID,
AccountID: l.AccountID,
RequestID: l.RequestID,
Model: requestedModel,
ServiceTier: l.ServiceTier,
ReasoningEffort: l.ReasoningEffort,
InboundEndpoint: l.InboundEndpoint,
GroupID: l.GroupID,
SubscriptionID: l.SubscriptionID,
InputTokens: l.InputTokens,
OutputTokens: l.OutputTokens,
CacheCreationTokens: l.CacheCreationTokens,
CacheReadTokens: l.CacheReadTokens,
CacheCreation5mTokens: l.CacheCreation5mTokens,
CacheCreation1hTokens: l.CacheCreation1hTokens,
InputCost: l.InputCost,
OutputCost: l.OutputCost,
CacheCreationCost: l.CacheCreationCost,
CacheReadCost: l.CacheReadCost,
TotalCost: l.TotalCost,
ActualCost: l.ActualCost,
RateMultiplier: l.RateMultiplier,
BillingType: l.BillingType,
RequestType: requestType.String(),
Stream: stream,
OpenAIWSMode: openAIWSMode,
DurationMs: l.DurationMs,
FirstTokenMs: l.FirstTokenMs,
ImageCount: l.ImageCount,
ImageSize: l.ImageSize,
ImageInputSize: l.ImageInputSize,
ImageOutputSize: l.ImageOutputSize,
ImageOutputTokens: l.ImageOutputTokens,
ImageOutputCost: l.ImageOutputCost,
ImageSizeSource: l.ImageSizeSource,
ImageSizeBreakdown: l.ImageSizeBreakdown,
MediaType: l.MediaType,
UserAgent: l.UserAgent,
IPAddress: l.IPAddress,
CacheTTLOverridden: l.CacheTTLOverridden,
BillingMode: l.BillingMode,
CreatedAt: l.CreatedAt,
User: UserFromServiceShallow(l.User),
APIKey: APIKeyFromService(l.APIKey),
Group: GroupFromServiceShallow(l.Group),
Subscription: UserSubscriptionFromService(l.Subscription),
ID: l.ID,
UserID: l.UserID,
APIKeyID: l.APIKeyID,
AccountID: l.AccountID,
RequestID: l.RequestID,
Model: requestedModel,
ServiceTier: l.ServiceTier,
ReasoningEffort: l.ReasoningEffort,
InboundEndpoint: l.InboundEndpoint,
GroupID: l.GroupID,
SubscriptionID: l.SubscriptionID,
InputTokens: l.InputTokens,
OutputTokens: l.OutputTokens,
CacheCreationTokens: l.CacheCreationTokens,
CacheReadTokens: l.CacheReadTokens,
CacheCreation5mTokens: l.CacheCreation5mTokens,
CacheCreation1hTokens: l.CacheCreation1hTokens,
InputCost: l.InputCost,
OutputCost: l.OutputCost,
CacheCreationCost: l.CacheCreationCost,
CacheReadCost: l.CacheReadCost,
TotalCost: l.TotalCost,
ActualCost: l.ActualCost,
RateMultiplier: l.RateMultiplier,
LongContextBillingApplied: l.LongContextBillingApplied,
BillingType: l.BillingType,
RequestType: requestType.String(),
Stream: stream,
OpenAIWSMode: openAIWSMode,
DurationMs: l.DurationMs,
FirstTokenMs: l.FirstTokenMs,
ImageCount: l.ImageCount,
ImageSize: l.ImageSize,
ImageInputSize: l.ImageInputSize,
ImageOutputSize: l.ImageOutputSize,
ImageOutputTokens: l.ImageOutputTokens,
ImageOutputCost: l.ImageOutputCost,
ImageSizeSource: l.ImageSizeSource,
ImageSizeBreakdown: l.ImageSizeBreakdown,
MediaType: l.MediaType,
UserAgent: l.UserAgent,
IPAddress: l.IPAddress,
CacheTTLOverridden: l.CacheTTLOverridden,
BillingMode: l.BillingMode,
CreatedAt: l.CreatedAt,
User: UserFromServiceShallow(l.User),
APIKey: APIKeyFromService(l.APIKey),
Group: GroupFromServiceShallow(l.Group),
Subscription: UserSubscriptionFromService(l.Subscription),
}
}
+8 -7
View File
@@ -482,13 +482,14 @@ type UsageLog struct {
CacheCreation5mTokens int `json:"cache_creation_5m_tokens"`
CacheCreation1hTokens int `json:"cache_creation_1h_tokens"`
InputCost float64 `json:"input_cost"`
OutputCost float64 `json:"output_cost"`
CacheCreationCost float64 `json:"cache_creation_cost"`
CacheReadCost float64 `json:"cache_read_cost"`
TotalCost float64 `json:"total_cost"`
ActualCost float64 `json:"actual_cost"`
RateMultiplier float64 `json:"rate_multiplier"`
InputCost float64 `json:"input_cost"`
OutputCost float64 `json:"output_cost"`
CacheCreationCost float64 `json:"cache_creation_cost"`
CacheReadCost float64 `json:"cache_read_cost"`
TotalCost float64 `json:"total_cost"`
ActualCost float64 `json:"actual_cost"`
RateMultiplier float64 `json:"rate_multiplier"`
LongContextBillingApplied bool `json:"long_context_billing_applied"`
BillingType int8 `json:"billing_type"`
RequestType string `json:"request_type"`
+7 -1
View File
@@ -24,6 +24,8 @@ const (
EndpointImagesGenerations = "/v1/images/generations"
EndpointImagesEdits = "/v1/images/edits"
EndpointVideosGenerations = "/v1/videos/generations"
EndpointVideosEdits = "/v1/videos/edits"
EndpointVideosExtensions = "/v1/videos/extensions"
EndpointVideos = "/v1/videos"
EndpointGeminiModels = "/v1beta/models"
)
@@ -88,6 +90,10 @@ func NormalizeInboundEndpoint(path string) string {
return EndpointImagesEdits
case strings.Contains(path, EndpointVideosGenerations) || strings.Contains(path, "/videos/generations"):
return EndpointVideosGenerations
case strings.Contains(path, EndpointVideosEdits) || strings.Contains(path, "/videos/edits"):
return EndpointVideosEdits
case strings.Contains(path, EndpointVideosExtensions) || strings.Contains(path, "/videos/extensions"):
return EndpointVideosExtensions
case strings.Contains(path, EndpointVideos) || strings.Contains(path, "/videos/"):
return EndpointVideos
case strings.Contains(path, EndpointResponsesCompact) || isResponsesCompactAliasPath(path):
@@ -173,7 +179,7 @@ func DeriveUpstreamEndpoint(inbound, rawRequestPath, platform string) string {
switch platform {
case service.PlatformOpenAI, service.PlatformGrok:
if inbound == EndpointEmbeddings || inbound == EndpointAlphaSearch || inbound == EndpointImagesGenerations || inbound == EndpointImagesEdits || inbound == EndpointVideosGenerations || inbound == EndpointVideos {
if inbound == EndpointEmbeddings || inbound == EndpointAlphaSearch || inbound == EndpointImagesGenerations || inbound == EndpointImagesEdits || inbound == EndpointVideosGenerations || inbound == EndpointVideosEdits || inbound == EndpointVideosExtensions || inbound == EndpointVideos {
return inbound
}
// OpenAI forwards everything to the Responses API.
+7 -4
View File
@@ -29,7 +29,8 @@ const (
)
const (
// maxSameAccountRetries 同账号重试次数上限(针对 RetryableOnSameAccount 错误)
// maxSameAccountRetries 同账号重试次数默认上限(针对 RetryableOnSameAccount 错误)
// 生产调用方通常传入账号级配置 account.GetPoolModeRetryCount(),该常量仅作兜底/测试默认值。
maxSameAccountRetries = 3
// sameAccountRetryDelay 同账号重试间隔
sameAccountRetryDelay = 500 * time.Millisecond
@@ -67,6 +68,7 @@ func (s *FailoverState) HandleFailoverError(
gatewayService TempUnscheduler,
accountID int64,
platform string,
retryLimit int,
failoverErr *service.UpstreamFailoverError,
) FailoverAction {
s.LastFailoverErr = failoverErr
@@ -76,14 +78,15 @@ func (s *FailoverState) HandleFailoverError(
s.ForceCacheBilling = true
}
// 同账号重试:对 RetryableOnSameAccount 的临时性错误,先在同一账号上重试
if failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < maxSameAccountRetries {
// 同账号重试:对 RetryableOnSameAccount 的临时性错误,先在同一账号上重试
// 重试次数上限 retryLimit 由调用方传入(账号级 pool_mode_retry_count 配置)。
if failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < retryLimit {
s.SameAccountRetryCount[accountID]++
logger.FromContext(ctx).Warn("gateway.failover_same_account_retry",
zap.Int64("account_id", accountID),
zap.Int("upstream_status", failoverErr.StatusCode),
zap.Int("same_account_retry_count", s.SameAccountRetryCount[accountID]),
zap.Int("same_account_retry_max", maxSameAccountRetries),
zap.Int("same_account_retry_max", retryLimit),
)
if !sleepWithContext(ctx, sameAccountRetryDelay) {
return FailoverCanceled
+87 -50
View File
@@ -133,7 +133,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
@@ -150,7 +150,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
err := newTestFailoverErr(500, false, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
@@ -166,7 +166,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
err := newTestFailoverErr(500, false, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
@@ -181,19 +181,19 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
// 第一次切换:0→1
err1 := newTestFailoverErr(500, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
// 第二次切换:1→2
err2 := newTestFailoverErr(502, false, false)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 2, fs.SwitchCount)
// 第三次已耗尽:SwitchCount(2) >= MaxSwitches(2)
err3 := newTestFailoverErr(503, false, false)
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", err3)
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, err3)
require.Equal(t, FailoverExhausted, action)
require.Equal(t, 2, fs.SwitchCount, "耗尽时不应继续递增")
@@ -212,7 +212,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
fs := NewFailoverState(0, false)
err := newTestFailoverErr(500, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverExhausted, action)
require.Equal(t, 0, fs.SwitchCount)
require.Contains(t, fs.FailedAccountIDs, int64(100))
@@ -229,7 +229,7 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
fs := NewFailoverState(3, true) // hasBoundSession=true
err := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.True(t, fs.ForceCacheBilling)
})
@@ -238,7 +238,7 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, true) // ForceCacheBilling=true
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.True(t, fs.ForceCacheBilling)
})
@@ -247,7 +247,7 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.False(t, fs.ForceCacheBilling)
})
@@ -257,12 +257,12 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
// 第一次:ForceCacheBilling=true → 设置
err1 := newTestFailoverErr(500, false, true)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.True(t, fs.ForceCacheBilling)
// 第二次:ForceCacheBilling=false → 仍然保持 true
err2 := newTestFailoverErr(502, false, false)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.True(t, fs.ForceCacheBilling, "ForceCacheBilling 一旦设置不应被重置")
})
}
@@ -278,7 +278,7 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
@@ -297,7 +297,7 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
for i := 1; i <= maxSameAccountRetries; i++ {
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, i, fs.SameAccountRetryCount[100])
}
@@ -311,12 +311,12 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
for i := 0; i < maxSameAccountRetries; i++ {
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
}
require.Equal(t, maxSameAccountRetries, fs.SameAccountRetryCount[100])
// 第 maxSameAccountRetries+1 次:重试耗尽,应切换账号
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
require.Contains(t, fs.FailedAccountIDs, int64(100))
@@ -333,12 +333,12 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
// 账号 100 第一次重试
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[100])
// 账号 200 第一次重试(独立计数)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", err)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[200])
require.Equal(t, 1, fs.SameAccountRetryCount[100], "账号 100 的计数不应受影响")
@@ -351,17 +351,54 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
// 耗尽账号 100 的重试
for i := 0; i < maxSameAccountRetries; i++ {
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
}
// 第 maxSameAccountRetries+1 次: 重试耗尽 → 切换
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
// 再次遇到账号 100,计数仍为 maxSameAccountRetries,条件不满足 → 直接切换
action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Len(t, mock.calls, 2, "第二次耗尽也应调用 TempUnschedule")
})
t.Run("尊重账号级retryLimit_配置1次只重试1次", func(t *testing.T) {
// 回归测试:Anthropic 等路径此前硬编码同账号重试 3 次,忽略账号
// pool_mode_retry_count 配置。此处验证传入 retryLimit=1 时只重试 1 次即切换。
mock := &mockTempUnscheduler{}
fs := NewFailoverState(5, false)
err := newTestFailoverErr(403, true, false)
const retryLimit = 1
// 第 1 次:同账号重试
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryLimit, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[100])
require.Equal(t, 0, fs.SwitchCount, "首次重试不应切换账号")
require.Empty(t, mock.calls, "未耗尽前不应 TempUnschedule")
// 第 2 次:已达上限 1 → 不再同账号重试,直接切换 + TempUnschedule
action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryLimit, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[100], "重试计数不应超过 retryLimit")
require.Equal(t, 1, fs.SwitchCount, "重试耗尽应切换账号")
require.Contains(t, fs.FailedAccountIDs, int64(100))
require.Len(t, mock.calls, 1, "重试耗尽应触发 TempUnschedule")
})
t.Run("retryLimit为0时立即切换不重试", func(t *testing.T) {
// pool_mode_retry_count=0 表示关闭同账号重试(如 GPT Image 账号)。
mock := &mockTempUnscheduler{}
fs := NewFailoverState(5, false)
err := newTestFailoverErr(403, true, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", 0, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 0, fs.SameAccountRetryCount[100], "retryLimit=0 不应发生同账号重试")
require.Equal(t, 1, fs.SwitchCount, "应立即切换账号")
require.Len(t, mock.calls, 1, "应立即 TempUnschedule")
})
}
// ---------------------------------------------------------------------------
@@ -374,7 +411,7 @@ func TestHandleFailoverError_TempUnschedule(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, false) // RetryableOnSameAccount=false
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Empty(t, mock.calls)
})
@@ -384,10 +421,10 @@ func TestHandleFailoverError_TempUnschedule(t *testing.T) {
err := newTestFailoverErr(502, true, false)
for i := 0; i < maxSameAccountRetries; i++ {
fs.HandleFailoverError(context.Background(), mock, 42, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 42, "openai", maxSameAccountRetries, err)
}
// 再次触发时才会执行 TempUnschedule + 切换
fs.HandleFailoverError(context.Background(), mock, 42, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 42, "openai", maxSameAccountRetries, err)
require.Len(t, mock.calls, 1)
require.Equal(t, int64(42), mock.calls[0].accountID)
@@ -410,7 +447,7 @@ func TestHandleFailoverError_ContextCanceled(t *testing.T) {
cancel() // 立即取消
start := time.Now()
action := fs.HandleFailoverError(ctx, mock, 100, "openai", err)
action := fs.HandleFailoverError(ctx, mock, 100, "openai", maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverCanceled, action)
@@ -429,7 +466,7 @@ func TestHandleFailoverError_ContextCanceled(t *testing.T) {
cancel() // 立即取消
start := time.Now()
action := fs.HandleFailoverError(ctx, mock, 100, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(ctx, mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverCanceled, action)
@@ -446,10 +483,10 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(3, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
require.Contains(t, fs.FailedAccountIDs, int64(100))
fs.HandleFailoverError(context.Background(), mock, 200, "openai", newTestFailoverErr(502, false, false))
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, newTestFailoverErr(502, false, false))
require.Contains(t, fs.FailedAccountIDs, int64(200))
require.Len(t, fs.FailedAccountIDs, 2)
})
@@ -458,7 +495,7 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(0, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
require.Equal(t, FailoverExhausted, action)
require.Contains(t, fs.FailedAccountIDs, int64(100))
})
@@ -467,7 +504,7 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(3, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(400, true, false))
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(400, true, false))
require.Equal(t, FailoverContinue, action)
require.NotContains(t, fs.FailedAccountIDs, int64(100))
})
@@ -476,8 +513,8 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(5, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
require.Len(t, fs.FailedAccountIDs, 1, "map 天然去重")
})
}
@@ -492,11 +529,11 @@ func TestHandleFailoverError_LastFailoverErr(t *testing.T) {
fs := NewFailoverState(3, false)
err1 := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.Equal(t, err1, fs.LastFailoverErr)
err2 := newTestFailoverErr(502, false, false)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.Equal(t, err2, fs.LastFailoverErr)
})
@@ -505,7 +542,7 @@ func TestHandleFailoverError_LastFailoverErr(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(400, true, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, err, fs.LastFailoverErr)
})
}
@@ -522,30 +559,30 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
// 1. 账号 100 遇到可重试错误,同账号重试 maxSameAccountRetries 次
retryErr := newTestFailoverErr(400, true, false)
for i := 0; i < maxSameAccountRetries; i++ {
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryErr)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
require.Equal(t, FailoverContinue, action)
}
require.True(t, fs.ForceCacheBilling, "hasBoundSession=true 应设置 ForceCacheBilling")
// 2. 账号 100 超过重试上限 → TempUnschedule + 切换
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryErr)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
require.Len(t, mock.calls, 1)
// 3. 账号 200 遇到不可重试错误 → 直接切换
switchErr := newTestFailoverErr(500, false, false)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", switchErr)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, switchErr)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 2, fs.SwitchCount)
// 4. 账号 300 遇到不可重试错误 → 再切换
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", switchErr)
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, switchErr)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 3, fs.SwitchCount)
// 5. 账号 400 → 已耗尽 (SwitchCount=3 >= MaxSwitches=3)
action = fs.HandleFailoverError(context.Background(), mock, 400, "openai", switchErr)
action = fs.HandleFailoverError(context.Background(), mock, 400, "openai", maxSameAccountRetries, switchErr)
require.Equal(t, FailoverExhausted, action)
// 最终状态验证
@@ -563,21 +600,21 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
// 第一次切换:delay = 0s
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
require.Less(t, elapsed, 200*time.Millisecond, "第一次切换延迟为 0")
// 第二次切换:delay = 1s
start = time.Now()
action = fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, err)
action = fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed = time.Since(start)
require.Equal(t, FailoverContinue, action)
require.GreaterOrEqual(t, elapsed, 800*time.Millisecond, "第二次切换延迟约 1s")
// 第三次:耗尽(无延迟,因为在检查延迟之前就返回了)
start = time.Now()
action = fs.HandleFailoverError(context.Background(), mock, 300, service.PlatformAntigravity, err)
action = fs.HandleFailoverError(context.Background(), mock, 300, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed = time.Since(start)
require.Equal(t, FailoverExhausted, action)
require.Less(t, elapsed, 200*time.Millisecond, "耗尽时不应有延迟")
@@ -589,17 +626,17 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
// 第一次:ForceCacheBilling=false
err1 := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.False(t, fs.ForceCacheBilling)
// 第二次:ForceCacheBilling=trueAntigravity 粘性会话切换)
err2 := newTestFailoverErr(500, false, true)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.True(t, fs.ForceCacheBilling, "错误标志应触发 ForceCacheBilling")
// 第三次:ForceCacheBilling=false,但状态仍保持 true
err3 := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 300, "openai", err3)
fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, err3)
require.True(t, fs.ForceCacheBilling, "不应重置")
})
}
@@ -614,7 +651,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(0, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
})
@@ -623,7 +660,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, true, false)
action := fs.HandleFailoverError(context.Background(), mock, 0, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 0, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[0])
})
@@ -633,7 +670,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, true, false)
action := fs.HandleFailoverError(context.Background(), mock, -1, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, -1, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[-1])
})
@@ -645,7 +682,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
err := newTestFailoverErr(500, false, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, "", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "", maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
+10 -9
View File
@@ -448,7 +448,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
h.handleFailoverExhausted(c, failoverErr, service.PlatformGemini, true)
return
}
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -868,7 +868,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
h.handleFailoverExhausted(c, failoverErr, account.Platform, true)
return
}
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -1454,13 +1454,14 @@ func (h *GatewayHandler) usageUnrestricted(c *gin.Context, ctx context.Context,
remaining := h.calculateSubscriptionRemaining(apiKey.Group, subscription)
resp["remaining"] = remaining
resp["subscription"] = gin.H{
"daily_usage_usd": subscription.DailyUsageUSD,
"weekly_usage_usd": subscription.WeeklyUsageUSD,
"monthly_usage_usd": subscription.MonthlyUsageUSD,
"daily_limit_usd": apiKey.Group.DailyLimitUSD,
"weekly_limit_usd": apiKey.Group.WeeklyLimitUSD,
"monthly_limit_usd": apiKey.Group.MonthlyLimitUSD,
"expires_at": subscription.ExpiresAt,
"daily_usage_usd": subscription.DailyUsageUSD,
"weekly_usage_usd": subscription.WeeklyUsageUSD,
"monthly_usage_usd": subscription.MonthlyUsageUSD,
"daily_limit_usd": apiKey.Group.DailyLimitUSD,
"weekly_limit_usd": apiKey.Group.WeeklyLimitUSD,
"monthly_limit_usd": apiKey.Group.MonthlyLimitUSD,
"weekly_window_start": subscription.WeeklyWindowStart,
"expires_at": subscription.ExpiresAt,
}
}
@@ -254,7 +254,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
h.handleCCFailoverExhausted(c, failoverErr, true)
return
}
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -233,7 +233,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
h.handleResponsesFailoverExhausted(c, failoverErr, true)
return
}
action := fs.HandleFailoverError(requestCtx, h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(requestCtx, h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -0,0 +1,51 @@
package handler
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestUsageUnrestrictedIncludesWeeklyWindowStart(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/usage", nil)
weeklyWindowStart := time.Date(2026, time.July, 13, 0, 30, 0, 0, time.FixedZone("UTC+8", 8*60*60))
c.Set(string(middleware.ContextKeySubscription), &service.UserSubscription{
WeeklyWindowStart: &weeklyWindowStart,
})
handler := &GatewayHandler{}
handler.usageUnrestricted(
c,
context.Background(),
&service.APIKey{Group: &service.Group{
Name: "Weekly plan",
SubscriptionType: service.SubscriptionTypeSubscription,
}},
middleware.AuthSubject{},
nil,
nil,
nil,
)
require.Equal(t, http.StatusOK, recorder.Code)
var response struct {
Subscription struct {
WeeklyWindowStart *time.Time `json:"weekly_window_start"`
} `json:"subscription"`
}
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response))
require.NotNil(t, response.Subscription.WeeklyWindowStart)
require.True(t, weeklyWindowStart.Equal(*response.Subscription.WeeklyWindowStart))
}
@@ -220,6 +220,15 @@ func (h *ConcurrencyHelper) TryAcquireUserSlotForAPIKey(ctx context.Context, use
return h.withAPIKeySlot(ctx, apiKeyID, releaseFunc), true, nil
}
// AcquireOpenAIWSIngressLease bounds the whole client WebSocket lifecycle,
// independently from per-turn user and account slots.
func (h *ConcurrencyHelper) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int) (*service.OpenAIWSIngressLease, bool, error) {
if h == nil || h.concurrencyService == nil {
return nil, false, fmt.Errorf("concurrency service is unavailable")
}
return h.concurrencyService.AcquireOpenAIWSIngressLease(ctx, apiKeyID, maxConnections)
}
// TryAcquireAccountSlot 尝试立即获取账号并发槽位。
// 返回值: (releaseFunc, acquired, error)
func (h *ConcurrencyHelper) TryAcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int) (func(), bool, error) {
@@ -11,10 +11,13 @@ import (
)
type concurrencyCacheMock struct {
acquireUserSlotFn func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error)
acquireAccountSlotFn func(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error)
releaseUserCalled int32
releaseAccountCalled int32
acquireUserSlotFn func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error)
acquireAccountSlotFn func(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error)
acquireIngressLeaseFn func(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error)
releaseIngressLeaseFn func(ctx context.Context, apiKeyID int64, leaseID string) error
releaseUserCalled int32
releaseAccountCalled int32
releaseIngressCalled int32
}
func (m *concurrencyCacheMock) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
@@ -97,6 +100,25 @@ func (m *concurrencyCacheMock) CleanupStaleProcessSlots(ctx context.Context, act
return nil
}
func (m *concurrencyCacheMock) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error) {
if m.acquireIngressLeaseFn != nil {
return m.acquireIngressLeaseFn(ctx, apiKeyID, maxConnections, leaseID)
}
return false, nil
}
func (m *concurrencyCacheMock) RefreshOpenAIWSIngressLease(context.Context, int64, string) (bool, error) {
return true, nil
}
func (m *concurrencyCacheMock) ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error {
atomic.AddInt32(&m.releaseIngressCalled, 1)
if m.releaseIngressLeaseFn != nil {
return m.releaseIngressLeaseFn(ctx, apiKeyID, leaseID)
}
return nil
}
func TestConcurrencyHelper_TryAcquireUserSlot(t *testing.T) {
cache := &concurrencyCacheMock{
acquireUserSlotFn: func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
@@ -482,7 +482,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
if err != nil {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
failoverAction := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
failoverAction := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch failoverAction {
case FailoverContinue:
continue
+11 -1
View File
@@ -31,6 +31,16 @@ func (h *OpenAIGatewayHandler) GrokVideoGeneration(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideosGenerations, "")
}
// GrokVideoEdit handles asynchronous xAI video edits through Grok groups.
func (h *OpenAIGatewayHandler) GrokVideoEdit(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideosEdits, "")
}
// GrokVideoExtension handles asynchronous xAI video extensions through Grok groups.
func (h *OpenAIGatewayHandler) GrokVideoExtension(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideosExtensions, "")
}
// GrokVideoStatus handles xAI video status retrieval through Grok groups.
func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id"))
@@ -298,7 +308,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
if endpoint == service.GrokMediaEndpointVideosGenerations && strings.TrimSpace(result.ResponseID) != "" {
if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" {
if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil {
reqLog.Warn("grok_media.bind_video_request_account_failed",
zap.Int64("account_id", account.ID),
@@ -15,11 +15,13 @@ import (
// Codex CLI and the Codex desktop app refresh their model picker from
// GET {base_url}/models?client_version=... (custom provider mode) or
// GET /backend-api/codex/models (chatgpt_base_url mode). Both routes land
// here. The manifest is proxied verbatim from the ChatGPT backend with a
// schedulable OAuth account's credentials, so clients pointed at the gateway
// see the account's real, always-current model entitlements instead of a
// frozen local cache.
// here. The manifest is proxied verbatim from the selected account's ChatGPT
// backend or custom API key upstream. API key manifests use a short-lived,
// asynchronously revalidated cache to tolerate canceled client requests.
func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
if c.Request.Context().Err() != nil {
return
}
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
if !ok || apiKey.Group == nil {
h.errorResponse(c, http.StatusUnauthorized, "invalid_request_error", "API key group is required")
@@ -30,24 +32,54 @@ func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
return
}
account, err := h.gatewayService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, "", "")
if err != nil {
h.errorResponse(c, http.StatusServiceUnavailable, "upstream_error", "No available OpenAI accounts")
return
maxAccountSwitches := h.maxAccountSwitches
if maxAccountSwitches <= 0 {
maxAccountSwitches = 3
}
failedAccountIDs := make(map[int64]struct{})
switchCount := 0
var lastUpstreamErr error
manifest, err := h.gatewayService.FetchCodexModelsManifest(c.Request.Context(), account, c.Query("client_version"), c.GetHeader("If-None-Match"))
if err != nil {
h.errorResponse(c, infraerrors.Code(err), "upstream_error", infraerrors.Message(err))
return
}
for {
account, err := h.gatewayService.SelectAccountForModelWithExclusions(c.Request.Context(), apiKey.GroupID, "", "", failedAccountIDs)
if err != nil {
if c.Request.Context().Err() != nil {
return
}
if lastUpstreamErr != nil {
h.errorResponse(c, infraerrors.Code(lastUpstreamErr), "upstream_error", infraerrors.Message(lastUpstreamErr))
return
}
h.errorResponse(c, http.StatusServiceUnavailable, "upstream_error", "No available OpenAI accounts")
return
}
if manifest.ETag != "" {
c.Header("ETag", manifest.ETag)
}
if manifest.NotModified {
c.Status(http.StatusNotModified)
manifest, err := h.gatewayService.FetchCodexModelsManifest(c.Request.Context(), account, c.Query("client_version"), c.GetHeader("If-None-Match"))
if err != nil {
if c.Request.Context().Err() != nil {
return
}
if service.IsRetryableCodexModelsManifestError(err) && switchCount < maxAccountSwitches {
failedAccountIDs[account.ID] = struct{}{}
switchCount++
lastUpstreamErr = err
continue
}
h.errorResponse(c, infraerrors.Code(err), "upstream_error", infraerrors.Message(err))
return
}
if c.Request.Context().Err() != nil {
return
}
if manifest.ETag != "" {
c.Header("ETag", manifest.ETag)
}
if manifest.NotModified {
c.Status(http.StatusNotModified)
return
}
c.Data(http.StatusOK, "application/json", manifest.Body)
return
}
c.Data(http.StatusOK, "application/json", manifest.Body)
}
@@ -0,0 +1,288 @@
package handler
import (
"context"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
type codexModelsFailoverAccountRepo struct {
service.AccountRepository
accounts []service.Account
}
func (r codexModelsFailoverAccountRepo) GetByID(_ context.Context, id int64) (*service.Account, error) {
for i := range r.accounts {
if r.accounts[i].ID == id {
account := r.accounts[i]
return &account, nil
}
}
return nil, service.ErrNoAvailableAccounts
}
func (r codexModelsFailoverAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]service.Account, error) {
accounts := make([]service.Account, 0, len(r.accounts))
for _, account := range r.accounts {
if account.Platform == platform {
accounts = append(accounts, account)
}
}
return accounts, nil
}
type codexModelsFailoverHTTPUpstream struct {
service.HTTPUpstream
mu sync.Mutex
accountIDs []int64
firstErr error
firstStatus int
statuses map[int64]int
}
func (u *codexModelsFailoverHTTPUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
u.mu.Lock()
u.accountIDs = append(u.accountIDs, accountID)
u.mu.Unlock()
status, hasStatus := u.statuses[accountID]
if accountID == 1 || hasStatus {
if u.firstErr != nil {
return nil, u.firstErr
}
if !hasStatus {
status = u.firstStatus
}
if status == 0 {
status = http.StatusServiceUnavailable
}
return &http.Response{
StatusCode: status,
Status: http.StatusText(status),
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(
`{"error":{"message":"No available OpenAI accounts","type":"upstream_error"}}`,
)),
}, nil
}
return &http.Response{
StatusCode: http.StatusOK,
Status: "200 OK",
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(`{"models":[{"slug":"gpt-5.6-sol"}]}`)),
}, nil
}
func (u *codexModelsFailoverHTTPUpstream) calls() []int64 {
u.mu.Lock()
defer u.mu.Unlock()
return append([]int64(nil), u.accountIDs...)
}
func TestCodexModelsCanceledRequestDoesNotWriteResponse(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
ctx, cancel := context.WithCancel(context.Background())
cancel()
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil).WithContext(ctx)
h := &OpenAIGatewayHandler{}
h.CodexModels(c)
if c.Writer.Written() {
t.Fatalf("canceled request wrote an HTTP response: status=%d body=%q", recorder.Code, recorder.Body.String())
}
}
func TestCodexModelsFailsOverFromRetryableUpstreamStatus(t *testing.T) {
retryableStatuses := []int{
http.StatusTooManyRequests,
http.StatusInternalServerError,
http.StatusBadGateway,
http.StatusServiceUnavailable,
http.StatusGatewayTimeout,
}
for _, status := range retryableStatuses {
t.Run(http.StatusText(status), func(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(status)
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusOK {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
}
if got, want := recorder.Body.String(), `{"models":[{"slug":"gpt-5.6-sol"}]}`; got != want {
t.Fatalf("body: got %q, want %q", got, want)
}
})
}
}
func TestCodexModelsFailsOverFromUpstreamTransportError(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
upstream.firstErr = &net.OpError{
Op: "read",
Net: "tcp",
Err: errors.New("connection reset"),
}
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusOK {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
}
}
func TestCodexModelsDoesNotFailOverFromPermanentUpstreamStatus(t *testing.T) {
statuses := []int{
http.StatusBadRequest,
http.StatusUnauthorized,
http.StatusForbidden,
http.StatusNotFound,
600,
}
for _, status := range statuses {
t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(status)
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusBadGateway {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
})
}
}
func TestCodexModelsDoesNotFailOverFromUpstreamConfigurationError(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
upstream.firstErr = errors.New("invalid proxy URL")
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusBadGateway {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
}
func TestCodexModelsReturnsLastUpstreamErrorWhenAccountsAreExhausted(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
upstream.statuses = map[int64]int{
1: http.StatusServiceUnavailable,
2: http.StatusGatewayTimeout,
}
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusBadGateway {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
if body := recorder.Body.String(); !strings.Contains(body, "upstream error 504") {
t.Fatalf("body does not preserve the last upstream error: %s", body)
}
}
func TestCodexModelsHonorsAccountSwitchLimit(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandlerWithAccountCount(http.StatusServiceUnavailable, 4, 2)
upstream.statuses = map[int64]int{
1: http.StatusServiceUnavailable,
2: http.StatusBadGateway,
3: http.StatusGatewayTimeout,
4: http.StatusInternalServerError,
}
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1, 2, 3}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusBadGateway {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
if body := recorder.Body.String(); !strings.Contains(body, "upstream error 504") {
t.Fatalf("body does not preserve the limit-ending upstream error: %s", body)
}
}
func newCodexModelsFailoverTestHandler(firstStatus int) (*OpenAIGatewayHandler, *codexModelsFailoverHTTPUpstream, int64) {
return newCodexModelsFailoverTestHandlerWithAccountCount(firstStatus, 2, 3)
}
func newCodexModelsFailoverTestHandlerWithAccountCount(firstStatus, accountCount, maxSwitches int) (*OpenAIGatewayHandler, *codexModelsFailoverHTTPUpstream, int64) {
gin.SetMode(gin.TestMode)
groupID := int64(42)
accounts := make([]service.Account, 0, accountCount)
for i := 1; i <= accountCount; i++ {
accounts = append(accounts, service.Account{
ID: int64(i),
Name: fmt.Sprintf("upstream-%d", i),
Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey,
Status: service.StatusActive,
Schedulable: true,
Priority: i - 1,
Concurrency: 1,
Credentials: map[string]any{
"api_key": fmt.Sprintf("sk-%d", i),
"base_url": fmt.Sprintf("https://upstream-%d.example/v1", i),
},
})
}
upstream := &codexModelsFailoverHTTPUpstream{firstStatus: firstStatus}
cfg := &config.Config{RunMode: config.RunModeSimple}
gatewayService := service.NewOpenAIGatewayService(
codexModelsFailoverAccountRepo{accounts: accounts},
nil, nil, nil, nil, nil, nil, cfg, nil, nil, nil, nil, nil,
upstream,
nil, nil, nil, nil, nil, nil, nil, nil,
)
return &OpenAIGatewayHandler{gatewayService: gatewayService, maxAccountSwitches: maxSwitches}, upstream, groupID
}
func performCodexModelsRequest(t *testing.T, handler *OpenAIGatewayHandler, groupID int64) *httptest.ResponseRecorder {
t.Helper()
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models?client_version=0.144.0", nil)
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI},
})
handler.CodexModels(c)
return recorder
}
func equalInt64Slices(got, want []int64) bool {
if len(got) != len(want) {
return false
}
for i := range got {
if got[i] != want[i] {
return false
}
}
return true
}
@@ -1299,10 +1299,36 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
wsConn.SetReadLimit(service.ResolveOpenAIWSClientReadLimitBytes(h.cfg))
ctx := c.Request.Context()
maxIngressConnections := 0
if h.cfg != nil {
maxIngressConnections = h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey
}
ingressLease, ingressLeaseAcquired, ingressLeaseErr := h.concurrencyHelper.AcquireOpenAIWSIngressLease(ctx, apiKey.ID, maxIngressConnections)
if ingressLeaseErr != nil {
reqLog.Error("openai.websocket_ingress_lease_acquire_failed", zap.Error(ingressLeaseErr))
closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to reserve websocket ingress capacity")
return
}
if !ingressLeaseAcquired {
reqLog.Info("openai.websocket_ingress_capacity_rejected", zap.Int("max_ingress_connections_per_api_key", maxIngressConnections))
closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "too many open websocket connections, please retry later")
return
}
if ingressLease != nil {
defer ingressLease.Release()
ctx = ingressLease.Context()
c.Request = c.Request.WithContext(ctx)
}
readCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
msgType, firstMessage, err := wsConn.Read(readCtx)
cancel()
if err != nil {
if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) {
reqLog.Warn("openai.websocket_ingress_lease_lost_before_first_message", zap.Error(err))
closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "websocket ingress capacity lease lost; please reconnect")
return
}
closeStatus, closeReason := summarizeWSCloseErrorForLog(err)
reqLog.Warn("openai.websocket_read_first_message_failed",
zap.Error(err),
@@ -1692,6 +1718,25 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
continue
}
if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) {
reqLog.Warn("openai.websocket_ingress_lease_lost",
zap.Int64("account_id", account.ID),
zap.Error(err),
)
closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "websocket ingress capacity lease lost; please reconnect")
return
}
var closeErr *service.OpenAIWSClientCloseError
if errors.As(err, &closeErr) && closeErr.StatusCode() == coderws.StatusNormalClosure {
reqLog.Info("openai.websocket_ingress_closed_normally",
zap.Int64("account_id", account.ID),
zap.String("reason", closeErr.Reason()),
)
closeOpenAIClientWS(wsConn, closeErr.StatusCode(), closeErr.Reason())
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
closeStatus, closeReason := summarizeWSCloseErrorForLog(err)
reqLog.Warn("openai.websocket_proxy_failed",
@@ -1700,7 +1745,6 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
zap.String("close_status", closeStatus),
zap.String("close_reason", closeReason),
)
var closeErr *service.OpenAIWSClientCloseError
if errors.As(err, &closeErr) {
closeOpenAIClientWS(wsConn, closeErr.StatusCode(), closeErr.Reason())
return
@@ -2047,7 +2091,8 @@ func openAIForwardErrorAlreadyCommunicated(c *gin.Context, writerSizeBeforeForwa
}
// 与快照同口径:排除 compact 心跳字节,避免"仅心跳写出"被误判为
// 响应已写出(#3887)。
if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward {
if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward ||
service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward {
return false
}
@@ -8,6 +8,7 @@ import (
"net/http/httptest"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
@@ -414,11 +415,13 @@ func TestResolveOpenAIMessagesDispatchMappedModel(t *testing.T) {
SonnetMappedModel: "gpt-5.2",
ExactModelMappings: map[string]string{
"claude-sonnet-4-5-20250929": "gpt-5.4-mini-high",
"claude-fable-5": "gpt-5.6-sol",
},
},
},
}
require.Equal(t, "gpt-5.4-mini", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-sonnet-4-5-20250929"))
require.Equal(t, "gpt-5.6-sol", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-fable-5"))
})
t.Run("uses_family_default_when_no_override", func(t *testing.T) {
@@ -711,6 +714,68 @@ func TestOpenAIResponsesWebSocket_InvalidUpgradeDoesNotSetTransport(t *testing.T
require.Equal(t, service.OpenAIClientTransportUnknown, service.GetOpenAIClientTransport(c))
}
func TestOpenAIResponsesWebSocket_IngressCapacityRejected(t *testing.T) {
gin.SetMode(gin.TestMode)
cache := &concurrencyCacheMock{
acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
return false, nil
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
h.cfg = &config.Config{}
h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
defer wsServer.Close()
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
cancelDial()
require.NoError(t, err)
defer func() { _ = clientConn.CloseNow() }()
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, _, err = clientConn.Read(readCtx)
cancelRead()
var closeErr coderws.CloseError
require.ErrorAs(t, err, &closeErr)
require.Equal(t, coderws.StatusTryAgainLater, closeErr.Code)
}
func TestOpenAIResponsesWebSocket_IngressLeaseReleasedOnEarlyReturn(t *testing.T) {
gin.SetMode(gin.TestMode)
cache := &concurrencyCacheMock{
acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
return true, nil
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
h.cfg = &config.Config{}
h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
defer wsServer.Close()
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
cancelDial()
require.NoError(t, err)
defer func() { _ = clientConn.CloseNow() }()
writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
err = clientConn.Write(writeCtx, coderws.MessageBinary, []byte("not a response.create frame"))
cancelWrite()
require.NoError(t, err)
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, _, err = clientConn.Read(readCtx)
cancelRead()
var closeErr coderws.CloseError
require.ErrorAs(t, err, &closeErr)
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
require.Eventually(t, func() bool {
return atomic.LoadInt32(&cache.releaseIngressCalled) == 1
}, time.Second, 10*time.Millisecond)
}
func TestOpenAIResponsesWebSocket_RejectsMessageIDAsPreviousResponseID(t *testing.T) {
gin.SetMode(gin.TestMode)
+16 -2
View File
@@ -142,6 +142,9 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
failedAccountIDs := make(map[int64]struct{})
sameAccountRetryCount := make(map[int64]int)
var lastFailoverErr *service.UpstreamFailoverError
stopJSONKeepalive := func() {}
jsonKeepaliveStarted := false
defer func() { stopJSONKeepalive() }()
for {
reqLog.Debug("openai.images.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs)))
@@ -210,8 +213,12 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
}
service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds())
if !parsed.Stream && !jsonKeepaliveStarted {
stopJSONKeepalive = service.StartOpenAIImagesJSONKeepalive(c, h.openAIImagesJSONKeepaliveInterval())
jsonKeepaliveStarted = true
}
forwardStart := time.Now()
writerSizeBeforeForward := c.Writer.Size()
writerSizeBeforeForward := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)
result, err := func() (*service.OpenAIForwardResult, error) {
defer func() {
if accountReleaseFunc != nil {
@@ -258,7 +265,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
if c.Writer.Size() != writerSizeBeforeForward {
if service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward {
reqLog.Warn("openai.images.upstream_failover_skipped_after_flush",
zap.Int64("account_id", account.ID),
zap.Int("upstream_status", failoverErr.StatusCode),
@@ -383,6 +390,13 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
}
}
func (h *OpenAIGatewayHandler) openAIImagesJSONKeepaliveInterval() time.Duration {
if h.cfg == nil || h.cfg.Gateway.ImageNonstreamKeepaliveInterval <= 0 {
return 0
}
return time.Duration(h.cfg.Gateway.ImageNonstreamKeepaliveInterval) * time.Second
}
func isMultipartImagesContentType(contentType string) bool {
return strings.HasPrefix(strings.ToLower(strings.TrimSpace(contentType)), "multipart/form-data")
}
+1 -15
View File
@@ -9,7 +9,6 @@ import (
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/internal/payment"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -19,15 +18,13 @@ import (
// PaymentHandler handles user-facing payment requests.
type PaymentHandler struct {
channelService *service.ChannelService
paymentService *service.PaymentService
configService *service.PaymentConfigService
}
// NewPaymentHandler creates a new PaymentHandler.
func NewPaymentHandler(paymentService *service.PaymentService, configService *service.PaymentConfigService, channelService *service.ChannelService) *PaymentHandler {
func NewPaymentHandler(paymentService *service.PaymentService, configService *service.PaymentConfigService) *PaymentHandler {
return &PaymentHandler{
channelService: channelService,
paymentService: paymentService,
configService: configService,
}
@@ -91,17 +88,6 @@ func (h *PaymentHandler) GetPlans(c *gin.Context) {
response.Success(c, result)
}
// GetChannels returns enabled payment channels.
// GET /api/v1/payment/channels
func (h *PaymentHandler) GetChannels(c *gin.Context) {
channels, _, err := h.channelService.List(c.Request.Context(), pagination.PaginationParams{Page: 1, PageSize: 1000}, "active", "")
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, channels)
}
// GetCheckoutInfo returns all data the payment page needs in a single call:
// payment methods with limits, subscription plans, and configuration.
// GET /api/v1/payment/checkout-info
@@ -119,7 +119,7 @@ func TestVerifyOrderPublicReturnsLegacyOrderState(t *testing.T) {
require.NoError(t, err)
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, nil, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -219,7 +219,7 @@ func TestResolveOrderPublicByResumeTokenReturnsFrontendContractFields(t *testing
configSvc := service.NewPaymentConfigService(client, nil, []byte("0123456789abcdef0123456789abcdef"))
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, configSvc, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -307,7 +307,7 @@ func TestResolveOrderPublicByResumeTokenReturnsBadRequestForMismatchedToken(t *t
configSvc := service.NewPaymentConfigService(client, nil, []byte("0123456789abcdef0123456789abcdef"))
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, configSvc, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -347,7 +347,7 @@ func TestVerifyOrderPublicRejectsBlankOutTradeNo(t *testing.T) {
t.Cleanup(func() { _ = client.Close() })
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, nil, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
+1 -1
View File
@@ -164,7 +164,7 @@ var ProviderSet = wire.NewSet(
admin.NewDashboardHandler,
admin.NewUserHandler,
admin.NewGroupHandler,
admin.NewAccountHandler,
admin.ProvideAccountHandler,
admin.NewAnnouncementHandler,
admin.NewDataManagementHandler,
admin.NewBackupHandler,
+9 -9
View File
@@ -17,6 +17,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
)
// ForbiddenError 表示上游返回 403 Forbidden
@@ -279,7 +280,6 @@ func NewClient(proxyURL string) (*Client, error) {
}
client.Transport = transport
}
return &Client{
httpClient: client,
}, nil
@@ -341,7 +341,7 @@ func (c *Client) ExchangeCode(ctx context.Context, code, codeVerifier string) (*
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("token 交换请求失败: %w", err)
}
@@ -383,7 +383,7 @@ func (c *Client) RefreshToken(ctx context.Context, refreshToken string) (*TokenR
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("token 刷新请求失败: %w", err)
}
@@ -414,7 +414,7 @@ func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo
}
req.Header.Set("Authorization", "Bearer "+accessToken)
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("用户信息请求失败: %w", err)
}
@@ -465,7 +465,7 @@ func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadC
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
lastErr = fmt.Errorf("loadCodeAssist 请求失败: %w", err)
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
@@ -544,7 +544,7 @@ func (c *Client) OnboardUser(ctx context.Context, accessToken, tierID string) (s
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
lastErr = fmt.Errorf("onboardUser 请求失败: %w", err)
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
@@ -683,7 +683,7 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
resp, err := fetchClient.Do(req)
resp, err := servertiming.Do(fetchClient, req)
if err != nil {
lastErr = fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
@@ -842,7 +842,7 @@ func (c *Client) SetUserSettings(ctx context.Context, accessToken string) (*SetU
req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
req.Host = "daily-cloudcode-pa.googleapis.com"
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("setUserSettings 请求失败: %w", err)
}
@@ -885,7 +885,7 @@ func (c *Client) FetchUserInfo(ctx context.Context, accessToken, projectID strin
req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
req.Host = "daily-cloudcode-pa.googleapis.com"
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("fetchUserInfo 请求失败: %w", err)
}
@@ -718,7 +718,7 @@ func TestStreamingToolCallDoneWithoutDeltaEmitsArguments(t *testing.T) {
assert.Equal(t, "content_block_stop", events[1].Type)
}
func TestStreamingReadToolDropsEmptyPages(t *testing.T) {
func TestStreamingReadToolStreamsDeltas(t *testing.T) {
state := NewResponsesEventToAnthropicState()
ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
@@ -739,18 +739,17 @@ func TestStreamingReadToolDropsEmptyPages(t *testing.T) {
OutputIndex: 0,
Delta: `{"file_path":"/tmp/demo.py","limit":2000,"offset":0,"pages":""}`,
}, state)
assert.Len(t, events, 0)
require.Len(t, events, 1, "Read tool deltas must be streamed like any other tool")
assert.Equal(t, "content_block_delta", events[0].Type)
assert.Equal(t, "input_json_delta", events[0].Delta.Type)
events = ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
Type: "response.function_call_arguments.done",
OutputIndex: 0,
Arguments: `{"file_path":"/tmp/demo.py","limit":2000,"offset":0,"pages":""}`,
}, state)
require.Len(t, events, 2)
assert.Equal(t, "content_block_delta", events[0].Type)
assert.Equal(t, "input_json_delta", events[0].Delta.Type)
assert.JSONEq(t, `{"file_path":"/tmp/demo.py","limit":2000,"offset":0}`, events[0].Delta.PartialJSON)
assert.Equal(t, "content_block_stop", events[1].Type)
require.Len(t, events, 1, "after streaming deltas, .done should just close the block")
assert.Equal(t, "content_block_stop", events[0].Type)
}
func TestStreamingReasoning(t *testing.T) {
@@ -164,6 +164,8 @@ type AnthropicEventToResponsesState struct {
OutputTokens int
CacheReadInputTokens int
CacheCreationInputTokens int
StopReason string
}
// NewAnthropicEventToResponsesState returns an initialised stream state.
@@ -405,7 +407,6 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic
}
func anthToResHandleMessageDelta(evt *AnthropicStreamEvent, state *AnthropicEventToResponsesState) []ResponsesStreamEvent {
// Update usage
if evt.Usage != nil {
state.OutputTokens = evt.Usage.OutputTokens
if evt.Usage.InputTokens > 0 {
@@ -418,6 +419,9 @@ func anthToResHandleMessageDelta(evt *AnthropicStreamEvent, state *AnthropicEven
state.CacheCreationInputTokens = evt.Usage.CacheCreationInputTokens
}
}
if evt.Delta != nil && evt.Delta.StopReason != "" {
state.StopReason = evt.Delta.StopReason
}
return nil
}
@@ -428,15 +432,15 @@ func anthToResHandleMessageStop(state *AnthropicEventToResponsesState) []Respons
}
var events []ResponsesStreamEvent
// Close any open item
events = append(events, closeCurrentResponsesItem(state)...)
// Determine status
status := "completed"
var incompleteDetails *ResponsesIncompleteDetails
if state.StopReason == "max_tokens" {
status = "incomplete"
incompleteDetails = &ResponsesIncompleteDetails{Reason: "max_output_tokens"}
}
// Emit response.completed
events = append(events, makeResponsesCompletedEvent(state, status, incompleteDetails))
state.CompletedSent = true
return events
@@ -509,15 +513,20 @@ func makeResponsesCompletedEvent(
}
}
eventType := "response.completed"
if status == "incomplete" {
eventType = "response.incomplete"
}
return ResponsesStreamEvent{
Type: "response.completed",
Type: eventType,
SequenceNumber: seq,
Response: &ResponsesResponse{
ID: state.ResponseID,
Object: "response",
Model: state.Model,
Status: status,
Output: []ResponsesOutput{}, // Simplified; full output tracking would add complexity
Output: []ResponsesOutput{},
Usage: usage,
IncompleteDetails: incompleteDetails,
},
@@ -35,8 +35,12 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR
if req.Reasoning != nil {
out.ReasoningEffort = req.Reasoning.Effort
}
if len(req.Tools) > 0 {
tools, err := responsesToolsToChatTools(req.Tools)
effectiveTools, err := EffectiveResponsesTools(req)
if err != nil {
return nil, err
}
if len(effectiveTools) > 0 {
tools, err := responsesToolsToChatTools(effectiveTools)
if err != nil {
return nil, err
}
@@ -63,6 +67,44 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR
return out, nil
}
// EffectiveResponsesTools returns every client-executable tool declared by a
// Responses request. Newer Codex clients place their runtime tools in an
// input item shaped as {"type":"additional_tools","tools":[...]} instead of
// the top-level tools field. Chat-only upstreams must receive both forms.
func EffectiveResponsesTools(req *ResponsesRequest) ([]ResponsesTool, error) {
if req == nil {
return nil, nil
}
tools := append([]ResponsesTool(nil), req.Tools...)
inputRaw := bytesTrimSpace(req.Input)
if len(inputRaw) == 0 || string(inputRaw) == "null" || inputRaw[0] != '[' {
return tools, nil
}
var items []json.RawMessage
if err := json.Unmarshal(inputRaw, &items); err != nil {
return nil, fmt.Errorf("parse responses input for additional tools: %w", err)
}
for _, raw := range items {
raw = bytesTrimSpace(raw)
if len(raw) == 0 || raw[0] != '{' {
continue
}
var item struct {
Type string `json:"type"`
Tools []ResponsesTool `json:"tools"`
}
if err := json.Unmarshal(raw, &item); err != nil {
return nil, fmt.Errorf("parse responses additional tools item: %w", err)
}
if item.Type == "additional_tools" {
tools = append(tools, item.Tools...)
}
}
return tools, nil
}
// CustomToolNames 收集 Responses 请求中 custom/freeform 工具的名字。chat 桥回程时
// 需要据此把模型对这些工具的调用还原为 custom_tool_call 项(codex 只按该类型路由)。
func CustomToolNames(tools []ResponsesTool) map[string]bool {
@@ -34,6 +34,51 @@ func TestResponsesToChatCompletionsRequest_CustomToolBecomesFunctionTool(t *test
assert.Equal(t, "wait", out.Tools[1].Function.Name)
}
func TestResponsesToChatCompletionsRequest_AdditionalToolsItem(t *testing.T) {
req := &ResponsesRequest{
Model: "gpt-test",
Input: json.RawMessage(`[
{"type":"additional_tools","role":"developer","tools":[
{"type":"custom","name":"exec","description":"Run PowerShell","format":{"type":"text"}},
{"type":"function","name":"wait","parameters":{"type":"object","properties":{}}},
{"type":"namespace","name":"collaboration","tools":[
{"type":"function","name":"send_message","parameters":{"type":"object","properties":{}}}
]}
]},
{"type":"message","role":"user","content":[{"type":"input_text","text":"run Get-Location"}]}
]`),
ToolChoice: json.RawMessage(`"auto"`),
}
effective, err := EffectiveResponsesTools(req)
require.NoError(t, err)
require.Len(t, effective, 3)
assert.True(t, CustomToolNames(effective)["exec"])
assert.Equal(t, NamespacedToolName{Namespace: "collaboration", Name: "send_message"}, NamespaceToolNames(effective)["collaboration__send_message"])
out, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
require.Len(t, out.Tools, 3)
assert.Equal(t, "exec", out.Tools[0].Function.Name)
assert.Equal(t, "wait", out.Tools[1].Function.Name)
assert.Equal(t, "collaboration__send_message", out.Tools[2].Function.Name)
assert.JSONEq(t, `"auto"`, string(out.ToolChoice))
require.Len(t, out.Messages, 1, "additional_tools must not become a chat message")
assert.Equal(t, "user", out.Messages[0].Role)
}
func TestEffectiveResponsesTools_SkipsStringInputItems(t *testing.T) {
req := &ResponsesRequest{
Input: json.RawMessage(`["plain input",{"type":"additional_tools","tools":[{"type":"custom","name":"exec"}]}]`),
}
tools, err := EffectiveResponsesTools(req)
require.NoError(t, err)
require.Len(t, tools, 1)
assert.Equal(t, "exec", tools[0].Name)
}
func TestResponsesToChatCompletionsRequest_DropsToolChoiceWhenNoConvertibleTools(t *testing.T) {
req := &ResponsesRequest{
Model: "glm-5.2",
@@ -413,10 +413,6 @@ func resToAnthHandleFuncArgsDelta(evt *ResponsesStreamEvent, state *ResponsesEve
return nil
}
if state.CurrentBlockType == "tool_use" && state.CurrentToolName == "Read" {
state.CurrentToolArgs += evt.Delta
return nil
}
if state.CurrentBlockType == "tool_use" {
state.CurrentToolHadDelta = true
}
@@ -0,0 +1,84 @@
package apicompat
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestResToAnthFuncArgsDelta_ReadToolStreamsDeltas(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
state.CurrentBlockType = "tool_use"
state.CurrentToolName = "Read"
state.OutputIndexToBlockIdx = map[int]int{0: 0}
evt := &ResponsesStreamEvent{
Type: "response.function_call_arguments.delta",
OutputIndex: 0,
Delta: `{"file_path":"/tmp/test.go"}`,
}
events := ResponsesEventToAnthropicEvents(evt, state)
require.Len(t, events, 1, "Read tool delta must produce content_block_delta")
assert.Equal(t, "content_block_delta", events[0].Type)
assert.Equal(t, "input_json_delta", events[0].Delta.Type)
assert.Equal(t, `{"file_path":"/tmp/test.go"}`, events[0].Delta.PartialJSON)
assert.True(t, state.CurrentToolHadDelta, "Read deltas should set CurrentToolHadDelta")
}
func TestResToAnthFuncArgsDelta_ReadToolWithoutDone(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
state.ContentBlockIndex = 0
state.ContentBlockOpen = true
state.CurrentBlockType = "tool_use"
state.CurrentToolName = "Read"
state.OutputIndexToBlockIdx = map[int]int{0: 0}
delta := &ResponsesStreamEvent{
Type: "response.function_call_arguments.delta",
OutputIndex: 0,
Delta: `{"file_path":"/tmp/test.go"}`,
}
events := ResponsesEventToAnthropicEvents(delta, state)
require.Len(t, events, 1, "delta should be streamed")
completed := &ResponsesStreamEvent{
Type: "response.completed",
Response: &ResponsesResponse{
Status: "completed",
},
}
events = ResponsesEventToAnthropicEvents(completed, state)
hasStop := false
for _, e := range events {
if e.Type == "content_block_stop" {
hasStop = true
}
}
assert.True(t, hasStop, "block should be closed even without .done event")
}
func TestResToAnthFuncArgsDelta_NonReadToolUnchanged(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
state.CurrentBlockType = "tool_use"
state.CurrentToolName = "Write"
state.OutputIndexToBlockIdx = map[int]int{0: 0}
evt := &ResponsesStreamEvent{
Type: "response.function_call_arguments.delta",
OutputIndex: 0,
Delta: `{"file_path":"/tmp/out.txt","content":"hello"}`,
}
events := ResponsesEventToAnthropicEvents(evt, state)
require.Len(t, events, 1)
assert.Equal(t, "content_block_delta", events[0].Type)
assert.True(t, state.CurrentToolHadDelta)
}
@@ -89,8 +89,13 @@ func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatComp
func responsesStatusToChatFinishReason(status string, details *ResponsesIncompleteDetails, toolCalls []ChatToolCall) string {
switch status {
case "incomplete":
if details != nil && details.Reason == "max_output_tokens" {
return "length"
if details != nil {
switch details.Reason {
case "max_output_tokens":
return "length"
case "content_filter":
return "content_filter"
}
}
return "stop"
case "completed":
@@ -299,8 +304,13 @@ func resToChatHandleCompleted(evt *ResponsesStreamEvent, state *ResponsesEventTo
switch evt.Response.Status {
case "incomplete":
if evt.Response.IncompleteDetails != nil && evt.Response.IncompleteDetails.Reason == "max_output_tokens" {
finishReason = "length"
if evt.Response.IncompleteDetails != nil {
switch evt.Response.IncompleteDetails.Reason {
case "max_output_tokens":
finishReason = "length"
case "content_filter":
finishReason = "content_filter"
}
}
case "completed":
if state.SawToolCall {
@@ -0,0 +1,122 @@
package apicompat
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAnthropicStreamingMaxTokens_MapsToIncomplete(t *testing.T) {
state := NewAnthropicEventToResponsesState()
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_start",
Message: &AnthropicResponse{ID: "msg_test", Model: "claude-opus-4-6", Role: "assistant"},
}, state)
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_delta",
Delta: &AnthropicDelta{
StopReason: "max_tokens",
},
Usage: &AnthropicUsage{OutputTokens: 4096},
}, state)
require.Equal(t, "max_tokens", state.StopReason)
events := AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_stop",
}, state)
var completed *ResponsesStreamEvent
for i := range events {
if events[i].Type == "response.completed" || events[i].Type == "response.incomplete" {
completed = &events[i]
break
}
}
require.NotNil(t, completed, "should have terminal event")
assert.Equal(t, "response.incomplete", completed.Type)
require.NotNil(t, completed.Response)
assert.Equal(t, "incomplete", completed.Response.Status)
require.NotNil(t, completed.Response.IncompleteDetails)
assert.Equal(t, "max_output_tokens", completed.Response.IncompleteDetails.Reason)
}
func TestAnthropicStreamingEndTurn_MapsToCompleted(t *testing.T) {
state := NewAnthropicEventToResponsesState()
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_start",
Message: &AnthropicResponse{ID: "msg_test", Model: "claude-opus-4-6", Role: "assistant"},
}, state)
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_delta",
Delta: &AnthropicDelta{StopReason: "end_turn"},
Usage: &AnthropicUsage{OutputTokens: 100},
}, state)
events := AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_stop",
}, state)
var completed *ResponsesStreamEvent
for i := range events {
if events[i].Type == "response.completed" {
completed = &events[i]
break
}
}
require.NotNil(t, completed)
assert.Equal(t, "completed", completed.Response.Status)
assert.Nil(t, completed.Response.IncompleteDetails)
}
func TestResponsesToChatCompletions_ContentFilter(t *testing.T) {
resp := &ResponsesResponse{
ID: "resp_cf",
Status: "incomplete",
IncompleteDetails: &ResponsesIncompleteDetails{
Reason: "content_filter",
},
Output: []ResponsesOutput{{
Type: "message",
Content: []ResponsesContentPart{{Type: "output_text", Text: "partial"}},
}},
Usage: &ResponsesUsage{InputTokens: 10, OutputTokens: 5},
}
cc := ResponsesToChatCompletions(resp, "gpt-5.5")
require.Len(t, cc.Choices, 1)
assert.Equal(t, "content_filter", cc.Choices[0].FinishReason)
}
func TestResponsesToChatCompletionsStreaming_ContentFilter(t *testing.T) {
state := NewResponsesEventToChatState()
state.ID = "resp_cf"
state.Model = "gpt-5.5"
state.SentRole = true
events := ResponsesEventToChatChunks(&ResponsesStreamEvent{
Type: "response.completed",
Response: &ResponsesResponse{
ID: "resp_cf",
Status: "incomplete",
IncompleteDetails: &ResponsesIncompleteDetails{
Reason: "content_filter",
},
},
}, state)
hasContentFilter := false
for _, chunk := range events {
for _, choice := range chunk.Choices {
if choice.FinishReason != nil && *choice.FinishReason == "content_filter" {
hasContentFilter = true
}
}
}
assert.True(t, hasContentFilter, "streaming content_filter should map to finish_reason content_filter")
}
+2
View File
@@ -25,6 +25,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
)
@@ -92,6 +93,7 @@ func buildClient(opts Options) (*http.Client, error) {
if opts.ValidateResolvedIP && !opts.AllowPrivateHosts {
rt = newValidatedTransport(transport)
}
rt = servertiming.WrapRoundTripper(rt)
return &http.Client{
Transport: rt,
Timeout: opts.Timeout,
@@ -38,7 +38,7 @@ func (p PaginationParams) Offset() int {
if p.Page < 1 {
p.Page = 1
}
return (p.Page - 1) * p.PageSize
return (p.Page - 1) * p.Limit()
}
// Limit 获取限制数
@@ -69,3 +69,30 @@ func TestPaginationParamsLimit(t *testing.T) {
})
}
}
func TestPaginationParamsOffsetUsesNormalizedLimit(t *testing.T) {
t.Parallel()
tests := []struct {
name string
page int
pageSize int
want int
}{
{name: "invalid page uses first page", page: 0, pageSize: 50, want: 0},
{name: "zero page size uses default", page: 2, pageSize: 0, want: 20},
{name: "negative page size uses default", page: 2, pageSize: -1, want: 20},
{name: "normal values", page: 3, pageSize: 50, want: 100},
{name: "page size beyond max is clamped", page: 2, pageSize: 1500, want: 1000},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
params := PaginationParams{Page: tt.page, PageSize: tt.pageSize}
if got := params.Offset(); got != tt.want {
t.Fatalf("Offset() for Page=%d, PageSize=%d = %d, want %d", tt.page, tt.pageSize, got, tt.want)
}
})
}
}
@@ -0,0 +1,348 @@
package servertiming
import (
"context"
"fmt"
"sort"
"strconv"
"strings"
"sync"
"time"
)
const (
HeaderName = "Server-Timing"
AdminUIHeader = "X-Admin-UI-Request"
MetricDatabase = "db"
MetricRedis = "redis"
dependencyPrefix = "dep_"
maxMetricNameLength = 48
maxIntervals = 2048
maxHeaderLength = 4096
)
type contextKey struct{}
type interval struct {
start time.Time
end time.Time
}
type metric struct {
count int64
intervals []interval
}
// Collector stores request-scoped timing samples. It is safe for concurrent use.
type Collector struct {
startedAt time.Time
mu sync.Mutex
metrics map[string]*metric
cacheStatus string
}
// New creates a collector whose total duration starts at startedAt.
func New(startedAt time.Time) *Collector {
if startedAt.IsZero() {
startedAt = time.Now()
}
return &Collector{
startedAt: startedAt,
metrics: make(map[string]*metric),
}
}
// WithCollector attaches a collector to a context.
func WithCollector(ctx context.Context, collector *Collector) context.Context {
if ctx == nil {
ctx = context.Background()
}
if collector == nil {
return ctx
}
return context.WithValue(ctx, contextKey{}, collector)
}
// FromContext returns the request timing collector, when one is active.
func FromContext(ctx context.Context) (*Collector, bool) {
if ctx == nil {
return nil, false
}
collector, ok := ctx.Value(contextKey{}).(*Collector)
return collector, ok && collector != nil
}
// Active reports whether timing collection is enabled for this request.
func Active(ctx context.Context) bool {
_, ok := FromContext(ctx)
return ok
}
// Record adds a completed interval and operation count to a metric.
func Record(ctx context.Context, name string, startedAt, endedAt time.Time, count int) {
collector, ok := FromContext(ctx)
if !ok {
return
}
collector.Record(name, startedAt, endedAt, count)
}
// RecordInterval adds timing without incrementing the operation count. It is
// useful when one logical operation has multiple blocking driver calls.
func RecordInterval(ctx context.Context, name string, startedAt, endedAt time.Time) {
collector, ok := FromContext(ctx)
if !ok {
return
}
collector.record(name, startedAt, endedAt, 0)
}
// Record adds a completed interval directly to the collector.
func (c *Collector) Record(name string, startedAt, endedAt time.Time, count int) {
if count <= 0 {
count = 1
}
c.record(name, startedAt, endedAt, count)
}
func (c *Collector) record(name string, startedAt, endedAt time.Time, count int) {
name = normalizeMetricName(name)
if c == nil || name == "" || startedAt.IsZero() || endedAt.Before(startedAt) {
return
}
if count < 0 {
count = 0
}
c.mu.Lock()
m := c.metrics[name]
if m == nil {
m = &metric{}
c.metrics[name] = m
}
m.count += int64(count)
if len(m.intervals) < maxIntervals {
m.intervals = append(m.intervals, interval{start: startedAt, end: endedAt})
}
c.mu.Unlock()
}
// Observe starts a metric span and returns an idempotent completion function.
func Observe(ctx context.Context, name string) func() {
collector, ok := FromContext(ctx)
name = normalizeMetricName(name)
if !ok || name == "" {
return func() {}
}
startedAt := time.Now()
var once sync.Once
return func() {
once.Do(func() {
collector.Record(name, startedAt, time.Now(), 1)
})
}
}
// ObserveDependency starts a named external dependency span.
func ObserveDependency(ctx context.Context, module string) func() {
return Observe(ctx, dependencyMetricName(module))
}
// RecordDependency records a completed external dependency interval.
func RecordDependency(ctx context.Context, module string, startedAt, endedAt time.Time) {
Record(ctx, dependencyMetricName(module), startedAt, endedAt, 1)
}
// SetCacheStatus records the response-cache outcome for the request.
func SetCacheStatus(ctx context.Context, status string) {
collector, ok := FromContext(ctx)
if !ok {
return
}
status = normalizeCacheStatus(status)
if status == "" {
return
}
collector.mu.Lock()
collector.cacheStatus = status
collector.mu.Unlock()
}
// HeaderValue renders a bounded, deterministic Server-Timing header.
func HeaderValue(ctx context.Context, endedAt time.Time, cacheStatus string) string {
collector, ok := FromContext(ctx)
if !ok {
return ""
}
return collector.HeaderValue(endedAt, cacheStatus)
}
// HeaderValue renders a bounded, deterministic Server-Timing header.
func (c *Collector) HeaderValue(endedAt time.Time, cacheStatus string) string {
if c == nil {
return ""
}
if endedAt.IsZero() {
endedAt = time.Now()
}
if endedAt.Before(c.startedAt) {
endedAt = c.startedAt
}
c.mu.Lock()
metrics := make(map[string]metric, len(c.metrics))
allIntervals := make([]interval, 0)
dependencyIntervals := make([]interval, 0)
var dependencyCount int64
for name, source := range c.metrics {
copied := metric{count: source.count, intervals: append([]interval(nil), source.intervals...)}
metrics[name] = copied
allIntervals = append(allIntervals, copied.intervals...)
if strings.HasPrefix(name, dependencyPrefix) {
dependencyIntervals = append(dependencyIntervals, copied.intervals...)
dependencyCount += copied.count
}
}
storedCacheStatus := c.cacheStatus
c.mu.Unlock()
total := endedAt.Sub(c.startedAt)
blocked := unionDuration(allIntervals, c.startedAt, endedAt)
app := total - blocked
if app < 0 {
app = 0
}
cacheStatus = normalizeCacheStatus(cacheStatus)
if cacheStatus == "" {
cacheStatus = normalizeCacheStatus(storedCacheStatus)
}
if cacheStatus == "" {
cacheStatus = "bypass"
}
database := metrics[MetricDatabase]
redisMetric := metrics[MetricRedis]
parts := []string{
"total;dur=" + formatDuration(total),
"app;dur=" + formatDuration(app),
fmt.Sprintf("db;dur=%s;desc=\"queries=%d\"", formatDuration(unionDuration(database.intervals, c.startedAt, endedAt)), database.count),
fmt.Sprintf("redis;dur=%s;desc=\"commands=%d\"", formatDuration(unionDuration(redisMetric.intervals, c.startedAt, endedAt)), redisMetric.count),
"cache;desc=\"" + cacheStatus + "\"",
fmt.Sprintf("deps;dur=%s;desc=\"calls=%d\"", formatDuration(unionDuration(dependencyIntervals, c.startedAt, endedAt)), dependencyCount),
}
dependencyNames := make([]string, 0)
for name := range metrics {
if strings.HasPrefix(name, dependencyPrefix) {
dependencyNames = append(dependencyNames, name)
}
}
sort.Strings(dependencyNames)
for _, name := range dependencyNames {
m := metrics[name]
part := fmt.Sprintf("%s;dur=%s;desc=\"calls=%d\"", name, formatDuration(unionDuration(m.intervals, c.startedAt, endedAt)), m.count)
candidate := strings.Join(append(parts, part), ", ")
if len(candidate) > maxHeaderLength {
break
}
parts = append(parts, part)
}
return strings.Join(parts, ", ")
}
func dependencyMetricName(module string) string {
module = normalizeMetricName(module)
module = strings.TrimPrefix(module, dependencyPrefix)
if module == "" {
module = "http"
}
return dependencyPrefix + module
}
func normalizeMetricName(name string) string {
name = strings.ToLower(strings.TrimSpace(name))
if name == "" {
return ""
}
var b strings.Builder
b.Grow(min(len(name), maxMetricNameLength))
for _, r := range name {
if b.Len() >= maxMetricNameLength {
break
}
switch {
case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
_, _ = b.WriteRune(r)
case r == '_' || r == '-':
_ = b.WriteByte('_')
}
}
return strings.Trim(b.String(), "_")
}
func normalizeCacheStatus(status string) string {
switch strings.ToLower(strings.TrimSpace(status)) {
case "hit":
return "hit"
case "miss":
return "miss"
case "bypass":
return "bypass"
default:
return ""
}
}
func unionDuration(intervals []interval, lowerBound, upperBound time.Time) time.Duration {
if len(intervals) == 0 || !upperBound.After(lowerBound) {
return 0
}
normalized := make([]interval, 0, len(intervals))
for _, item := range intervals {
start := item.start
end := item.end
if start.Before(lowerBound) {
start = lowerBound
}
if end.After(upperBound) {
end = upperBound
}
if end.After(start) {
normalized = append(normalized, interval{start: start, end: end})
}
}
if len(normalized) == 0 {
return 0
}
sort.Slice(normalized, func(i, j int) bool {
return normalized[i].start.Before(normalized[j].start)
})
currentStart := normalized[0].start
currentEnd := normalized[0].end
var total time.Duration
for _, item := range normalized[1:] {
if !item.start.After(currentEnd) {
if item.end.After(currentEnd) {
currentEnd = item.end
}
continue
}
total += currentEnd.Sub(currentStart)
currentStart = item.start
currentEnd = item.end
}
total += currentEnd.Sub(currentStart)
return total
}
func formatDuration(value time.Duration) string {
if value < 0 {
value = 0
}
return strconv.FormatFloat(float64(value)/float64(time.Millisecond), 'f', 1, 64)
}
@@ -0,0 +1,129 @@
package servertiming
import (
"context"
"fmt"
"strings"
"sync"
"testing"
"time"
)
func TestCollectorHeaderValueAggregatesIntervals(t *testing.T) {
startedAt := time.Unix(100, 0)
collector := New(startedAt)
collector.Record(MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(40*time.Millisecond), 2)
collector.Record(MetricRedis, startedAt.Add(30*time.Millisecond), startedAt.Add(50*time.Millisecond), 3)
collector.Record(dependencyMetricName("openai"), startedAt.Add(70*time.Millisecond), startedAt.Add(100*time.Millisecond), 1)
collector.Record(dependencyMetricName("github"), startedAt.Add(60*time.Millisecond), startedAt.Add(90*time.Millisecond), 1)
got := collector.HeaderValue(startedAt.Add(120*time.Millisecond), "miss")
want := `total;dur=120.0, app;dur=40.0, db;dur=30.0;desc="queries=2", redis;dur=20.0;desc="commands=3", cache;desc="miss", deps;dur=40.0;desc="calls=2", dep_github;dur=30.0;desc="calls=1", dep_openai;dur=30.0;desc="calls=1"`
if got != want {
t.Fatalf("HeaderValue() = %q, want %q", got, want)
}
}
func TestRecordIntervalDoesNotIncrementCount(t *testing.T) {
startedAt := time.Unix(200, 0)
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
Record(ctx, MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(20*time.Millisecond), 1)
RecordInterval(ctx, MetricDatabase, startedAt.Add(30*time.Millisecond), startedAt.Add(40*time.Millisecond))
header := HeaderValue(ctx, startedAt.Add(100*time.Millisecond), "hit")
if !strings.Contains(header, `db;dur=20.0;desc="queries=1"`) {
t.Fatalf("header %q does not contain one query with both blocking intervals", header)
}
if !strings.Contains(header, "app;dur=80.0") {
t.Fatalf("header %q does not subtract the interval union from app time", header)
}
}
func TestCollectorCacheStatusFallback(t *testing.T) {
startedAt := time.Unix(300, 0)
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
SetCacheStatus(ctx, " HIT ")
if got := HeaderValue(ctx, startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="hit"`) {
t.Fatalf("HeaderValue() = %q, want stored cache hit", got)
}
other := New(startedAt)
if got := other.HeaderValue(startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="bypass"`) {
t.Fatalf("HeaderValue() = %q, want cache bypass", got)
}
}
func TestCollectorSanitizesDependencyMetric(t *testing.T) {
startedAt := time.Unix(400, 0)
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
RecordDependency(ctx, "GitHub API\r\nInjected;dur=999", startedAt, startedAt.Add(time.Millisecond))
header := HeaderValue(ctx, startedAt.Add(2*time.Millisecond), "bypass")
if strings.ContainsAny(header, "\r\n") || strings.Contains(header, ";dur=999") {
t.Fatalf("unsafe metric content reached header: %q", header)
}
if !strings.Contains(header, "dep_githubapiinjecteddur999;dur=1.0") {
t.Fatalf("sanitized dependency metric missing from header: %q", header)
}
}
func TestCollectorBoundsHeaderLength(t *testing.T) {
startedAt := time.Unix(500, 0)
collector := New(startedAt)
for i := 0; i < 300; i++ {
collector.Record(
dependencyMetricName(fmt.Sprintf("module_%03d_with_a_deliberately_long_name", i)),
startedAt,
startedAt.Add(time.Millisecond),
1,
)
}
header := collector.HeaderValue(startedAt.Add(2*time.Millisecond), "bypass")
if len(header) > maxHeaderLength {
t.Fatalf("header length = %d, want <= %d", len(header), maxHeaderLength)
}
if !strings.Contains(header, "total;dur=2.0") || !strings.Contains(header, "deps;dur=1.0") {
t.Fatalf("bounded header lost fixed metrics: %q", header)
}
}
func TestCollectorConcurrentRecording(t *testing.T) {
startedAt := time.Now()
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
const workers = 25
const recordsPerWorker = 100
var wg sync.WaitGroup
wg.Add(workers)
for i := 0; i < workers; i++ {
go func() {
defer wg.Done()
for j := 0; j < recordsPerWorker; j++ {
Record(ctx, MetricDatabase, startedAt, startedAt.Add(time.Microsecond), 1)
}
}()
}
wg.Wait()
header := HeaderValue(ctx, startedAt.Add(time.Millisecond), "bypass")
want := fmt.Sprintf(`queries=%d`, workers*recordsPerWorker)
if !strings.Contains(header, want) {
t.Fatalf("header %q does not contain %q", header, want)
}
}
func TestContextHelpersHandleMissingCollector(t *testing.T) {
if Active(context.Background()) {
t.Fatal("context without collector reported active")
}
if got := HeaderValue(context.Background(), time.Now(), "hit"); got != "" {
t.Fatalf("HeaderValue() = %q without collector, want empty", got)
}
}
+104
View File
@@ -0,0 +1,104 @@
package servertiming
import (
"context"
"net/http"
"strings"
"time"
)
type dependencyModuleKey struct{}
type timingRoundTripper struct {
base http.RoundTripper
}
// WithDependencyModule overrides the safe module name used for an outbound call.
func WithDependencyModule(ctx context.Context, module string) context.Context {
if ctx == nil {
ctx = context.Background()
}
module = strings.TrimPrefix(normalizeMetricName(module), dependencyPrefix)
if module == "" {
return ctx
}
return context.WithValue(ctx, dependencyModuleKey{}, module)
}
// WrapRoundTripper records outbound response-header latency for active requests.
func WrapRoundTripper(base http.RoundTripper) http.RoundTripper {
if base == nil {
base = http.DefaultTransport
}
if _, ok := base.(*timingRoundTripper); ok {
return base
}
return &timingRoundTripper{base: base}
}
// InstrumentClient returns a shallow client copy with an instrumented transport.
func InstrumentClient(client *http.Client) *http.Client {
if client == nil {
client = &http.Client{}
}
copyClient := *client
copyClient.Transport = WrapRoundTripper(copyClient.Transport)
return &copyClient
}
// Do records response-header latency without changing the client's transport
// type. Use it for clients whose callers inspect or configure *http.Transport.
func Do(client *http.Client, req *http.Request) (*http.Response, error) {
if client == nil {
client = http.DefaultClient
}
if req == nil || !Active(req.Context()) {
return client.Do(req)
}
startedAt := time.Now()
response, err := client.Do(req)
RecordDependency(req.Context(), dependencyModule(req), startedAt, time.Now())
return response, err
}
func (t *timingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
if req == nil || !Active(req.Context()) {
return t.base.RoundTrip(req)
}
startedAt := time.Now()
response, err := t.base.RoundTrip(req)
RecordDependency(req.Context(), dependencyModule(req), startedAt, time.Now())
return response, err
}
func dependencyModule(req *http.Request) string {
if req != nil {
if module, ok := req.Context().Value(dependencyModuleKey{}).(string); ok && module != "" {
return module
}
}
if req == nil || req.URL == nil {
return "http"
}
host := strings.ToLower(req.URL.Hostname())
switch {
case strings.Contains(host, "github"):
return "github"
case strings.Contains(host, "openai"):
return "openai"
case strings.Contains(host, "anthropic"):
return "anthropic"
case strings.Contains(host, "generativelanguage") || strings.Contains(host, "gemini"):
return "gemini"
case strings.Contains(host, "cloudcode") || strings.Contains(host, "antigravity"):
return "antigravity"
case strings.Contains(host, "googleapis") || strings.Contains(host, "google"):
return "google"
case strings.Contains(host, "amazonaws") || strings.Contains(host, "cloudflarestorage") || strings.Contains(host, "s3"):
return "s3"
case strings.Contains(host, "stripe") || strings.Contains(host, "airwallex") || strings.Contains(host, "alipay") || strings.Contains(host, "wechatpay") || strings.Contains(host, "paypal"):
return "payment"
default:
return "http"
}
}
@@ -0,0 +1,168 @@
package servertiming
import (
"context"
"io"
"net/http"
"strings"
"testing"
"time"
)
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return f(req)
}
type trackingBody struct {
read bool
}
func (b *trackingBody) Read(_ []byte) (int, error) {
b.read = true
return 0, io.EOF
}
func (b *trackingBody) Close() error { return nil }
func TestWrapRoundTripperRecordsResponseHeaderLatency(t *testing.T) {
startedAt := time.Now()
collector := New(startedAt)
body := &trackingBody{}
baseCalled := false
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
baseCalled = true
return &http.Response{
StatusCode: http.StatusOK,
Body: body,
Header: make(http.Header),
Request: req,
}, nil
})
req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.github.com/repos/example/project", nil)
if err != nil {
t.Fatal(err)
}
resp, err := WrapRoundTripper(base).RoundTrip(req)
if err != nil {
t.Fatal(err)
}
defer func() { _ = resp.Body.Close() }()
if !baseCalled {
t.Fatal("base RoundTripper was not called")
}
if body.read {
t.Fatal("RoundTripper instrumentation read the response body; timing must stop at response headers")
}
header := collector.HeaderValue(time.Now(), "bypass")
if !strings.Contains(header, `dep_github;dur=`) || !strings.Contains(header, `deps;dur=`) {
t.Fatalf("dependency metrics missing from header: %q", header)
}
}
func TestWrapRoundTripperUsesContextModuleOverride(t *testing.T) {
collector := New(time.Now())
ctx := WithDependencyModule(WithCollector(context.Background(), collector), "data-managementd")
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://private.example.test/path", nil)
if err != nil {
t.Fatal(err)
}
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
})
if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil {
t.Fatal(err)
}
header := collector.HeaderValue(time.Now(), "bypass")
if !strings.Contains(header, "dep_data_managementd") {
t.Fatalf("module override missing from header: %q", header)
}
if strings.Contains(header, "private.example") {
t.Fatalf("raw host leaked into header: %q", header)
}
}
func TestWrapRoundTripperSkipsInactiveContext(t *testing.T) {
called := false
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
called = true
return &http.Response{StatusCode: http.StatusNoContent, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
})
req, err := http.NewRequest(http.MethodGet, "https://api.openai.com/v1/models", nil)
if err != nil {
t.Fatal(err)
}
if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil {
t.Fatal(err)
}
if !called {
t.Fatal("inactive request did not reach base RoundTripper")
}
}
func TestDoRecordsWithoutChangingTransportType(t *testing.T) {
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
})
client := &http.Client{Transport: base}
collector := New(time.Now())
req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.openai.com/v1/models", nil)
if err != nil {
t.Fatal(err)
}
if _, err := Do(client, req); err != nil {
t.Fatal(err)
}
if _, ok := client.Transport.(roundTripFunc); !ok {
t.Fatalf("Do changed client transport type to %T", client.Transport)
}
if header := collector.HeaderValue(time.Now(), "bypass"); !strings.Contains(header, "dep_openai;dur=") {
t.Fatalf("dependency metric missing from header: %q", header)
}
}
func TestDependencyModuleClassification(t *testing.T) {
tests := map[string]string{
"https://api.github.com/repos/a/b": "github",
"https://api.openai.com/v1/models": "openai",
"https://api.anthropic.com/v1/messages": "anthropic",
"https://generativelanguage.googleapis.com/v1/models": "gemini",
"https://cloudcode-pa.googleapis.com/v1internal": "antigravity",
"https://storage.googleapis.com/bucket/object": "google",
"https://bucket.s3.amazonaws.com/object": "s3",
"https://api.stripe.com/v1/refunds": "payment",
"https://dependency.example.test/path": "http",
}
for rawURL, want := range tests {
req, err := http.NewRequest(http.MethodGet, rawURL, nil)
if err != nil {
t.Fatalf("NewRequest(%q): %v", rawURL, err)
}
if got := dependencyModule(req); got != want {
t.Errorf("dependencyModule(%q) = %q, want %q", rawURL, got, want)
}
}
}
func TestClientInstrumentationDoesNotMutateOriginal(t *testing.T) {
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
})
original := &http.Client{Transport: base, Timeout: time.Second}
instrumented := InstrumentClient(original)
if instrumented == original {
t.Fatal("InstrumentClient returned the original client")
}
if _, ok := original.Transport.(roundTripFunc); !ok {
t.Fatalf("InstrumentClient mutated the original transport to %T", original.Transport)
}
if instrumented.Timeout != original.Timeout {
t.Fatal("InstrumentClient did not preserve client settings")
}
if WrapRoundTripper(instrumented.Transport) != instrumented.Transport {
t.Fatal("WrapRoundTripper wrapped an already instrumented transport twice")
}
}
+372
View File
@@ -0,0 +1,372 @@
package xai
import (
"encoding/json"
"fmt"
"math"
"net/http"
"strconv"
"strings"
"time"
)
const (
// CLI client identity required by cli-chat-proxy billing endpoints.
CLITokenAuthHeader = "x-xai-token-auth"
CLITokenAuthValue = "xai-grok-cli"
CLIClientVersionHeader = "x-grok-client-version"
// Keep in sync with https://x.ai/cli/stable.
CLIClientVersion = "0.2.93"
CLIUserAgent = "grok-pager/" + CLIClientVersion + " grok-shell/" + CLIClientVersion + " (macos; aarch64)"
BillingWeeklyPath = "/billing?format=credits"
BillingMonthlyPath = "/billing"
SuperGrokLimitCents = 15_000 // $150.00
SuperGrokHeavyLimitCents = 150_000 // $1,500.00
)
// BillingPeriod describes the current weekly/monthly window.
type BillingPeriod struct {
Type string `json:"type,omitempty"`
Start string `json:"start,omitempty"`
End string `json:"end,omitempty"`
}
// BillingProductUsage is per-product usage inside the weekly credits window.
type BillingProductUsage struct {
Product string `json:"product,omitempty"`
UsagePercent *float64 `json:"usagePercent,omitempty"`
}
// BillingConfig is the nested config object from /v1/billing responses.
type BillingConfig struct {
CurrentPeriod *BillingPeriod `json:"currentPeriod,omitempty"`
CreditUsagePercent *float64 `json:"creditUsagePercent,omitempty"`
ProductUsage []BillingProductUsage `json:"productUsage,omitempty"`
MonthlyLimit json.RawMessage `json:"monthlyLimit,omitempty"`
Used json.RawMessage `json:"used,omitempty"`
BillingPeriodStart string `json:"billingPeriodStart,omitempty"`
BillingPeriodEnd string `json:"billingPeriodEnd,omitempty"`
}
// BillingPayload is the top-level body from /v1/billing.
type BillingPayload struct {
Config *BillingConfig `json:"config,omitempty"`
}
// BillingProductSummary is a normalized product usage row for UI.
type BillingProductSummary struct {
Product string `json:"product"`
UsagePercent *float64 `json:"usage_percent,omitempty"`
}
// BillingSummary is the merged weekly + monthly billing view.
type BillingSummary struct {
PeriodType string `json:"period_type,omitempty"` // weekly | monthly | unknown
UsagePercent *float64 `json:"usage_percent,omitempty"`
PeriodStart string `json:"period_start,omitempty"`
PeriodEnd string `json:"period_end,omitempty"`
ProductUsage []BillingProductSummary `json:"product_usage,omitempty"`
MonthlyLimitCents *float64 `json:"monthly_limit_cents,omitempty"`
UsedCents *float64 `json:"used_cents,omitempty"`
IncludedUsedCents *float64 `json:"included_used_cents,omitempty"`
BillingPeriodStart string `json:"billing_period_start,omitempty"`
BillingPeriodEnd string `json:"billing_period_end,omitempty"`
UsedPercent *float64 `json:"used_percent,omitempty"`
Plan string `json:"plan,omitempty"` // SuperGrok | SuperGrok Heavy | ""
StatusCode int `json:"status_code,omitempty"`
Source string `json:"source,omitempty"`
FetchedAt string `json:"fetched_at,omitempty"`
UpdatedAt string `json:"updated_at,omitempty"`
WeeklyUpdatedAt string `json:"weekly_updated_at,omitempty"`
MonthlyUpdatedAt string `json:"monthly_updated_at,omitempty"`
Partial bool `json:"partial,omitempty"`
FailedWindows []string `json:"failed_windows,omitempty"`
}
// BuildBillingURL builds weekly or monthly billing URL against the CLI chat proxy.
func BuildBillingURL(formatCredits bool) string {
base := strings.TrimRight(DefaultCLIBaseURL, "/")
if formatCredits {
return base + BillingWeeklyPath
}
return base + BillingMonthlyPath
}
// ApplyCLIBillingHeaders sets Authorization + CLI identity headers for billing GETs.
func ApplyCLIBillingHeaders(req *http.Request, accessToken string) {
if req == nil {
return
}
token := strings.TrimSpace(accessToken)
if token != "" {
req.Header.Set("Authorization", "Bearer "+token)
}
req.Header.Set("Accept", "application/json")
req.Header.Set("Content-Type", "application/json")
req.Header.Set(CLITokenAuthHeader, CLITokenAuthValue)
req.Header.Set(CLIClientVersionHeader, CLIClientVersion)
req.Header.Set("User-Agent", CLIUserAgent)
}
// ParseBillingPayload unmarshals a billing API response body.
func ParseBillingPayload(body []byte) (*BillingPayload, error) {
if len(body) == 0 {
return nil, fmt.Errorf("empty billing body")
}
var payload BillingPayload
if err := json.Unmarshal(body, &payload); err != nil {
return nil, err
}
return &payload, nil
}
// BuildBillingSummary normalizes a billing config into a UI-friendly summary.
func BuildBillingSummary(config *BillingConfig) *BillingSummary {
if config == nil {
return nil
}
summary := &BillingSummary{}
period := config.CurrentPeriod
periodType := resolvePeriodType(period)
creditUsage := cloneFloat(config.CreditUsagePercent)
periodStart := ""
periodEnd := ""
if period != nil {
periodStart = strings.TrimSpace(period.Start)
periodEnd = strings.TrimSpace(period.End)
}
if periodStart == "" {
periodStart = strings.TrimSpace(config.BillingPeriodStart)
}
if periodEnd == "" {
periodEnd = strings.TrimSpace(config.BillingPeriodEnd)
}
products := make([]BillingProductSummary, 0, len(config.ProductUsage))
for _, item := range config.ProductUsage {
product := strings.TrimSpace(item.Product)
if product == "" {
continue
}
products = append(products, BillingProductSummary{
Product: product,
UsagePercent: cloneFloat(item.UsagePercent),
})
}
monthlyLimit := parseCentValue(config.MonthlyLimit)
used := parseCentValue(config.Used)
billingStart := strings.TrimSpace(config.BillingPeriodStart)
billingEnd := strings.TrimSpace(config.BillingPeriodEnd)
var includedUsed *float64
if used != nil {
if monthlyLimit != nil && *monthlyLimit > 0 {
v := math.Min(*used, *monthlyLimit)
includedUsed = &v
} else {
includedUsed = cloneFloat(used)
}
}
var usedPercent *float64
if monthlyLimit != nil && *monthlyLimit > 0 && includedUsed != nil {
v := (*includedUsed / *monthlyLimit) * 100
usedPercent = &v
}
hasWeekly := creditUsage != nil || periodType == "weekly" || len(products) > 0
hasMonthly := monthlyLimit != nil || used != nil || (!hasWeekly && billingEnd != "")
if !hasWeekly && !hasMonthly {
return nil
}
if hasWeekly {
if periodType == "unknown" {
periodType = "weekly"
}
summary.PeriodType = periodType
summary.UsagePercent = creditUsage
summary.PeriodStart = periodStart
summary.PeriodEnd = periodEnd
} else {
// Monthly-only: do not put monthly % into UsagePercent (weekly bar field).
// Frontend weekly bar only renders when PeriodType == weekly.
summary.PeriodType = "monthly"
summary.PeriodStart = billingStart
summary.PeriodEnd = billingEnd
}
summary.ProductUsage = products
summary.MonthlyLimitCents = monthlyLimit
summary.UsedCents = used
summary.IncludedUsedCents = includedUsed
if hasMonthly {
summary.BillingPeriodStart = billingStart
summary.BillingPeriodEnd = billingEnd
}
summary.UsedPercent = usedPercent
summary.Plan = resolvePlan(monthlyLimit)
return summary
}
// MergeBillingProbeResult updates successful billing domains while retaining
// the previous value for any domain that could not be refreshed.
func MergeBillingProbeResult(previous, weekly, monthly *BillingSummary, weeklyOK, monthlyOK bool) *BillingSummary {
var out BillingSummary
if previous != nil {
out = *previous
previousUpdatedAt := previous.UpdatedAt
if previousUpdatedAt == "" {
previousUpdatedAt = previous.FetchedAt
}
if out.WeeklyUpdatedAt == "" && (out.UsagePercent != nil || len(out.ProductUsage) > 0) {
out.WeeklyUpdatedAt = previousUpdatedAt
}
if out.MonthlyUpdatedAt == "" && (out.MonthlyLimitCents != nil || out.UsedPercent != nil) {
out.MonthlyUpdatedAt = previousUpdatedAt
}
}
now := time.Now().UTC().Format(time.RFC3339)
if weeklyOK && weekly != nil {
out.PeriodType = weekly.PeriodType
out.UsagePercent = weekly.UsagePercent
out.PeriodStart = weekly.PeriodStart
out.PeriodEnd = weekly.PeriodEnd
out.ProductUsage = weekly.ProductUsage
out.WeeklyUpdatedAt = now
}
if monthlyOK && monthly != nil {
if out.PeriodType == "" {
out.PeriodType = "monthly"
}
out.MonthlyLimitCents = monthly.MonthlyLimitCents
out.UsedCents = monthly.UsedCents
out.IncludedUsedCents = monthly.IncludedUsedCents
out.BillingPeriodStart = monthly.BillingPeriodStart
out.BillingPeriodEnd = monthly.BillingPeriodEnd
out.UsedPercent = monthly.UsedPercent
out.Plan = monthly.Plan
out.MonthlyUpdatedAt = now
}
out.Partial = !weeklyOK || !monthlyOK
out.FailedWindows = nil
if !weeklyOK {
out.FailedWindows = append(out.FailedWindows, "weekly")
}
if !monthlyOK {
out.FailedWindows = append(out.FailedWindows, "monthly")
}
if !weeklyOK && !monthlyOK && previous == nil {
return nil
}
return &out
}
// StampBillingSummary sets fetch metadata.
func StampBillingSummary(summary *BillingSummary, statusCode int, source string) *BillingSummary {
if summary == nil {
return nil
}
now := time.Now().UTC().Format(time.RFC3339)
summary.StatusCode = statusCode
summary.Source = source
summary.FetchedAt = now
summary.UpdatedAt = now
return summary
}
func resolvePeriodType(period *BillingPeriod) string {
if period == nil {
return "unknown"
}
raw := strings.ToLower(strings.TrimSpace(period.Type))
if strings.Contains(raw, "weekly") {
return "weekly"
}
if strings.Contains(raw, "monthly") {
return "monthly"
}
return "unknown"
}
func resolvePlan(monthlyLimitCents *float64) string {
if monthlyLimitCents == nil {
return ""
}
// Allow small float noise.
limit := math.Round(*monthlyLimitCents)
switch limit {
case SuperGrokLimitCents:
return "SuperGrok"
case SuperGrokHeavyLimitCents:
return "SuperGrok Heavy"
default:
return ""
}
}
func parseCentValue(raw json.RawMessage) *float64 {
if len(raw) == 0 || string(raw) == "null" {
return nil
}
// Object form: {"val": 123}
var obj struct {
Val any `json:"val"`
}
if err := json.Unmarshal(raw, &obj); err == nil && obj.Val != nil {
return anyToFloat(obj.Val)
}
// Bare number / string
var n any
if err := json.Unmarshal(raw, &n); err != nil {
return nil
}
return anyToFloat(n)
}
func anyToFloat(v any) *float64 {
switch n := v.(type) {
case float64:
return &n
case float32:
f := float64(n)
return &f
case int:
f := float64(n)
return &f
case int64:
f := float64(n)
return &f
case json.Number:
f, err := n.Float64()
if err != nil {
return nil
}
return &f
case string:
s := strings.TrimSpace(n)
if s == "" {
return nil
}
f, err := strconv.ParseFloat(s, 64)
if err != nil {
return nil
}
return &f
default:
return nil
}
}
func cloneFloat(v *float64) *float64 {
if v == nil {
return nil
}
f := *v
return &f
}
+127
View File
@@ -0,0 +1,127 @@
package xai
import (
"encoding/json"
"net/http"
"testing"
"github.com/stretchr/testify/require"
)
func TestBuildBillingURL(t *testing.T) {
t.Parallel()
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing?format=credits", BuildBillingURL(true))
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing", BuildBillingURL(false))
}
func TestApplyCLIBillingHeaders(t *testing.T) {
t.Parallel()
req, err := http.NewRequest(http.MethodGet, BuildBillingURL(true), nil)
require.NoError(t, err)
ApplyCLIBillingHeaders(req, " token ")
require.Equal(t, "Bearer token", req.Header.Get("Authorization"))
require.Equal(t, CLITokenAuthValue, req.Header.Get(CLITokenAuthHeader))
require.Equal(t, CLIClientVersion, req.Header.Get(CLIClientVersionHeader))
require.Equal(t, "grok-pager/"+CLIClientVersion+" grok-shell/"+CLIClientVersion+" (macos; aarch64)", req.UserAgent())
}
func TestBuildBillingSummaryWeeklyAndMonthly(t *testing.T) {
t.Parallel()
weeklyBody := []byte(`{
"config": {
"currentPeriod": {"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"},
"creditUsagePercent": 2.0,
"productUsage": [{"product":"Api","usagePercent":2.0}]
}
}`)
monthlyBody := []byte(`{
"config": {
"monthlyLimit": {"val": 15000},
"used": {"val": 78},
"billingPeriodStart": "2026-07-01T00:00:00Z",
"billingPeriodEnd": "2026-08-01T00:00:00Z"
}
}`)
weeklyPayload, err := ParseBillingPayload(weeklyBody)
require.NoError(t, err)
monthlyPayload, err := ParseBillingPayload(monthlyBody)
require.NoError(t, err)
weekly := BuildBillingSummary(weeklyPayload.Config)
monthly := BuildBillingSummary(monthlyPayload.Config)
require.NotNil(t, weekly)
require.NotNil(t, monthly)
require.Equal(t, "weekly", weekly.PeriodType)
require.InDelta(t, 2.0, *weekly.UsagePercent, 1e-9)
require.Equal(t, "Api", weekly.ProductUsage[0].Product)
require.Equal(t, "SuperGrok", monthly.Plan)
require.InDelta(t, 15000, *monthly.MonthlyLimitCents, 1e-9)
require.InDelta(t, 78, *monthly.UsedCents, 1e-9)
require.InDelta(t, 0.52, *monthly.UsedPercent, 1e-2)
merged := MergeBillingProbeResult(nil, weekly, monthly, true, true)
require.Equal(t, "weekly", merged.PeriodType)
require.InDelta(t, 2.0, *merged.UsagePercent, 1e-9)
require.Equal(t, "SuperGrok", merged.Plan)
require.InDelta(t, 15000, *merged.MonthlyLimitCents, 1e-9)
require.Equal(t, "2026-08-01T00:00:00Z", merged.BillingPeriodEnd)
}
func TestParseCentValueBareNumber(t *testing.T) {
t.Parallel()
raw, _ := json.Marshal(15000)
v := parseCentValue(raw)
require.NotNil(t, v)
require.InDelta(t, 15000, *v, 1e-9)
}
func TestBuildBillingSummaryMonthlyOnlyKeepsWeeklyUsageEmpty(t *testing.T) {
t.Parallel()
payload, err := ParseBillingPayload([]byte(`{"config":{"monthlyLimit":{"val":15000},"used":{"val":7500},"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"}}`))
require.NoError(t, err)
summary := BuildBillingSummary(payload.Config)
require.NotNil(t, summary)
require.Equal(t, "monthly", summary.PeriodType)
require.Nil(t, summary.UsagePercent)
require.InDelta(t, 50, *summary.UsedPercent, 1e-9)
}
func TestMergeBillingProbeResultRetainsFailedWindow(t *testing.T) {
t.Parallel()
previous := &BillingSummary{
PeriodType: "weekly",
UsagePercent: floatPointer(100),
PeriodEnd: "2026-07-16T00:00:00Z",
MonthlyLimitCents: floatPointer(15000),
UsedPercent: floatPointer(20),
BillingPeriodEnd: "2026-08-01T00:00:00Z",
WeeklyUpdatedAt: "2026-07-10T00:00:00Z",
MonthlyUpdatedAt: "2026-07-10T00:00:00Z",
FailedWindows: []string{"monthly"},
}
monthly := &BillingSummary{
PeriodType: "monthly",
MonthlyLimitCents: floatPointer(15000),
UsedPercent: floatPointer(30),
BillingPeriodEnd: "2026-08-01T00:00:00Z",
}
merged := MergeBillingProbeResult(previous, nil, monthly, false, true)
require.Equal(t, "weekly", merged.PeriodType)
require.InDelta(t, 100, *merged.UsagePercent, 1e-9)
require.Equal(t, previous.WeeklyUpdatedAt, merged.WeeklyUpdatedAt)
require.InDelta(t, 30, *merged.UsedPercent, 1e-9)
require.NotEqual(t, previous.MonthlyUpdatedAt, merged.MonthlyUpdatedAt)
require.True(t, merged.Partial)
require.Equal(t, []string{"weekly"}, merged.FailedWindows)
require.Equal(t, []string{"monthly"}, previous.FailedWindows)
}
func floatPointer(value float64) *float64 {
return &value
}
+30 -1
View File
@@ -191,7 +191,7 @@ func RuntimeSanity() RuntimeSanityReport {
UnsafeURLOverrides: AllowUnsafeURLOverrides(),
UnsafeHighConcurrency: AllowUnsafeHighConcurrency(),
PublicGatewayScope: "responses_only",
ProxyPolicy: "account_proxy_optional; upstream URL allowlists enforced unless unsafe overrides are enabled",
ProxyPolicy: "account_proxy_optional; OAuth URLs use trusted-host allowlists; API-key base URLs require public HTTPS unless unsafe overrides are enabled",
}
}
@@ -252,6 +252,19 @@ func ValidateOAuthEndpointURL(raw string) (string, error) {
}
func ValidateBaseURL(raw string) (string, error) {
if AllowUnsafeURLOverrides() {
return urlvalidator.ValidateURLFormat(raw, true)
}
normalized, err := urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{
AllowPrivate: false,
})
if err != nil {
return "", err
}
return normalizeKnownBaseURLPath(normalized)
}
func ValidateTrustedBaseURL(raw string) (string, error) {
if AllowUnsafeURLOverrides() {
return urlvalidator.ValidateURLFormat(raw, true)
}
@@ -461,6 +474,22 @@ func BuildVideosGenerationsURL(baseURL string) (string, error) {
return validatedBaseURL + "/videos/generations", nil
}
func BuildVideosEditsURL(baseURL string) (string, error) {
validatedBaseURL, err := ValidatedBaseURL(baseURL)
if err != nil {
return "", fmt.Errorf("invalid base url: %w", err)
}
return validatedBaseURL + "/videos/edits", nil
}
func BuildVideosExtensionsURL(baseURL string) (string, error) {
validatedBaseURL, err := ValidatedBaseURL(baseURL)
if err != nil {
return "", fmt.Errorf("invalid base url: %w", err)
}
return validatedBaseURL + "/videos/extensions", nil
}
func BuildVideoURL(baseURL, requestID string) (string, error) {
validatedBaseURL, err := ValidatedBaseURL(baseURL)
if err != nil {
+19 -4
View File
@@ -129,6 +129,14 @@ func TestBuildGrokMediaURLs(t *testing.T) {
require.NoError(t, err)
require.Equal(t, DefaultBaseURL+"/videos/generations", videosURL)
videoEditsURL, err := BuildVideosEditsURL(DefaultBaseURL)
require.NoError(t, err)
require.Equal(t, DefaultBaseURL+"/videos/edits", videoEditsURL)
videoExtensionsURL, err := BuildVideosExtensionsURL(DefaultBaseURL)
require.NoError(t, err)
require.Equal(t, DefaultBaseURL+"/videos/extensions", videoExtensionsURL)
videoURL, err := BuildVideoURL(DefaultBaseURL, "req 123")
require.NoError(t, err)
require.Equal(t, DefaultBaseURL+"/videos/req%20123", videoURL)
@@ -137,13 +145,10 @@ func TestBuildGrokMediaURLs(t *testing.T) {
require.Error(t, err)
}
func TestValidateXAIURLsRejectArbitraryHostsByDefault(t *testing.T) {
func TestValidateXAIURLsRejectUntrustedOAuthAndUnsafeBaseURLsByDefault(t *testing.T) {
_, err := ValidateOAuthEndpointURL("https://auth.example.test/oauth2/token")
require.Error(t, err)
_, err = ValidateBaseURL("https://xai.test/v1")
require.Error(t, err)
_, err = ValidateBaseURL("http://127.0.0.1:8080/v1")
require.Error(t, err)
@@ -151,6 +156,15 @@ func TestValidateXAIURLsRejectArbitraryHostsByDefault(t *testing.T) {
require.Error(t, err)
}
func TestValidateBaseURLAllowsPublicThirdPartyGrokAPI(t *testing.T) {
baseURL, err := ValidateBaseURL("https://grok.example.test/v1/")
require.NoError(t, err)
require.Equal(t, "https://grok.example.test/v1", baseURL)
_, err = ValidateTrustedBaseURL("https://grok.example.test/v1")
require.Error(t, err)
}
func TestValidateXAIURLsAllowUnsafeDevOverride(t *testing.T) {
t.Setenv(EnvAllowUnsafeURLOverrides, "true")
@@ -182,6 +196,7 @@ func TestRuntimeSanityReportsSafeDefaults(t *testing.T) {
require.False(t, report.UnsafeHighConcurrency)
require.Equal(t, "responses_only", report.PublicGatewayScope)
require.Contains(t, report.ProxyPolicy, "account_proxy_optional")
require.Contains(t, report.ProxyPolicy, "API-key base URLs require public HTTPS")
}
func TestRuntimeSanityReportsInvalidOverridesWithoutSecrets(t *testing.T) {
+418
View File
@@ -0,0 +1,418 @@
package xai
import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"sort"
"strconv"
"strings"
"time"
)
const (
SSOBuildScope = "openid profile email offline_access grok-cli:access api:access conversations:read conversations:write"
SSOAccountsURL = "https://accounts.x.ai/"
SSODeviceURL = OAuthIssuer + "/oauth2/device/code"
SSOVerifyURL = OAuthIssuer + "/oauth2/device/verify"
SSOApproveURL = OAuthIssuer + "/oauth2/device/approve"
SSOTokenURL = OAuthIssuer + "/oauth2/token"
SSOConversionTimeout = 90 * time.Second
ssoMaxAuthBody = 2 << 20
ssoDefaultUA = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
ssoDefaultTokenTTL = 6 * time.Hour
)
var (
ErrSSOUnauthorized = errors.New("xai sso unauthorized")
ErrSSOAuthorizationDenied = errors.New("xai device authorization denied")
)
type SSOHTTPError struct{ Status int }
func (e SSOHTTPError) Error() string { return fmt.Sprintf("xAI OAuth HTTP %d", e.Status) }
type SSODeviceHTTPClient interface {
Do(*http.Request) (*http.Response, error)
}
type SSODeviceOptions struct {
HTTPClient SSODeviceHTTPClient
UserAgent string
Sleep func(context.Context, time.Duration) error
}
type ssoDeviceFlow struct {
client SSODeviceHTTPClient
userAgent string
cookies map[string]string
sleep func(context.Context, time.Duration) error
}
func ConvertSSOToBuild(ctx context.Context, ssoToken string, opts *SSODeviceOptions) (*TokenResponse, error) {
ssoToken = NormalizeSSOToken(ssoToken)
if ssoToken == "" {
return nil, ErrSSOUnauthorized
}
if opts == nil {
opts = &SSODeviceOptions{}
}
client := opts.HTTPClient
if client == nil {
client = &http.Client{
Timeout: SSOConversionTimeout,
CheckRedirect: func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
},
}
}
userAgent := strings.TrimSpace(opts.UserAgent)
if userAgent == "" {
userAgent = ssoDefaultUA
}
sleep := opts.Sleep
if sleep == nil {
sleep = sleepContext
}
flow := &ssoDeviceFlow{
client: client,
userAgent: userAgent,
cookies: map[string]string{"sso": ssoToken, "sso-rw": ssoToken},
sleep: sleep,
}
return flow.convert(ctx)
}
func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) {
status, finalURL, _, err := f.do(ctx, http.MethodGet, SSOAccountsURL, nil)
if err != nil {
return nil, err
}
if status == http.StatusUnauthorized || strings.Contains(finalURL, "sign-in") || strings.Contains(finalURL, "sign-up") {
return nil, ErrSSOUnauthorized
}
if status < 200 || status >= 400 {
return nil, fmt.Errorf("validate Grok Web SSO: %w", SSOHTTPError{Status: status})
}
status, _, body, err := f.do(ctx, http.MethodPost, SSODeviceURL, url.Values{
"client_id": {DefaultClientID},
"scope": {SSOBuildScope},
})
if err != nil {
return nil, err
}
if status < 200 || status >= 300 {
return nil, fmt.Errorf("start xAI device flow: %w", SSOHTTPError{Status: status})
}
var device struct {
DeviceCode string `json:"device_code"`
UserCode string `json:"user_code"`
VerificationURIComplete string `json:"verification_uri_complete"`
Interval int `json:"interval"`
ExpiresIn int `json:"expires_in"`
}
if err := json.Unmarshal(body, &device); err != nil {
return nil, fmt.Errorf("parse xAI device flow response: %w", err)
}
if device.DeviceCode == "" || device.UserCode == "" || !safeXAIAuthURL(device.VerificationURIComplete) {
return nil, errors.New("xAI device flow response is incomplete")
}
if device.Interval <= 0 {
device.Interval = 5
}
if device.ExpiresIn <= 0 {
device.ExpiresIn = 1800
}
status, _, _, err = f.do(ctx, http.MethodGet, device.VerificationURIComplete, nil)
if err != nil {
return nil, err
}
if status < 200 || status >= 400 {
return nil, fmt.Errorf("open xAI device verification page: %w", SSOHTTPError{Status: status})
}
status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOVerifyURL, url.Values{"user_code": {device.UserCode}})
if err != nil {
return nil, err
}
if status < 200 || status >= 400 {
return nil, fmt.Errorf("verify xAI device code: %w", SSOHTTPError{Status: status})
}
if !strings.Contains(finalURL, "consent") {
return nil, errors.New("xAI device verification did not reach consent page")
}
status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOApproveURL, url.Values{
"user_code": {device.UserCode},
"action": {"allow"},
"principal_type": {"User"},
"principal_id": {""},
})
if err != nil {
return nil, err
}
if status < 200 || status >= 400 {
return nil, fmt.Errorf("approve xAI device code: %w", SSOHTTPError{Status: status})
}
if !strings.Contains(finalURL, "done") {
return nil, errors.New("xAI device approval did not reach done page")
}
return f.pollToken(ctx, device.DeviceCode, time.Duration(device.Interval)*time.Second, time.Duration(device.ExpiresIn)*time.Second)
}
func (f *ssoDeviceFlow) pollToken(ctx context.Context, deviceCode string, interval, expiresIn time.Duration) (*TokenResponse, error) {
if interval < time.Second {
interval = time.Second
}
deadline := time.Now().Add(minDuration(expiresIn, 75*time.Second))
for time.Now().Before(deadline) {
if err := f.sleep(ctx, interval); err != nil {
return nil, err
}
status, _, body, err := f.do(ctx, http.MethodPost, SSOTokenURL, url.Values{
"grant_type": {"urn:ietf:params:oauth:grant-type:device_code"},
"client_id": {DefaultClientID},
"device_code": {deviceCode},
})
if err != nil {
return nil, err
}
var payload struct {
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
IDToken string `json:"id_token"`
TokenType string `json:"token_type"`
ExpiresIn int64 `json:"expires_in"`
Scope string `json:"scope"`
Error string `json:"error"`
ErrorDescription string `json:"error_description"`
}
if err := json.Unmarshal(body, &payload); err != nil {
return nil, fmt.Errorf("parse xAI token response: %w", err)
}
if status >= 200 && status < 300 && payload.AccessToken != "" {
if payload.ExpiresIn <= 0 {
payload.ExpiresIn = int64(ssoDefaultTokenTTL.Seconds())
}
if payload.TokenType == "" {
payload.TokenType = "Bearer"
}
return &TokenResponse{
AccessToken: payload.AccessToken,
RefreshToken: payload.RefreshToken,
IDToken: payload.IDToken,
TokenType: payload.TokenType,
ExpiresIn: payload.ExpiresIn,
Scope: payload.Scope,
}, nil
}
switch payload.Error {
case "authorization_pending":
continue
case "slow_down":
interval += 5 * time.Second
continue
case "access_denied", "expired_token":
return nil, ErrSSOAuthorizationDenied
default:
if status >= 400 {
return nil, fmt.Errorf("xAI token polling failed (%s): %w", firstNonEmpty(payload.ErrorDescription, payload.Error), SSOHTTPError{Status: status})
}
return nil, fmt.Errorf("xAI token polling failed: %s", firstNonEmpty(payload.ErrorDescription, payload.Error, strconv.Itoa(status)))
}
}
return nil, errors.New("xAI device flow token polling timed out")
}
func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form url.Values) (int, string, []byte, error) {
if !safeXAIAuthURL(endpoint) {
return 0, "", nil, errors.New("xAI OAuth URL is not trusted")
}
currentURL := endpoint
currentMethod := method
currentForm := form
for redirects := 0; redirects <= 8; redirects++ {
var body io.Reader
if currentForm != nil {
body = strings.NewReader(currentForm.Encode())
}
request, err := http.NewRequestWithContext(ctx, currentMethod, currentURL, body)
if err != nil {
return 0, currentURL, nil, err
}
request.Header.Set("Accept", "application/json, text/html;q=0.9, */*;q=0.8")
request.Header.Set("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8")
request.Header.Set("User-Agent", f.userAgent)
if cookie := f.cookieHeader(); cookie != "" {
request.Header.Set("Cookie", cookie)
}
if currentForm != nil {
request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
}
response, err := f.client.Do(request)
if err != nil {
return 0, currentURL, nil, err
}
f.captureCookies(response)
data, readErr := io.ReadAll(io.LimitReader(response.Body, ssoMaxAuthBody+1))
_ = response.Body.Close()
if readErr != nil {
return response.StatusCode, currentURL, nil, readErr
}
if len(data) > ssoMaxAuthBody {
return response.StatusCode, currentURL, nil, errors.New("xAI OAuth response exceeds 2 MiB")
}
if response.StatusCode < 300 || response.StatusCode > 399 {
return response.StatusCode, currentURL, data, nil
}
location := strings.TrimSpace(response.Header.Get("Location"))
if location == "" {
return response.StatusCode, currentURL, data, errors.New("xAI OAuth redirect missing Location")
}
base, _ := url.Parse(currentURL)
next, err := url.Parse(location)
if err != nil {
return response.StatusCode, currentURL, data, err
}
currentURL = base.ResolveReference(next).String()
if !safeXAIAuthURL(currentURL) {
return response.StatusCode, currentURL, data, errors.New("xAI OAuth redirected to untrusted host")
}
if response.StatusCode == http.StatusSeeOther || ((response.StatusCode == http.StatusMovedPermanently || response.StatusCode == http.StatusFound) && currentMethod != http.MethodGet && currentMethod != http.MethodHead) {
currentMethod = http.MethodGet
currentForm = nil
}
}
return 0, currentURL, nil, errors.New("xAI OAuth redirected too many times")
}
func (f *ssoDeviceFlow) captureCookies(response *http.Response) {
for _, cookie := range response.Cookies() {
name := strings.TrimSpace(cookie.Name)
value := strings.TrimSpace(cookie.Value)
if name == "" || len(name) > 128 || len(value) > 16384 || strings.ContainsAny(name+value, "\r\n\x00") {
continue
}
if cookie.MaxAge < 0 {
delete(f.cookies, name)
continue
}
f.cookies[name] = value
}
}
func (f *ssoDeviceFlow) cookieHeader() string {
keys := make([]string, 0, len(f.cookies))
for key := range f.cookies {
keys = append(keys, key)
}
sort.Strings(keys)
parts := make([]string, 0, len(keys))
for _, key := range keys {
parts = append(parts, key+"="+f.cookies[key])
}
return strings.Join(parts, "; ")
}
func safeXAIAuthURL(raw string) bool {
parsed, err := url.Parse(raw)
if err != nil || parsed.User != nil || parsed.Hostname() == "" {
return false
}
if AllowUnsafeURLOverrides() {
return parsed.Scheme != "" && parsed.Host != ""
}
if parsed.Scheme != "https" {
return false
}
host := strings.ToLower(parsed.Hostname())
return host == "x.ai" || strings.HasSuffix(host, ".x.ai")
}
func NormalizeSSOToken(value string) string {
value = strings.TrimSpace(value)
if strings.HasPrefix(strings.ToLower(value), "cookie:") {
value = strings.TrimSpace(value[len("cookie:"):])
}
for _, part := range strings.Split(value, ";") {
name, token, found := strings.Cut(strings.TrimSpace(part), "=")
if !found {
continue
}
switch strings.ToLower(strings.TrimSpace(name)) {
case "sso", "sso-rw":
return sanitizeSSOToken(token)
}
}
if token, _, found := strings.Cut(value, ";"); found {
value = strings.TrimSpace(token)
}
return sanitizeSSOToken(value)
}
func sanitizeSSOToken(value string) string {
return strings.NewReplacer("\r", "", "\n", "", "\x00", "").Replace(strings.TrimSpace(value))
}
func DecodeJWTClaims(token string) map[string]any {
parts := strings.Split(token, ".")
if len(parts) < 2 {
return nil
}
payload, err := base64.RawURLEncoding.DecodeString(parts[1])
if err != nil {
return nil
}
var claims map[string]any
if err := json.Unmarshal(payload, &claims); err != nil {
return nil
}
return claims
}
func JWTClaimString(claims map[string]any, key string) string {
value, _ := claims[key].(string)
return strings.TrimSpace(value)
}
func sleepContext(ctx context.Context, d time.Duration) error {
timer := time.NewTimer(d)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
}
func minDuration(a, b time.Duration) time.Duration {
if a <= 0 {
return b
}
if a < b {
return a
}
return b
}
func firstNonEmpty(values ...string) string {
for _, value := range values {
if strings.TrimSpace(value) != "" {
return strings.TrimSpace(value)
}
}
return ""
}
+115
View File
@@ -0,0 +1,115 @@
//go:build unit
package xai
import (
"context"
"io"
"net/http"
"net/url"
"strings"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type ssoDeviceFakeClient struct {
t *testing.T
tokenCalls int
cookieHeaders []string
}
func (c *ssoDeviceFakeClient) Do(req *http.Request) (*http.Response, error) {
c.cookieHeaders = append(c.cookieHeaders, req.Header.Get("Cookie"))
switch req.URL.String() {
case SSOAccountsURL:
require.Equal(c.t, http.MethodGet, req.Method)
return ssoDeviceResponse(http.StatusOK, http.Header{"Set-Cookie": {"session=web-session; Path=/"}}, `{}`), nil
case SSODeviceURL:
require.Equal(c.t, http.MethodPost, req.Method)
values := readSSODeviceForm(c.t, req)
require.Equal(c.t, DefaultClientID, values.Get("client_id"))
require.Equal(c.t, SSOBuildScope, values.Get("scope"))
return ssoDeviceResponse(http.StatusOK, http.Header{"Set-Cookie": {"csrf=csrf-token; Path=/"}}, `{"device_code":"device-1","user_code":"USER-1","verification_uri_complete":"https://auth.x.ai/oauth2/device/complete","interval":1,"expires_in":60}`), nil
case "https://auth.x.ai/oauth2/device/complete":
require.Equal(c.t, http.MethodGet, req.Method)
return ssoDeviceResponse(http.StatusOK, nil, `<html>ok</html>`), nil
case SSOVerifyURL:
require.Equal(c.t, http.MethodPost, req.Method)
values := readSSODeviceForm(c.t, req)
require.Equal(c.t, "USER-1", values.Get("user_code"))
return ssoDeviceResponse(http.StatusFound, http.Header{"Location": {"/oauth2/device/consent"}}, ``), nil
case "https://auth.x.ai/oauth2/device/consent":
require.Equal(c.t, http.MethodGet, req.Method)
return ssoDeviceResponse(http.StatusOK, nil, `<html>consent</html>`), nil
case SSOApproveURL:
require.Equal(c.t, http.MethodPost, req.Method)
values := readSSODeviceForm(c.t, req)
require.Equal(c.t, "USER-1", values.Get("user_code"))
require.Equal(c.t, "allow", values.Get("action"))
require.Equal(c.t, "User", values.Get("principal_type"))
return ssoDeviceResponse(http.StatusSeeOther, http.Header{"Location": {"/oauth2/device/done"}}, ``), nil
case "https://auth.x.ai/oauth2/device/done":
require.Equal(c.t, http.MethodGet, req.Method)
return ssoDeviceResponse(http.StatusOK, nil, `<html>done</html>`), nil
case SSOTokenURL:
require.Equal(c.t, http.MethodPost, req.Method)
c.tokenCalls++
values := readSSODeviceForm(c.t, req)
require.Equal(c.t, "urn:ietf:params:oauth:grant-type:device_code", values.Get("grant_type"))
require.Equal(c.t, "device-1", values.Get("device_code"))
return ssoDeviceResponse(http.StatusOK, nil, `{"access_token":"access-token","refresh_token":"refresh-token","id_token":"id-token","token_type":"Bearer","expires_in":3600,"scope":"`+SSOBuildScope+`"}`), nil
default:
c.t.Fatalf("unexpected request: %s %s", req.Method, req.URL.String())
return nil, nil
}
}
func TestConvertSSOToBuildCompletesDeviceFlow(t *testing.T) {
t.Setenv(EnvClientID, "")
client := &ssoDeviceFakeClient{t: t}
token, err := ConvertSSOToBuild(context.Background(), "sso=sso-token; ignored=1", &SSODeviceOptions{
HTTPClient: client,
Sleep: func(context.Context, time.Duration) error {
return nil
},
})
require.NoError(t, err)
require.Equal(t, "access-token", token.AccessToken)
require.Equal(t, "refresh-token", token.RefreshToken)
require.Equal(t, "id-token", token.IDToken)
require.Equal(t, SSOBuildScope, token.Scope)
require.Equal(t, 1, client.tokenCalls)
require.Contains(t, client.cookieHeaders[0], "sso=sso-token")
require.Contains(t, client.cookieHeaders[0], "sso-rw=sso-token")
require.Contains(t, client.cookieHeaders[len(client.cookieHeaders)-1], "session=web-session")
require.Contains(t, client.cookieHeaders[len(client.cookieHeaders)-1], "csrf=csrf-token")
}
func TestNormalizeSSOTokenAcceptsCookieHeader(t *testing.T) {
require.Equal(t, "token-1", NormalizeSSOToken("Cookie: foo=bar; sso=token-1; sso-rw=token-2"))
require.Equal(t, "token-2", NormalizeSSOToken("sso-rw=token-2; foo=bar"))
require.Equal(t, "raw-token", NormalizeSSOToken(" raw-token ; ignored=1"))
}
func ssoDeviceResponse(status int, header http.Header, body string) *http.Response {
if header == nil {
header = http.Header{}
}
return &http.Response{
StatusCode: status,
Header: header,
Body: io.NopCloser(strings.NewReader(body)),
}
}
func readSSODeviceForm(t *testing.T, req *http.Request) url.Values {
t.Helper()
data, err := io.ReadAll(req.Body)
require.NoError(t, err)
values, err := url.ParseQuery(string(data))
require.NoError(t, err)
return values
}
+23 -7
View File
@@ -61,6 +61,7 @@ var schedulerNeutralExtraKeyPrefixes = []string{
var schedulerNeutralExtraKeys = map[string]struct{}{
"codex_usage_updated_at": {},
"grok_billing_snapshot": {},
"session_window_utilization": {},
}
@@ -1553,7 +1554,7 @@ func (r *accountRepository) SetSchedulable(ctx context.Context, id int64, schedu
}
func (r *accountRepository) AutoPauseExpiredAccounts(ctx context.Context, now time.Time) (int64, error) {
result, err := r.sql.ExecContext(ctx, `
rows, err := r.sql.QueryContext(ctx, `
UPDATE accounts
SET schedulable = FALSE,
updated_at = NOW()
@@ -1562,20 +1563,35 @@ func (r *accountRepository) AutoPauseExpiredAccounts(ctx context.Context, now ti
AND auto_pause_on_expired = TRUE
AND expires_at IS NOT NULL
AND expires_at <= $1
RETURNING id
`, now)
if err != nil {
return 0, err
}
rows, err := result.RowsAffected()
if err != nil {
defer func() {
_ = rows.Close()
}()
accountIDs := make([]int64, 0)
for rows.Next() {
var accountID int64
if err := rows.Scan(&accountID); err != nil {
return 0, err
}
accountIDs = append(accountIDs, accountID)
}
if err := rows.Err(); err != nil {
return 0, err
}
if rows > 0 {
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventFullRebuild, nil, nil, nil); err != nil {
logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue auto pause rebuild failed: err=%v", err)
if len(accountIDs) > 0 {
// 只刷新本次暂停的账号及其所属分组,避免少量账号到期触发所有调度桶重建。
payload := map[string]any{"account_ids": accountIDs}
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil {
logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue auto pause account changes failed: err=%v", err)
}
}
return rows, nil
return int64(len(accountIDs)), nil
}
func (r *accountRepository) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
@@ -0,0 +1,72 @@
package repository
import (
"context"
"database/sql/driver"
"encoding/json"
"reflect"
"regexp"
"testing"
"time"
"github.com/DATA-DOG/go-sqlmock"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
type accountIDsPayloadMatcher struct {
want []int64
}
func (m accountIDsPayloadMatcher) Match(value driver.Value) bool {
raw, ok := value.([]byte)
if !ok {
return false
}
var payload struct {
AccountIDs []int64 `json:"account_ids"`
}
if err := json.Unmarshal(raw, &payload); err != nil {
return false
}
return reflect.DeepEqual(m.want, payload.AccountIDs)
}
func TestAutoPauseExpiredAccountsEnqueuesAffectedAccounts(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
now := time.Now()
mock.ExpectQuery(`(?s)UPDATE accounts.*RETURNING id`).
WithArgs(now).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(11)).AddRow(int64(29)))
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)")).
WithArgs(service.SchedulerOutboxEventAccountBulkChanged, nil, nil, accountIDsPayloadMatcher{want: []int64{11, 29}}).
WillReturnResult(sqlmock.NewResult(1, 1))
repo := newAccountRepositoryWithSQL(nil, db, nil)
updated, err := repo.AutoPauseExpiredAccounts(context.Background(), now)
require.NoError(t, err)
require.EqualValues(t, 2, updated)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestAutoPauseExpiredAccountsSkipsOutboxWithoutChanges(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
now := time.Now()
mock.ExpectQuery(`(?s)UPDATE accounts.*RETURNING id`).
WithArgs(now).
WillReturnRows(sqlmock.NewRows([]string{"id"}))
repo := newAccountRepositoryWithSQL(nil, db, nil)
updated, err := repo.AutoPauseExpiredAccounts(context.Background(), now)
require.NoError(t, err)
require.Zero(t, updated)
require.NoError(t, mock.ExpectationsWereMet())
}
@@ -0,0 +1,16 @@
package repository
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestGrokBillingSnapshotIsSchedulerNeutral(t *testing.T) {
t.Parallel()
require.True(t, isSchedulerNeutralExtraKey("grok_billing_snapshot"))
require.False(t, shouldEnqueueSchedulerOutboxForExtraUpdates(map[string]any{
"grok_billing_snapshot": map[string]any{"usage_percent": 50},
}))
}
+12 -10
View File
@@ -525,17 +525,19 @@ func (r *apiKeyRepository) latestUsageLogIPs(ctx context.Context, apiKeyIDs []in
func latestUsageLogIPsQuery(apiKeyIDs []int64, dialectName string) (string, []any) {
if dialectName == dialect.Postgres {
// Keep each key lookup bounded to one ordered index probe instead of ranking its full history.
return `
SELECT api_key_id, ip_address
FROM (
SELECT api_key_id, ip_address,
ROW_NUMBER() OVER (PARTITION BY api_key_id ORDER BY created_at DESC, id DESC) AS rn
FROM usage_logs
WHERE api_key_id = ANY($1::bigint[])
AND ip_address IS NOT NULL
AND ip_address <> ''
) ranked
WHERE rn = 1`, []any{pq.Array(apiKeyIDs)}
SELECT requested.api_key_id, latest.ip_address
FROM unnest($1::bigint[]) AS requested(api_key_id)
CROSS JOIN LATERAL (
SELECT ul.ip_address
FROM usage_logs AS ul
WHERE ul.api_key_id = requested.api_key_id
AND ul.ip_address IS NOT NULL
AND ul.ip_address <> ''
ORDER BY ul.created_at DESC, ul.id DESC
LIMIT 1
) AS latest`, []any{pq.Array(apiKeyIDs)}
}
placeholders := make([]string, len(apiKeyIDs))
@@ -3,6 +3,7 @@ package repository
import (
"context"
"database/sql"
"strings"
"testing"
"time"
@@ -125,6 +126,20 @@ func TestAPIKeyRepositoryListByUserIDAttachesLastUsedIP(t *testing.T) {
require.Nil(t, byID[noLogs.ID].LastUsedIP)
}
func TestLatestUsageLogIPsQueryPostgresUsesPerKeyLateralLookup(t *testing.T) {
query, args := latestUsageLogIPsQuery([]int64{11, 22}, dialect.Postgres)
normalizedQuery := strings.Join(strings.Fields(query), " ")
require.Contains(t, normalizedQuery, "FROM unnest($1::bigint[]) AS requested(api_key_id)")
require.Contains(t, normalizedQuery, "CROSS JOIN LATERAL")
require.Contains(t, normalizedQuery, "WHERE ul.api_key_id = requested.api_key_id")
require.Contains(t, normalizedQuery, "AND ul.ip_address IS NOT NULL")
require.Contains(t, normalizedQuery, "AND ul.ip_address <> ''")
require.Contains(t, normalizedQuery, "ORDER BY ul.created_at DESC, ul.id DESC LIMIT 1")
require.NotContains(t, normalizedQuery, "ROW_NUMBER")
require.Len(t, args, 1)
}
func TestAPIKeyRepository_CreateWithLastUsedAt(t *testing.T) {
repo, client := newAPIKeyRepoSQLite(t)
ctx := context.Background()
@@ -13,6 +13,7 @@ import (
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/Wei-Shaw/sub2api/internal/service"
)
@@ -63,12 +64,14 @@ func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader,
return 0, fmt.Errorf("read body: %w", err)
}
finish := servertiming.ObserveDependency(ctx, "s3")
_, err = s.client.PutObject(ctx, &s3.PutObjectInput{
Bucket: &s.bucket,
Key: &key,
Body: bytes.NewReader(data),
ContentType: &contentType,
})
finish()
if err != nil {
return 0, fmt.Errorf("S3 PutObject: %w", err)
}
@@ -76,10 +79,12 @@ func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader,
}
func (s *S3BackupStore) Download(ctx context.Context, key string) (io.ReadCloser, error) {
finish := servertiming.ObserveDependency(ctx, "s3")
result, err := s.client.GetObject(ctx, &s3.GetObjectInput{
Bucket: &s.bucket,
Key: &key,
})
finish()
if err != nil {
return nil, fmt.Errorf("S3 GetObject: %w", err)
}
@@ -87,10 +92,12 @@ func (s *S3BackupStore) Download(ctx context.Context, key string) (io.ReadCloser
}
func (s *S3BackupStore) Delete(ctx context.Context, key string) error {
finish := servertiming.ObserveDependency(ctx, "s3")
_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
Bucket: &s.bucket,
Key: &key,
})
finish()
return err
}
@@ -107,9 +114,11 @@ func (s *S3BackupStore) PresignURL(ctx context.Context, key string, expiry time.
}
func (s *S3BackupStore) HeadBucket(ctx context.Context) error {
finish := servertiming.ObserveDependency(ctx, "s3")
_, err := s.client.HeadBucket(ctx, &s3.HeadBucketInput{
Bucket: &s.bucket,
})
finish()
if err != nil {
return fmt.Errorf("S3 HeadBucket failed: %w", err)
}
@@ -276,5 +276,5 @@ func createReqClient(proxyURL string) (*req.Client, error) {
client.SetProxyURL(trimmed)
}
return client, nil
return instrumentReqClient(client), nil
}
@@ -30,6 +30,10 @@ const (
userSlotKeyPrefix = "concurrency:user:"
// 格式: concurrency:api_key:{apiKeyID}
apiKeySlotKeyPrefix = "concurrency:api_key:"
// API-key-scoped client WebSocket ingress leases use a shorter TTL than
// ordinary request slots, because idle ingress sessions do not hold a turn slot.
openAIWSIngressLeaseKeyPrefix = "concurrency:openai_ws_ingress:api_key:"
openAIWSIngressLeaseTTLSeconds = 60
// 等待队列计数器格式: concurrency:wait:{userID}
waitQueueKeyPrefix = "concurrency:wait:"
// 账号级等待队列计数器格式: wait:account:{accountID}
@@ -138,6 +142,49 @@ var (
return 1
`)
// acquireOpenAIWSIngressLeaseScript atomically reaps crashed members and
// acquires or refreshes one API-key-scoped ingress lease using Redis TIME.
acquireOpenAIWSIngressLeaseScript = redis.NewScript(`
redis.replicate_commands()
local key = KEYS[1]
local maxConnections = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])
local leaseID = ARGV[3]
local now = tonumber(redis.call('TIME')[1])
local expireBefore = now - ttl
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
if redis.call('ZSCORE', key, leaseID) ~= false then
redis.call('ZADD', key, now, leaseID)
redis.call('EXPIRE', key, ttl)
return 1
end
if redis.call('ZCARD', key) < maxConnections then
redis.call('ZADD', key, now, leaseID)
redis.call('EXPIRE', key, ttl)
return 1
end
return 0
`)
// refreshOpenAIWSIngressLeaseScript does not recreate a missing member: a
// process that lost its lease must terminate its local WebSocket instead of
// silently continuing beyond the distributed cap.
refreshOpenAIWSIngressLeaseScript = redis.NewScript(`
redis.replicate_commands()
local key = KEYS[1]
local ttl = tonumber(ARGV[1])
local leaseID = ARGV[2]
local now = tonumber(redis.call('TIME')[1])
local expireBefore = now - ttl
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
if redis.call('ZSCORE', key, leaseID) == false then
return 0
end
redis.call('ZADD', key, now, leaseID)
redis.call('EXPIRE', key, ttl)
return 1
`)
// incrementWaitScript - refreshes TTL on each increment to keep queue depth accurate
// KEYS[1] = wait queue key
// ARGV[1] = maxWait
@@ -283,6 +330,10 @@ func apiKeySlotKey(apiKeyID int64) string {
return fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID)
}
func openAIWSIngressLeaseKey(apiKeyID int64) string {
return fmt.Sprintf("%s%d", openAIWSIngressLeaseKeyPrefix, apiKeyID)
}
func waitQueueKey(userID int64) string {
return fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID)
}
@@ -623,6 +674,48 @@ func (c *concurrencyCache) ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64
return c.rdb.ZRem(ctx, key, requestID).Err()
}
func (c *concurrencyCache) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error) {
if c == nil || c.rdb == nil || apiKeyID <= 0 || maxConnections <= 0 || leaseID == "" {
return false, nil
}
result, err := acquireOpenAIWSIngressLeaseScript.Run(
ctx,
c.rdb,
[]string{openAIWSIngressLeaseKey(apiKeyID)},
maxConnections,
openAIWSIngressLeaseTTLSeconds,
leaseID,
).Int()
if err != nil {
return false, err
}
return result == 1, nil
}
func (c *concurrencyCache) RefreshOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) (bool, error) {
if c == nil || c.rdb == nil || apiKeyID <= 0 || leaseID == "" {
return false, nil
}
result, err := refreshOpenAIWSIngressLeaseScript.Run(
ctx,
c.rdb,
[]string{openAIWSIngressLeaseKey(apiKeyID)},
openAIWSIngressLeaseTTLSeconds,
leaseID,
).Int()
if err != nil {
return false, err
}
return result == 1, nil
}
func (c *concurrencyCache) ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error {
if c == nil || c.rdb == nil || apiKeyID <= 0 || leaseID == "" {
return nil
}
return c.rdb.ZRem(ctx, openAIWSIngressLeaseKey(apiKeyID), leaseID).Err()
}
func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) {
if len(apiKeyIDs) == 0 {
return map[int64]int{}, nil
@@ -50,6 +50,53 @@ func (s *ConcurrencyCacheSuite) apiKeyConcurrencyCache() apiKeyConcurrencyCacheF
return cache
}
func (s *ConcurrencyCacheSuite) TestOpenAIWSIngressAPIKeySlot_HardLimitRefreshAndRelease() {
apiKeyID := int64(9011)
firstLeaseID := "ingress-first"
secondLeaseID := "ingress-second"
ok, err := s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 1, firstLeaseID)
require.NoError(s.T(), err)
require.True(s.T(), ok)
ok, err = s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 1, secondLeaseID)
require.NoError(s.T(), err)
require.False(s.T(), ok, "a second live session must not exceed the API key limit")
ok, err = s.rawCache.RefreshOpenAIWSIngressLease(s.ctx, apiKeyID, firstLeaseID)
require.NoError(s.T(), err)
require.True(s.T(), ok, "the current owner must be able to refresh its lease")
require.NoError(s.T(), s.rawCache.ReleaseOpenAIWSIngressLease(s.ctx, apiKeyID, firstLeaseID))
ok, err = s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 1, secondLeaseID)
require.NoError(s.T(), err)
require.True(s.T(), ok, "released capacity must become available immediately")
}
func (s *ConcurrencyCacheSuite) TestOpenAIWSIngressAPIKeySlot_ReapsCrashedLeaseWithoutDeletingLiveOtherInstance() {
apiKeyID := int64(9012)
key := openAIWSIngressLeaseKey(apiKeyID)
now, err := s.rawCache.redisUnixSeconds(s.ctx)
require.NoError(s.T(), err)
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, key,
redis.Z{Score: float64(now - openAIWSIngressLeaseTTLSeconds - 1), Member: "crashed-instance"},
redis.Z{Score: float64(now), Member: "live-other-instance"},
).Err())
require.NoError(s.T(), s.rdb.Expire(s.ctx, key, time.Duration(openAIWSIngressLeaseTTLSeconds)*time.Second).Err())
ok, err := s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 2, "new-instance")
require.NoError(s.T(), err)
require.True(s.T(), ok, "the crashed member should be reaped before enforcing the limit")
_, err = s.rdb.ZScore(s.ctx, key, "crashed-instance").Result()
require.ErrorIs(s.T(), err, redis.Nil)
_, err = s.rdb.ZScore(s.ctx, key, "live-other-instance").Result()
require.NoError(s.T(), err, "a live lease owned by another instance must be preserved")
count, err := s.rdb.ZCard(s.ctx, key).Result()
require.NoError(s.T(), err)
require.Equal(s.T(), int64(2), count)
}
func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() {
accountID := int64(10)
reqID1, reqID2, reqID3 := "req1", "req2", "req3"
+14 -4
View File
@@ -15,7 +15,7 @@ import (
"entgo.io/ent/dialect"
entsql "entgo.io/ent/dialect/sql"
_ "github.com/lib/pq" // PostgreSQL 驱动,通过副作用导入注册驱动
"github.com/lib/pq"
)
// InitEnt 初始化 Ent ORM 客户端并返回客户端实例和底层的 *sql.DB。
@@ -48,9 +48,19 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) {
// 使用 Ent 的 SQL 驱动打开 PostgreSQL 连接。
// dialect.Postgres 指定使用 PostgreSQL 方言进行 SQL 生成。
drv, err := entsql.Open(dialect.Postgres, dsn)
if err != nil {
return nil, nil, err
var drv *entsql.Driver
if cfg.Server.EnableServerTiming {
connector, err := pq.NewConnector(dsn)
if err != nil {
return nil, nil, err
}
drv = entsql.OpenDB(dialect.Postgres, sql.OpenDB(newServerTimingConnector(connector)))
} else {
var err error
drv, err = entsql.Open(dialect.Postgres, dsn)
if err != nil {
return nil, nil, err
}
}
applyDBPoolSettings(drv.DB(), cfg)
@@ -2,12 +2,14 @@ package repository
import (
"context"
"errors"
"net/http"
"net/url"
"strings"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
sharedhttp "github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/Wei-Shaw/sub2api/internal/util/logredact"
@@ -88,6 +90,21 @@ func (c *grokOAuthClient) RefreshToken(ctx context.Context, refreshToken, proxyU
return &tokenResp, nil
}
func (c *grokOAuthClient) ConvertSSOToBuild(ctx context.Context, ssoToken, proxyURL string) (*xai.TokenResponse, error) {
client, err := createGrokSSOHTTPClient(proxyURL)
if err != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "GROK_SSO_CLIENT_INIT_FAILED", "create HTTP client: %v", err)
}
requestCtx, cancel := context.WithTimeout(ctx, xai.SSOConversionTimeout)
defer cancel()
tokenResp, err := xai.ConvertSSOToBuild(requestCtx, ssoToken, &xai.SSODeviceOptions{HTTPClient: client})
if err != nil {
return nil, grokSSOConversionError(err)
}
return tokenResp, nil
}
func createGrokReqClient(proxyURL string) (*req.Client, error) {
return getSharedReqClient(reqClientOptions{
ProxyURL: proxyURL,
@@ -95,6 +112,43 @@ func createGrokReqClient(proxyURL string) (*req.Client, error) {
})
}
func createGrokSSOHTTPClient(proxyURL string) (*http.Client, error) {
client, err := sharedhttp.GetClient(sharedhttp.Options{
ProxyURL: proxyURL,
Timeout: xai.SSOConversionTimeout,
ResponseHeaderTimeout: 30 * time.Second,
})
if err != nil {
return nil, err
}
clone := *client
clone.CheckRedirect = func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
}
return &clone, nil
}
func grokSSOConversionError(err error) error {
if errors.Is(err, xai.ErrSSOUnauthorized) {
return infraerrors.New(http.StatusUnauthorized, "GROK_SSO_UNAUTHORIZED", "Grok Web SSO cookie is invalid or expired")
}
if errors.Is(err, xai.ErrSSOAuthorizationDenied) {
return infraerrors.New(http.StatusForbidden, "GROK_SSO_AUTHORIZATION_DENIED", "xAI device authorization was denied or expired")
}
var statusErr xai.SSOHTTPError
if errors.As(err, &statusErr) {
statusCode := http.StatusBadGateway
if statusErr.Status == http.StatusForbidden {
statusCode = http.StatusForbidden
}
return infraerrors.Newf(statusCode, "GROK_SSO_UPSTREAM_FAILED", "xAI SSO conversion failed: %v", err)
}
if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
return infraerrors.Newf(http.StatusGatewayTimeout, "GROK_SSO_TIMEOUT", "xAI SSO conversion timed out: %v", err)
}
return infraerrors.Newf(http.StatusBadGateway, "GROK_SSO_CONVERSION_FAILED", "xAI SSO conversion failed: %v", err)
}
func grokOAuthStatusError(code, message string, resp *req.Response) error {
statusCode := http.StatusBadGateway
errorCode := code
+32 -2
View File
@@ -21,10 +21,12 @@ import (
"github.com/andybalholm/brotli"
"github.com/klauspost/compress/zstd"
"golang.org/x/net/http2"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
@@ -59,6 +61,13 @@ const (
defaultOpenAIHTTP2FallbackErrorThreshold = 2
defaultOpenAIHTTP2FallbackWindow = 60 * time.Second
defaultOpenAIHTTP2FallbackTTL = 10 * time.Minute
// OpenAI HTTP/2 连接健康探测:Codex 上游改走 HTTP/2 后,池化连接被代理/NAT
// 静默掐断会成为“死连接”(两端都以为存活),请求落上去会挂到 TCP 重传超时
// (分钟级)。Go 的 http2.Transport 默认 ReadIdleTimeout=0(不发健康 PING),
// 无法检测。启用主动 PING 探测:连接空闲 ReadIdleTimeout 后发 PINGPingTimeout
// 内无响应即判定死连接并关闭,从源头避免请求挂在死连接上。
openAIHTTP2ReadIdleTimeout = 15 * time.Second
openAIHTTP2PingTimeout = 15 * time.Second
// The Grok CLI proxy rejects requests that do not identify a supported
// client version. Keep a known-good stable version in the binary while
@@ -186,7 +195,7 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i
}
// 执行请求
resp, err := entry.client.Do(req)
resp, err := servertiming.Do(entry.client, req)
if err != nil {
s.recordOpenAIHTTP2Failure(profile, entry.protocolMode, entry.proxyKey, err)
// 请求失败,立即减少计数
@@ -243,7 +252,7 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
return nil, err
}
resp, err := entry.client.Do(req)
resp, err := servertiming.Do(entry.client, req)
if err != nil {
atomic.AddInt64(&entry.inFlight, -1)
atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
@@ -1101,6 +1110,11 @@ func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMo
switch protocolMode {
case upstreamProtocolModeOpenAIH2:
transport.ForceAttemptHTTP2 = true
// 显式配置 http2 并启用 PING 健康探测,剔除代理/NAT 静默掐断的死连接,
// 避免请求挂在死连接上直到 TCP 重传超时(分钟级)。
if _, err := enableOpenAIHTTP2KeepAlive(transport); err != nil {
return nil, err
}
case upstreamProtocolModeOpenAIH1:
transport.ForceAttemptHTTP2 = false
transport.TLSNextProto = make(map[string]func(string, *tls.Conn) http.RoundTripper)
@@ -1115,6 +1129,22 @@ func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMo
return transport, nil
}
// enableOpenAIHTTP2KeepAlive 在 http.Transport 上显式配置 HTTP/2 并启用连接健康探测。
// Go 默认惰性配置 http2 且 ReadIdleTimeout=0(不发健康 PING),无法检测被代理/NAT
// 静默掐断的死连接。此处主动设置 ReadIdleTimeout/PingTimeout,让死连接被提前 PING
// 出并关闭,请求得以重建连接而非挂到 TCP 重传超时。返回底层 *http2.Transport 便于测试。
func enableOpenAIHTTP2KeepAlive(transport *http.Transport) (*http2.Transport, error) {
h2, err := http2.ConfigureTransports(transport)
if err != nil {
return nil, err
}
if h2 != nil {
h2.ReadIdleTimeout = openAIHTTP2ReadIdleTimeout
h2.PingTimeout = openAIHTTP2PingTimeout
}
return h2, nil
}
// buildUpstreamTransportWithTLSFingerprint 构建带 TLS 指纹伪装的 Transport
// 使用 utls 库模拟 Claude CLI 的 TLS 指纹
//
@@ -0,0 +1,67 @@
package repository
import (
"net/http"
"net/url"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func http2KeepAliveTestPoolSettings() poolSettings {
return poolSettings{
maxIdleConns: 10,
maxIdleConnsPerHost: 5,
maxConnsPerHost: 10,
idleConnTimeout: 90 * time.Second,
responseHeaderTimeout: time.Minute,
}
}
// Codex/OpenAI 上游改走 HTTP/2 后,池化连接被代理/NAT 静默掐断会成为“死连接”:
// 两端都以为连接存活,请求落上去会挂到 TCP 重传超时(分钟级)才失败。Go 的
// http2.Transport 默认 ReadIdleTimeout=0(不发健康 PING),无法检测这种死连接。
// 必须显式启用主动 PING 探测,让死连接被提前剔除,而不是只靠 ResponseHeaderTimeout
// 事后兜底。
func TestEnableOpenAIHTTP2KeepAlive_EnablesPingHealthCheck(t *testing.T) {
tr := &http.Transport{}
h2, err := enableOpenAIHTTP2KeepAlive(tr)
require.NoError(t, err)
require.NotNil(t, h2, "必须返回已配置的 *http2.Transport")
require.Positive(t, h2.ReadIdleTimeout, "必须启用空闲 PING 探测以剔除死连接")
require.Equal(t, openAIHTTP2ReadIdleTimeout, h2.ReadIdleTimeout)
require.Equal(t, openAIHTTP2PingTimeout, h2.PingTimeout, "PING 无响应必须有超时判定")
require.NotNil(t, tr.TLSNextProto["h2"], "http2 必须已挂到底层 http.Transport 上")
}
// openai_h2 模式构建的 Transport 必须带上 H2 PING 健康探测,从源头剔除死连接。
func TestBuildUpstreamTransport_OpenAIH2_EnablesPingHealthCheck(t *testing.T) {
tr, err := buildUpstreamTransport(http2KeepAliveTestPoolSettings(), nil, upstreamProtocolModeOpenAIH2)
require.NoError(t, err)
require.True(t, tr.ForceAttemptHTTP2, "openai_h2 必须启用 HTTP/2")
require.NotNil(t, tr.TLSNextProto["h2"], "openai_h2 必须显式配置 http2 以启用 ReadIdleTimeout")
}
// 非 H2 模式(default/h1)不应因本次改动被误配置:default 走 Go 自动 H2(惰性配置,
// 构建时 TLSNextProto 仍为空),h1 模式显式禁用 H2。避免波及 Claude/Gemini 热路径。
func TestBuildUpstreamTransport_NonOpenAIH2_NotEagerlyConfigured(t *testing.T) {
tr, err := buildUpstreamTransport(http2KeepAliveTestPoolSettings(), nil, upstreamProtocolModeDefault)
require.NoError(t, err)
require.Nil(t, tr.TLSNextProto["h2"], "default 模式不应在构建期主动配置 http2 keepalive")
}
// 死连接在经 HTTP 代理(CONNECT 隧道)时最高发,这是带 proxy 账号的真实生产路径:
// 显式 http2 配置须与 Transport.Proxy 同时正确生效,不能相互干扰。
func TestBuildUpstreamTransport_OpenAIH2_WithHTTPProxy_EnablesKeepAlive(t *testing.T) {
proxyURL, err := url.Parse("http://127.0.0.1:8080")
require.NoError(t, err)
tr, err := buildUpstreamTransport(http2KeepAliveTestPoolSettings(), proxyURL, upstreamProtocolModeOpenAIH2)
require.NoError(t, err)
require.True(t, tr.ForceAttemptHTTP2)
require.NotNil(t, tr.TLSNextProto["h2"], "经代理的 openai_h2 也必须启用 http2 keepalive")
require.NotNil(t, tr.Proxy, "HTTP 代理仍须通过 Transport.Proxy 生效")
}
@@ -55,6 +55,8 @@ const paymentOrdersOutTradeNoUniqueMigration = "120_enforce_payment_orders_out_t
const paymentOrdersOutTradeNoUniqueIndex = "paymentorder_out_trade_no_unique"
const schedulerOutboxPendingDedupKeyMigration = "153_scheduler_outbox_pending_dedup_key_index_notx.sql"
const schedulerOutboxPendingDedupKeyIndex = "idx_scheduler_outbox_pending_dedup_key"
const latestAPIKeyIPIndexMigration = "174_add_usage_logs_api_key_latest_ip_index_notx.sql"
const latestAPIKeyIPIndex = "idx_usage_logs_api_key_latest_ip"
type migrationChecksumCompatibilityRule struct {
fileChecksum string
@@ -264,6 +266,8 @@ func prepareNonTransactionalMigration(ctx context.Context, db *sql.DB, name stri
return preparePaymentOrdersOutTradeNoUniqueMigration(ctx, db)
case schedulerOutboxPendingDedupKeyMigration:
return dropInvalidIndexIfPresent(ctx, db, schedulerOutboxPendingDedupKeyIndex)
case latestAPIKeyIPIndexMigration:
return dropInvalidIndexIfPresent(ctx, db, latestAPIKeyIPIndex)
default:
return nil
}
@@ -116,6 +116,45 @@ CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_t_b ON t(b);
require.NoError(t, mock.ExpectationsWereMet())
}
func TestApplyMigrationsFS_NonTransactionalMigration_LatestAPIKeyIPIndexDropsInvalidIndexBeforeRetry(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
defer func() { _ = db.Close() }()
prepareMigrationsBootstrapExpectations(mock)
mock.ExpectQuery("SELECT checksum FROM schema_migrations WHERE filename = \\$1").
WithArgs(latestAPIKeyIPIndexMigration).
WillReturnError(sql.ErrNoRows)
mock.ExpectQuery("SELECT EXISTS \\(").
WithArgs(latestAPIKeyIPIndex).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
mock.ExpectExec("DROP INDEX CONCURRENTLY IF EXISTS idx_usage_logs_api_key_latest_ip").
WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectExec("CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_api_key_latest_ip").
WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectExec("INSERT INTO schema_migrations \\(filename, checksum\\) VALUES \\(\\$1, \\$2\\)").
WithArgs(latestAPIKeyIPIndexMigration, sqlmock.AnyArg()).
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectExec("SELECT pg_advisory_unlock\\(\\$1\\)").
WithArgs(migrationsAdvisoryLockID).
WillReturnResult(sqlmock.NewResult(0, 1))
fsys := fstest.MapFS{
latestAPIKeyIPIndexMigration: &fstest.MapFile{
Data: []byte(`
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_api_key_latest_ip
ON usage_logs (api_key_id, created_at DESC, id DESC)
INCLUDE (ip_address)
WHERE ip_address IS NOT NULL AND ip_address <> '';
`),
},
}
err = applyMigrationsFS(context.Background(), db, fsys)
require.NoError(t, err)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestApplyMigrationsFS_PaymentOrdersOutTradeNoUniqueMigration_FailsFastOnDuplicatePrecheck(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
@@ -0,0 +1,159 @@
//go:build integration
package repository
import (
"context"
"testing"
dbmigrations "github.com/Wei-Shaw/sub2api/migrations"
"github.com/stretchr/testify/require"
)
func TestMigration175EnforcesOpenAILongContextBillingWriteInvariant(t *testing.T) {
tx := testTx(t)
ctx := context.Background()
migrationSQL, err := dbmigrations.FS.ReadFile("175_default_openai_long_context_billing.sql")
require.NoError(t, err)
_, err = tx.ExecContext(ctx, `
DROP TRIGGER IF EXISTS accounts_propagate_openai_long_context_billing_extra ON accounts;
DROP TRIGGER IF EXISTS accounts_enforce_openai_long_context_billing_extra ON accounts;
`)
require.NoError(t, err)
var ordinaryID int64
require.NoError(t, tx.QueryRowContext(ctx, `
INSERT INTO accounts (name, platform, type, extra)
VALUES ('migration-175-ordinary', 'openai', 'oauth', '{}'::jsonb)
RETURNING id
`).Scan(&ordinaryID))
var parentID int64
require.NoError(t, tx.QueryRowContext(ctx, `
INSERT INTO accounts (name, platform, type, extra)
VALUES ('migration-175-parent', 'openai', 'oauth', '{"openai_long_context_billing_enabled":false}'::jsonb)
RETURNING id
`).Scan(&parentID))
var shadowID int64
require.NoError(t, tx.QueryRowContext(ctx, `
INSERT INTO accounts (name, platform, type, extra, parent_account_id, quota_dimension)
VALUES ('migration-175-shadow', 'openai', 'oauth', '{}'::jsonb, $1, 'spark')
RETURNING id
`, parentID).Scan(&shadowID))
var malformedLegacyID int64
require.NoError(t, tx.QueryRowContext(ctx, `
INSERT INTO accounts (name, platform, type, extra)
VALUES ('migration-175-malformed-legacy', 'openai', 'oauth', '{"openai_long_context_billing_enabled":"false"}'::jsonb)
RETURNING id
`).Scan(&malformedLegacyID))
_, err = tx.ExecContext(ctx, string(migrationSQL))
require.NoError(t, err)
_, err = tx.ExecContext(ctx, string(migrationSQL))
require.NoError(t, err)
var ordinaryEnabled bool
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT (extra->>'openai_long_context_billing_enabled')::boolean
FROM accounts
WHERE id = $1
`, ordinaryID).Scan(&ordinaryEnabled))
require.False(t, ordinaryEnabled)
var shadowEnabled bool
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT (extra->>'openai_long_context_billing_enabled')::boolean
FROM accounts
WHERE id = $1
`, shadowID).Scan(&shadowEnabled))
require.False(t, shadowEnabled)
var initialShadowOutboxEvents int
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM scheduler_outbox
WHERE event_type = 'account_changed' AND account_id = $1
`, shadowID).Scan(&initialShadowOutboxEvents))
require.Equal(t, 1, initialShadowOutboxEvents)
var malformedLegacyEnabled bool
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT (extra->>'openai_long_context_billing_enabled')::boolean
FROM accounts
WHERE id = $1
`, malformedLegacyID).Scan(&malformedLegacyEnabled))
require.False(t, malformedLegacyEnabled)
_, err = tx.ExecContext(ctx, `
UPDATE accounts
SET extra = extra || '{"migration_175_unrelated_update":true}'::jsonb
WHERE id = $1
`, malformedLegacyID)
require.NoError(t, err)
_, err = tx.ExecContext(ctx, "TRUNCATE scheduler_outbox")
require.NoError(t, err)
_, err = tx.ExecContext(ctx, `
UPDATE accounts
SET extra = '{"legacy_writer_replaced_extra":true}'::jsonb
WHERE id = $1
`, parentID)
require.NoError(t, err)
var parentEnabled bool
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT (extra->>'openai_long_context_billing_enabled')::boolean
FROM accounts
WHERE id = $1
`, parentID).Scan(&parentEnabled))
require.False(t, parentEnabled)
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT (extra->>'openai_long_context_billing_enabled')::boolean
FROM accounts
WHERE id = $1
`, shadowID).Scan(&shadowEnabled))
require.False(t, shadowEnabled)
var preservedOptOutEvents int
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM scheduler_outbox
WHERE event_type = 'account_changed' AND account_id = $1
`, shadowID).Scan(&preservedOptOutEvents))
require.Zero(t, preservedOptOutEvents)
require.NoError(t, tx.QueryRowContext(ctx, `
INSERT INTO accounts (name, platform, type, extra)
VALUES ('migration-175-rolling-writer', 'openai', 'oauth', '{}'::jsonb)
RETURNING (extra->>'openai_long_context_billing_enabled')::boolean
`).Scan(&ordinaryEnabled))
require.False(t, ordinaryEnabled)
_, err = tx.ExecContext(ctx, "TRUNCATE scheduler_outbox")
require.NoError(t, err)
_, err = tx.ExecContext(ctx, `
UPDATE accounts
SET extra = jsonb_set(extra, '{openai_long_context_billing_enabled}', 'true'::jsonb, true)
WHERE id = $1
`, parentID)
require.NoError(t, err)
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT (extra->>'openai_long_context_billing_enabled')::boolean
FROM accounts
WHERE id = $1
`, shadowID).Scan(&shadowEnabled))
require.True(t, shadowEnabled)
var shadowOutboxEvents int
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM scheduler_outbox
WHERE event_type = 'account_changed' AND account_id = $1
`, shadowID).Scan(&shadowOutboxEvents))
require.Equal(t, 1, shadowOutboxEvents)
_, err = tx.ExecContext(ctx, `
INSERT INTO accounts (name, platform, type, extra)
VALUES ('migration-175-malformed', 'openai', 'oauth', '{"openai_long_context_billing_enabled":"false"}'::jsonb)
`)
require.ErrorContains(t, err, "openai_long_context_billing_enabled must be a boolean")
}
+10
View File
@@ -718,6 +718,7 @@ func (r *opsRepository) BatchInsertSystemLogs(ctx context.Context, inputs []*ser
stmt, err := tx.PrepareContext(ctx, pq.CopyIn(
"ops_system_logs",
"created_at",
"host",
"level",
"component",
"message",
@@ -760,6 +761,7 @@ func (r *opsRepository) BatchInsertSystemLogs(ctx context.Context, inputs []*ser
if _, err := stmt.ExecContext(
ctx,
createdAt.UTC(),
opsNullString(input.Host),
level,
component,
message,
@@ -827,6 +829,7 @@ func (r *opsRepository) ListSystemLogs(ctx context.Context, filter *service.OpsS
SELECT
l.id,
l.created_at,
COALESCE(l.host, ''),
l.level,
COALESCE(l.component, ''),
COALESCE(l.message, ''),
@@ -859,6 +862,7 @@ LIMIT $` + itoa(len(args)+1) + ` OFFSET $` + itoa(len(args)+2)
if err := rows.Scan(
&item.ID,
&item.CreatedAt,
&item.Host,
&item.Level,
&item.Component,
&item.Message,
@@ -1130,6 +1134,11 @@ func buildOpsSystemLogsWhere(filter *service.OpsSystemLogFilter) (string, []any,
hasConstraint = true
}
if filter != nil {
if v := strings.TrimSpace(filter.Host); v != "" {
args = append(args, v)
clauses = append(clauses, "l.host = $"+itoa(len(args)))
hasConstraint = true
}
if v := strings.ToLower(strings.TrimSpace(filter.Level)); v != "" {
args = append(args, v)
clauses = append(clauses, "LOWER(COALESCE(l.level,'')) = $"+itoa(len(args)))
@@ -1194,6 +1203,7 @@ func buildOpsSystemLogsCleanupWhere(filter *service.OpsSystemLogCleanupFilter) (
listFilter := &service.OpsSystemLogFilter{
StartTime: filter.StartTime,
EndTime: filter.EndTime,
Host: filter.Host,
Level: filter.Level,
Component: filter.Component,
RequestID: filter.RequestID,
@@ -18,6 +18,7 @@ func TestBuildOpsSystemLogsWhere_WithClientRequestIDAndUserID(t *testing.T) {
filter := &service.OpsSystemLogFilter{
StartTime: &start,
EndTime: &end,
Host: "api-node-1",
Level: "warn",
Component: "http.access",
RequestID: "req-1",
@@ -37,8 +38,11 @@ func TestBuildOpsSystemLogsWhere_WithClientRequestIDAndUserID(t *testing.T) {
if where == "" {
t.Fatalf("where should not be empty")
}
if len(args) != 12 {
t.Fatalf("args len = %d, want 12", len(args))
if len(args) != 13 {
t.Fatalf("args len = %d, want 13", len(args))
}
if !contains(where, "l.host = $") {
t.Fatalf("where should include host condition: %s", where)
}
if !contains(where, "COALESCE(l.client_request_id,'') = $") {
t.Fatalf("where should include client_request_id condition: %s", where)
@@ -68,6 +72,7 @@ func TestBuildOpsSystemLogsCleanupWhere_WithClientRequestIDAndUserID(t *testing.
userID := int64(9)
apiKeyID := int64(10)
filter := &service.OpsSystemLogCleanupFilter{
Host: "api-node-2",
ClientRequestID: "creq-9",
UserID: &userID,
APIKeyID: &apiKeyID,
@@ -77,8 +82,11 @@ func TestBuildOpsSystemLogsCleanupWhere_WithClientRequestIDAndUserID(t *testing.
if !hasConstraint {
t.Fatalf("expected hasConstraint=true")
}
if len(args) != 3 {
t.Fatalf("args len = %d, want 3", len(args))
if len(args) != 4 {
t.Fatalf("args len = %d, want 4", len(args))
}
if !contains(where, "l.host = $") {
t.Fatalf("where should include host condition: %s", where)
}
if !contains(where, "COALESCE(l.client_request_id,'') = $") {
t.Fatalf("where should include client_request_id condition: %s", where)

Some files were not shown because too many files have changed in this diff Show More