mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-21 14:19:18 +08:00
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:
@@ -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:
|
||||
|
||||
@@ -116,6 +116,8 @@ backend/.installed
|
||||
# 其他
|
||||
# ===================
|
||||
tests
|
||||
!deploy/tests/
|
||||
!deploy/tests/**
|
||||
CLAUDE.md
|
||||
.claude
|
||||
scripts
|
||||
|
||||
@@ -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
@@ -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 container(macOS)
|
||||
|
||||
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
@@ -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 container(macOS)
|
||||
|
||||
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 @@
|
||||
0.1.152
|
||||
0.1.153
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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
@@ -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))
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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=true(Antigravity 粘性会话切换)
|
||||
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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -164,7 +164,7 @@ var ProviderSet = wire.NewSet(
|
||||
admin.NewDashboardHandler,
|
||||
admin.NewUserHandler,
|
||||
admin.NewGroupHandler,
|
||||
admin.NewAccountHandler,
|
||||
admin.ProvideAccountHandler,
|
||||
admin.NewAnnouncementHandler,
|
||||
admin.NewDataManagementHandler,
|
||||
admin.NewBackupHandler,
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 ©Client
|
||||
}
|
||||
|
||||
// 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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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 ""
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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},
|
||||
}))
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 后发 PING,PingTimeout
|
||||
// 内无响应即判定死连接并关闭,从源头避免请求挂在死连接上。
|
||||
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)
|
||||
|
||||
+159
@@ -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")
|
||||
}
|
||||
@@ -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
Reference in New Issue
Block a user