diff --git a/.github/workflows/backend-ci.yml b/.github/workflows/backend-ci.yml index bb5cf692bc..8599465254 100644 --- a/.github/workflows/backend-ci.yml +++ b/.github/workflows/backend-ci.yml @@ -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: diff --git a/.gitignore b/.gitignore index bd2e3e6ddf..d45d27f79e 100644 --- a/.gitignore +++ b/.gitignore @@ -116,6 +116,8 @@ backend/.installed # 其他 # =================== tests +!deploy/tests/ +!deploy/tests/** CLAUDE.md .claude scripts diff --git a/README.md b/README.md index 3e43fd3ae4..6bb4068644 100644 --- a/README.md +++ b/README.md @@ -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://: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 diff --git a/README_CN.md b/README_CN.md index 88c8ce11b1..348e1c93c4 100644 --- a/README_CN.md +++ b/README_CN.md @@ -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)。 + +--- + +### 方式四：源码编译 从源码编译安装，适合开发或定制需求。 diff --git a/README_JA.md b/README_JA.md index 21d070e397..ac18b36387 100644 --- a/README_JA.md +++ b/README_JA.md @@ -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: ソースからビルド 開発やカスタマイズのためにソースコードからビルドして実行します。 diff --git a/backend/cmd/server/VERSION b/backend/cmd/server/VERSION index 611234586a..c64f02bbc3 100644 --- a/backend/cmd/server/VERSION +++ b/backend/cmd/server/VERSION @@ -1 +1 @@ -0.1.152 +0.1.153 diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index e148c5c363..75786b9c70 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -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) diff --git a/backend/ent/channelmonitor/channelmonitor.go b/backend/ent/channelmonitor/channelmonitor.go index afdc6957d6..711e6217e2 100644 --- a/backend/ent/channelmonitor/channelmonitor.go +++ b/backend/ent/channelmonitor/channelmonitor.go @@ -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) diff --git a/backend/ent/channelmonitorrequesttemplate/channelmonitorrequesttemplate.go b/backend/ent/channelmonitorrequesttemplate/channelmonitorrequesttemplate.go index db04aee106..5989d0e743 100644 --- a/backend/ent/channelmonitorrequesttemplate/channelmonitorrequesttemplate.go +++ b/backend/ent/channelmonitorrequesttemplate/channelmonitorrequesttemplate.go @@ -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) diff --git a/backend/ent/migrate/schema.go b/backend/ent/migrate/schema.go index 3441afec04..52229f9151 100644 --- a/backend/ent/migrate/schema.go +++ b/backend/ent/migrate/schema.go @@ -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]}, }, }, } diff --git a/backend/ent/mutation.go b/backend/ent/mutation.go index ab7c424a47..fb35531878 100644 --- a/backend/ent/mutation.go +++ b/backend/ent/mutation.go @@ -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 diff --git a/backend/ent/runtime/runtime.go b/backend/ent/runtime/runtime.go index 4cb3f800f8..867f1cbdbd 100644 --- a/backend/ent/runtime/runtime.go +++ b/backend/ent/runtime/runtime.go @@ -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() diff --git a/backend/ent/schema/channel_monitor.go b/backend/ent/schema/channel_monitor.go index d9594ab39c..cb62079316 100644 --- a/backend/ent/schema/channel_monitor.go +++ b/backend/ent/schema/channel_monitor.go @@ -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). diff --git a/backend/ent/schema/channel_monitor_request_template.go b/backend/ent/schema/channel_monitor_request_template.go index 0e0ce3a0b5..cf7fe05158 100644 --- a/backend/ent/schema/channel_monitor_request_template.go +++ b/backend/ent/schema/channel_monitor_request_template.go @@ -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). diff --git a/backend/ent/schema/usage_log.go b/backend/ent/schema/usage_log.go index e84cc1c140..6d8c2d4191 100644 --- a/backend/ent/schema/usage_log.go +++ b/backend/ent/schema/usage_log.go @@ -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"). diff --git a/backend/ent/usagelog.go b/backend/ent/usagelog.go index 4d374a8495..b13e29b2f7 100644 --- a/backend/ent/usagelog.go +++ b/backend/ent/usagelog.go @@ -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)) diff --git a/backend/ent/usagelog/usagelog.go b/backend/ent/usagelog/usagelog.go index a74a92c40f..a87d937195 100644 --- a/backend/ent/usagelog/usagelog.go +++ b/backend/ent/usagelog/usagelog.go @@ -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() diff --git a/backend/ent/usagelog/where.go b/backend/ent/usagelog/where.go index 4b08cc3425..a9462e0d0e 100644 --- a/backend/ent/usagelog/where.go +++ b/backend/ent/usagelog/where.go @@ -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)) diff --git a/backend/ent/usagelog_create.go b/backend/ent/usagelog_create.go index 3326f72fc0..31cf45328e 100644 --- a/backend/ent/usagelog_create.go +++ b/backend/ent/usagelog_create.go @@ -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) { diff --git a/backend/ent/usagelog_update.go b/backend/ent/usagelog_update.go index 00a65ccff1..2a60d6f44d 100644 --- a/backend/ent/usagelog_update.go +++ b/backend/ent/usagelog_update.go @@ -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) } diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index df3afb6c7e..1262845cea 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -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") } diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index 32aff543af..2f9defb80e 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -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") } diff --git a/backend/internal/handler/admin/account_codex_import.go b/backend/internal/handler/admin/account_codex_import.go index 01a5fbfa1c..271bd19470 100644 --- a/backend/internal/handler/admin/account_codex_import.go +++ b/backend/internal/handler/admin/account_codex_import.go @@ -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 diff --git a/backend/internal/handler/admin/account_codex_import_test.go b/backend/internal/handler/admin/account_codex_import_test.go index a52463aa86..96a033d8c3 100644 --- a/backend/internal/handler/admin/account_codex_import_test.go +++ b/backend/internal/handler/admin/account_codex_import_test.go @@ -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) { diff --git a/backend/internal/handler/admin/account_data.go b/backend/internal/handler/admin/account_data.go index bf872c4826..e44d726fd6 100644 --- a/backend/internal/handler/admin/account_data.go +++ b/backend/internal/handler/admin/account_data.go @@ -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++ } diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index a4b0773999..e4ed5b46b0 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -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, diff --git a/backend/internal/handler/admin/account_handler_long_context_billing_test.go b/backend/internal/handler/admin/account_handler_long_context_billing_test.go new file mode 100644 index 0000000000..d50513a3e8 --- /dev/null +++ b/backend/internal/handler/admin/account_handler_long_context_billing_test.go @@ -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) +} diff --git a/backend/internal/handler/admin/admin_service_stub_test.go b/backend/internal/handler/admin/admin_service_stub_test.go index 7a7cbb473e..5e9c4d517e 100644 --- a/backend/internal/handler/admin/admin_service_stub_test.go +++ b/backend/internal/handler/admin/admin_service_stub_test.go @@ -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 } diff --git a/backend/internal/handler/admin/channel_monitor_handler.go b/backend/internal/handler/admin/channel_monitor_handler.go index 4ef774e9e7..a69b835849 100644 --- a/backend/internal/handler/admin/channel_monitor_handler.go +++ b/backend/internal/handler/admin/channel_monitor_handler.go @@ -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"` diff --git a/backend/internal/handler/admin/channel_monitor_template_handler.go b/backend/internal/handler/admin/channel_monitor_template_handler.go index c842f465c8..497e3d195b 100644 --- a/backend/internal/handler/admin/channel_monitor_template_handler.go +++ b/backend/internal/handler/admin/channel_monitor_template_handler.go @@ -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"` diff --git a/backend/internal/handler/admin/grok_import_probe.go b/backend/internal/handler/admin/grok_import_probe.go new file mode 100644 index 0000000000..f1df15bba9 --- /dev/null +++ b/backend/internal/handler/admin/grok_import_probe.go @@ -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 +} diff --git a/backend/internal/handler/admin/grok_import_probe_handler_test.go b/backend/internal/handler/admin/grok_import_probe_handler_test.go new file mode 100644 index 0000000000..489a13ae6d --- /dev/null +++ b/backend/internal/handler/admin/grok_import_probe_handler_test.go @@ -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) +} diff --git a/backend/internal/handler/admin/grok_import_probe_test.go b/backend/internal/handler/admin/grok_import_probe_test.go new file mode 100644 index 0000000000..3b8fc0ca6e --- /dev/null +++ b/backend/internal/handler/admin/grok_import_probe_test.go @@ -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") +} diff --git a/backend/internal/handler/admin/grok_oauth_handler.go b/backend/internal/handler/admin/grok_oauth_handler.go index dafa3076b8..1a309b7c9a 100644 --- a/backend/internal/handler/admin/grok_oauth_handler.go +++ b/backend/internal/handler/admin/grok_oauth_handler.go @@ -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 diff --git a/backend/internal/handler/admin/grok_oauth_handler_test.go b/backend/internal/handler/admin/grok_oauth_handler_test.go index 0b0f0d1aba..64ea044aa3 100644 --- a/backend/internal/handler/admin/grok_oauth_handler_test.go +++ b/backend/internal/handler/admin/grok_oauth_handler_test.go @@ -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") +} diff --git a/backend/internal/handler/admin/openai_oauth_handler.go b/backend/internal/handler/admin/openai_oauth_handler.go index d7a756bd00..78d57299b6 100644 --- a/backend/internal/handler/admin/openai_oauth_handler.go +++ b/backend/internal/handler/admin/openai_oauth_handler.go @@ -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 diff --git a/backend/internal/handler/admin/ops_system_log_handler.go b/backend/internal/handler/admin/ops_system_log_handler.go index 9f3c8b893a..1b6af45976 100644 --- a/backend/internal/handler/admin/ops_system_log_handler.go +++ b/backend/internal/handler/admin/ops_system_log_handler.go @@ -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), diff --git a/backend/internal/handler/admin/ops_system_log_handler_test.go b/backend/internal/handler/admin/ops_system_log_handler_test.go index 9557fce442..3390fbe3cb 100644 --- a/backend/internal/handler/admin/ops_system_log_handler_test.go +++ b/backend/internal/handler/admin/ops_system_log_handler_test.go @@ -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) diff --git a/backend/internal/handler/admin/ops_ws_handler.go b/backend/internal/handler/admin/ops_ws_handler.go index 75fd7ea002..e4c42cc9c0 100644 --- a/backend/internal/handler/admin/ops_ws_handler.go +++ b/backend/internal/handler/admin/ops_ws_handler.go @@ -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 diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index e770bcf036..3c45c3b95e 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -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), } } diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index 0418e5bc3e..619926c1e4 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -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"` diff --git a/backend/internal/handler/endpoint.go b/backend/internal/handler/endpoint.go index 2871eef6ce..5e4d84ba72 100644 --- a/backend/internal/handler/endpoint.go +++ b/backend/internal/handler/endpoint.go @@ -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. diff --git a/backend/internal/handler/failover_loop.go b/backend/internal/handler/failover_loop.go index 6d8ddc7236..5838e58f48 100644 --- a/backend/internal/handler/failover_loop.go +++ b/backend/internal/handler/failover_loop.go @@ -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 diff --git a/backend/internal/handler/failover_loop_test.go b/backend/internal/handler/failover_loop_test.go index 2c65ebc2c8..9fabe75f25 100644 --- a/backend/internal/handler/failover_loop_test.go +++ b/backend/internal/handler/failover_loop_test.go @@ -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) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 116346b4a6..fb07b098d1 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -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, } } diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index f3805f3a53..af9bcdb344 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -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 diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index 5b49ca69a2..8a88d5fa57 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -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 diff --git a/backend/internal/handler/gateway_handler_usage_test.go b/backend/internal/handler/gateway_handler_usage_test.go new file mode 100644 index 0000000000..b6b0e0efe6 --- /dev/null +++ b/backend/internal/handler/gateway_handler_usage_test.go @@ -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)) +} diff --git a/backend/internal/handler/gateway_helper.go b/backend/internal/handler/gateway_helper.go index 48110da93f..8489b17f6d 100644 --- a/backend/internal/handler/gateway_helper.go +++ b/backend/internal/handler/gateway_helper.go @@ -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) { diff --git a/backend/internal/handler/gateway_helper_fastpath_test.go b/backend/internal/handler/gateway_helper_fastpath_test.go index fecb9b071d..7ae9cb513e 100644 --- a/backend/internal/handler/gateway_helper_fastpath_test.go +++ b/backend/internal/handler/gateway_helper_fastpath_test.go @@ -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) { diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 86be1062e9..b1653c6b3f 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -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 diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 4fd1411b23..b7092fcc42 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -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), diff --git a/backend/internal/handler/openai_codex_models_handler.go b/backend/internal/handler/openai_codex_models_handler.go index e64c555d14..1c1357cbfa 100644 --- a/backend/internal/handler/openai_codex_models_handler.go +++ b/backend/internal/handler/openai_codex_models_handler.go @@ -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) } diff --git a/backend/internal/handler/openai_codex_models_handler_test.go b/backend/internal/handler/openai_codex_models_handler_test.go new file mode 100644 index 0000000000..ba74a5869f --- /dev/null +++ b/backend/internal/handler/openai_codex_models_handler_test.go @@ -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 +} diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index afa2a5073a..231363a8a3 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -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 } diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index b7f43079ef..e4b594c0b8 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -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) diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 5868f7f35b..c5982fb7d1 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -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") } diff --git a/backend/internal/handler/payment_handler.go b/backend/internal/handler/payment_handler.go index a267d73724..1ad054da75 100644 --- a/backend/internal/handler/payment_handler.go +++ b/backend/internal/handler/payment_handler.go @@ -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 diff --git a/backend/internal/handler/payment_handler_resume_test.go b/backend/internal/handler/payment_handler_resume_test.go index 21fd8ad763..c902f390e0 100644 --- a/backend/internal/handler/payment_handler_resume_test.go +++ b/backend/internal/handler/payment_handler_resume_test.go @@ -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) diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index cfbb72554c..67380fed33 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -164,7 +164,7 @@ var ProviderSet = wire.NewSet( admin.NewDashboardHandler, admin.NewUserHandler, admin.NewGroupHandler, - admin.NewAccountHandler, + admin.ProvideAccountHandler, admin.NewAnnouncementHandler, admin.NewDataManagementHandler, admin.NewBackupHandler, diff --git a/backend/internal/pkg/antigravity/client.go b/backend/internal/pkg/antigravity/client.go index e318d1cdaf..39b6d2c90c 100644 --- a/backend/internal/pkg/antigravity/client.go +++ b/backend/internal/pkg/antigravity/client.go @@ -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) } diff --git a/backend/internal/pkg/apicompat/anthropic_responses_test.go b/backend/internal/pkg/apicompat/anthropic_responses_test.go index 8997835c2a..db6b49aa9b 100644 --- a/backend/internal/pkg/apicompat/anthropic_responses_test.go +++ b/backend/internal/pkg/apicompat/anthropic_responses_test.go @@ -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) { diff --git a/backend/internal/pkg/apicompat/anthropic_to_responses_response.go b/backend/internal/pkg/apicompat/anthropic_to_responses_response.go index 67c161bdd3..661b47cebe 100644 --- a/backend/internal/pkg/apicompat/anthropic_to_responses_response.go +++ b/backend/internal/pkg/apicompat/anthropic_to_responses_response.go @@ -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, }, diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go index 23178bf3a0..8aa9eab60a 100644 --- a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go @@ -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 { diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge_custom_tools_test.go b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge_custom_tools_test.go index 9271b50d1d..5b1d994eb3 100644 --- a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge_custom_tools_test.go +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge_custom_tools_test.go @@ -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", diff --git a/backend/internal/pkg/apicompat/responses_to_anthropic.go b/backend/internal/pkg/apicompat/responses_to_anthropic.go index 9c3b85b2ef..376f0d97da 100644 --- a/backend/internal/pkg/apicompat/responses_to_anthropic.go +++ b/backend/internal/pkg/apicompat/responses_to_anthropic.go @@ -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 } diff --git a/backend/internal/pkg/apicompat/responses_to_anthropic_read_tool_test.go b/backend/internal/pkg/apicompat/responses_to_anthropic_read_tool_test.go new file mode 100644 index 0000000000..72b60099fe --- /dev/null +++ b/backend/internal/pkg/apicompat/responses_to_anthropic_read_tool_test.go @@ -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) +} diff --git a/backend/internal/pkg/apicompat/responses_to_chatcompletions.go b/backend/internal/pkg/apicompat/responses_to_chatcompletions.go index 2ae6f8ac3f..a89a1b4203 100644 --- a/backend/internal/pkg/apicompat/responses_to_chatcompletions.go +++ b/backend/internal/pkg/apicompat/responses_to_chatcompletions.go @@ -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 { diff --git a/backend/internal/pkg/apicompat/streaming_stop_reason_test.go b/backend/internal/pkg/apicompat/streaming_stop_reason_test.go new file mode 100644 index 0000000000..c2889f0251 --- /dev/null +++ b/backend/internal/pkg/apicompat/streaming_stop_reason_test.go @@ -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") +} diff --git a/backend/internal/pkg/httpclient/pool.go b/backend/internal/pkg/httpclient/pool.go index 12804cc67d..22d3c65feb 100644 --- a/backend/internal/pkg/httpclient/pool.go +++ b/backend/internal/pkg/httpclient/pool.go @@ -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, diff --git a/backend/internal/pkg/pagination/pagination.go b/backend/internal/pkg/pagination/pagination.go index ce8e74b8ce..334ba809de 100644 --- a/backend/internal/pkg/pagination/pagination.go +++ b/backend/internal/pkg/pagination/pagination.go @@ -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 获取限制数 diff --git a/backend/internal/pkg/pagination/pagination_test.go b/backend/internal/pkg/pagination/pagination_test.go index 9a3b069d90..9704449e92 100644 --- a/backend/internal/pkg/pagination/pagination_test.go +++ b/backend/internal/pkg/pagination/pagination_test.go @@ -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) + } + }) + } +} diff --git a/backend/internal/pkg/servertiming/collector.go b/backend/internal/pkg/servertiming/collector.go new file mode 100644 index 0000000000..553edede31 --- /dev/null +++ b/backend/internal/pkg/servertiming/collector.go @@ -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) +} diff --git a/backend/internal/pkg/servertiming/collector_test.go b/backend/internal/pkg/servertiming/collector_test.go new file mode 100644 index 0000000000..1bb809f9bc --- /dev/null +++ b/backend/internal/pkg/servertiming/collector_test.go @@ -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) + } +} diff --git a/backend/internal/pkg/servertiming/http.go b/backend/internal/pkg/servertiming/http.go new file mode 100644 index 0000000000..e326e24302 --- /dev/null +++ b/backend/internal/pkg/servertiming/http.go @@ -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" + } +} diff --git a/backend/internal/pkg/servertiming/http_test.go b/backend/internal/pkg/servertiming/http_test.go new file mode 100644 index 0000000000..d37f378414 --- /dev/null +++ b/backend/internal/pkg/servertiming/http_test.go @@ -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") + } +} diff --git a/backend/internal/pkg/xai/billing.go b/backend/internal/pkg/xai/billing.go new file mode 100644 index 0000000000..15b9c7e50e --- /dev/null +++ b/backend/internal/pkg/xai/billing.go @@ -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 +} diff --git a/backend/internal/pkg/xai/billing_test.go b/backend/internal/pkg/xai/billing_test.go new file mode 100644 index 0000000000..1d863f6a39 --- /dev/null +++ b/backend/internal/pkg/xai/billing_test.go @@ -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 +} diff --git a/backend/internal/pkg/xai/oauth.go b/backend/internal/pkg/xai/oauth.go index 1b3aadc08b..a8a549c9dc 100644 --- a/backend/internal/pkg/xai/oauth.go +++ b/backend/internal/pkg/xai/oauth.go @@ -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 { diff --git a/backend/internal/pkg/xai/oauth_test.go b/backend/internal/pkg/xai/oauth_test.go index d3d3d5cb29..39200ffc39 100644 --- a/backend/internal/pkg/xai/oauth_test.go +++ b/backend/internal/pkg/xai/oauth_test.go @@ -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) { diff --git a/backend/internal/pkg/xai/sso_device.go b/backend/internal/pkg/xai/sso_device.go new file mode 100644 index 0000000000..e533e394d2 --- /dev/null +++ b/backend/internal/pkg/xai/sso_device.go @@ -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 "" +} diff --git a/backend/internal/pkg/xai/sso_device_test.go b/backend/internal/pkg/xai/sso_device_test.go new file mode 100644 index 0000000000..27dff15f4d --- /dev/null +++ b/backend/internal/pkg/xai/sso_device_test.go @@ -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, `ok`), 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, `consent`), 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, `done`), 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 +} diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 261be7c14b..8eb819aeab 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -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 { diff --git a/backend/internal/repository/account_repo_auto_pause_test.go b/backend/internal/repository/account_repo_auto_pause_test.go new file mode 100644 index 0000000000..0eb48a296d --- /dev/null +++ b/backend/internal/repository/account_repo_auto_pause_test.go @@ -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()) +} diff --git a/backend/internal/repository/account_repo_grok_billing_test.go b/backend/internal/repository/account_repo_grok_billing_test.go new file mode 100644 index 0000000000..fb41ae5ffa --- /dev/null +++ b/backend/internal/repository/account_repo_grok_billing_test.go @@ -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}, + })) +} diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index ee4ed4785f..4c0edf72b4 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -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)) diff --git a/backend/internal/repository/api_key_repo_last_used_unit_test.go b/backend/internal/repository/api_key_repo_last_used_unit_test.go index 839eda7f75..dbdf653f8a 100644 --- a/backend/internal/repository/api_key_repo_last_used_unit_test.go +++ b/backend/internal/repository/api_key_repo_last_used_unit_test.go @@ -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() diff --git a/backend/internal/repository/backup_s3_store.go b/backend/internal/repository/backup_s3_store.go index 5d419f574b..2104e1e5d7 100644 --- a/backend/internal/repository/backup_s3_store.go +++ b/backend/internal/repository/backup_s3_store.go @@ -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) } diff --git a/backend/internal/repository/claude_oauth_service.go b/backend/internal/repository/claude_oauth_service.go index 5c5f27c86a..ec2d426ecb 100644 --- a/backend/internal/repository/claude_oauth_service.go +++ b/backend/internal/repository/claude_oauth_service.go @@ -276,5 +276,5 @@ func createReqClient(proxyURL string) (*req.Client, error) { client.SetProxyURL(trimmed) } - return client, nil + return instrumentReqClient(client), nil } diff --git a/backend/internal/repository/concurrency_cache.go b/backend/internal/repository/concurrency_cache.go index b657c1ce8f..5341d411e9 100644 --- a/backend/internal/repository/concurrency_cache.go +++ b/backend/internal/repository/concurrency_cache.go @@ -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 diff --git a/backend/internal/repository/concurrency_cache_integration_test.go b/backend/internal/repository/concurrency_cache_integration_test.go index f7e27d1118..02821159ee 100644 --- a/backend/internal/repository/concurrency_cache_integration_test.go +++ b/backend/internal/repository/concurrency_cache_integration_test.go @@ -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" diff --git a/backend/internal/repository/ent.go b/backend/internal/repository/ent.go index 64d321924d..3abb528e98 100644 --- a/backend/internal/repository/ent.go +++ b/backend/internal/repository/ent.go @@ -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) diff --git a/backend/internal/repository/grok_oauth_client.go b/backend/internal/repository/grok_oauth_client.go index 435ced5a65..6c9c2c407f 100644 --- a/backend/internal/repository/grok_oauth_client.go +++ b/backend/internal/repository/grok_oauth_client.go @@ -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 diff --git a/backend/internal/repository/http_upstream.go b/backend/internal/repository/http_upstream.go index bb079b0789..58b1d345d4 100644 --- a/backend/internal/repository/http_upstream.go +++ b/backend/internal/repository/http_upstream.go @@ -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 指纹 // diff --git a/backend/internal/repository/http_upstream_http2_keepalive_test.go b/backend/internal/repository/http_upstream_http2_keepalive_test.go new file mode 100644 index 0000000000..ead61da3b2 --- /dev/null +++ b/backend/internal/repository/http_upstream_http2_keepalive_test.go @@ -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 生效") +} diff --git a/backend/internal/repository/migrations_runner.go b/backend/internal/repository/migrations_runner.go index 7c045fea74..a071967f65 100644 --- a/backend/internal/repository/migrations_runner.go +++ b/backend/internal/repository/migrations_runner.go @@ -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 } diff --git a/backend/internal/repository/migrations_runner_notx_test.go b/backend/internal/repository/migrations_runner_notx_test.go index c9f6a2cdf1..6bb7914b95 100644 --- a/backend/internal/repository/migrations_runner_notx_test.go +++ b/backend/internal/repository/migrations_runner_notx_test.go @@ -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) diff --git a/backend/internal/repository/openai_long_context_billing_migration_integration_test.go b/backend/internal/repository/openai_long_context_billing_migration_integration_test.go new file mode 100644 index 0000000000..5f50ed0729 --- /dev/null +++ b/backend/internal/repository/openai_long_context_billing_migration_integration_test.go @@ -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") +} diff --git a/backend/internal/repository/ops_repo.go b/backend/internal/repository/ops_repo.go index 2129a451c4..900abcf212 100644 --- a/backend/internal/repository/ops_repo.go +++ b/backend/internal/repository/ops_repo.go @@ -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, diff --git a/backend/internal/repository/ops_repo_system_logs_test.go b/backend/internal/repository/ops_repo_system_logs_test.go index 98199f4828..48be3e7256 100644 --- a/backend/internal/repository/ops_repo_system_logs_test.go +++ b/backend/internal/repository/ops_repo_system_logs_test.go @@ -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) diff --git a/backend/internal/repository/proxy_expiry_integration_test.go b/backend/internal/repository/proxy_expiry_integration_test.go index d0cdc913b6..88f4f1df93 100644 --- a/backend/internal/repository/proxy_expiry_integration_test.go +++ b/backend/internal/repository/proxy_expiry_integration_test.go @@ -4,6 +4,7 @@ package repository import ( "context" + "encoding/json" "testing" "time" @@ -70,6 +71,41 @@ func (s *ProxyExpirySuite) TestSweep_DirectMode() { s.Require().Equal(pid, *origin) } +func (s *ProxyExpirySuite) TestSweep_EnqueuesChangedAccountIDsWithoutFullRebuild() { + past := time.Now().Add(-time.Hour) + firstProxyID := s.mkProxy("p-bulk-first", service.FallbackModeDirect, &past, nil) + secondProxyID := s.mkProxy("p-bulk-second", service.FallbackModeDirect, &past, nil) + firstAccountID := s.mkAccountWithProxy(firstProxyID) + secondAccountID := s.mkAccountWithProxy(secondProxyID) + + changed, err := s.repo.SweepExpiredProxies(s.ctx, time.Now()) + s.Require().NoError(err) + s.Require().EqualValues(2, changed) + + var payloadRaw []byte + err = scanSingleRow(s.ctx, s.tx, ` + SELECT payload + FROM scheduler_outbox + WHERE event_type=$1 + ORDER BY id DESC + LIMIT 1`, []any{service.SchedulerOutboxEventAccountBulkChanged}, &payloadRaw) + s.Require().NoError(err) + + var payload struct { + AccountIDs []int64 `json:"account_ids"` + } + s.Require().NoError(json.Unmarshal(payloadRaw, &payload)) + s.Require().Equal([]int64{firstAccountID, secondAccountID}, payload.AccountIDs) + + var fullRebuildCount int + err = scanSingleRow(s.ctx, s.tx, ` + SELECT COUNT(*) + FROM scheduler_outbox + WHERE event_type=$1`, []any{service.SchedulerOutboxEventFullRebuild}, &fullRebuildCount) + s.Require().NoError(err) + s.Require().Zero(fullRebuildCount) +} + func (s *ProxyExpirySuite) TestSweep_ProxyMode_Healthy() { future := time.Now().Add(24 * time.Hour) past := time.Now().Add(-time.Hour) diff --git a/backend/internal/repository/proxy_expiry_test.go b/backend/internal/repository/proxy_expiry_test.go new file mode 100644 index 0000000000..7313241085 --- /dev/null +++ b/backend/internal/repository/proxy_expiry_test.go @@ -0,0 +1,27 @@ +package repository + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSortedUniqueAccountIDs(t *testing.T) { + tests := []struct { + name string + input []int64 + want []int64 + }{ + {name: "unsorted duplicates", input: []int64{12, 3, 12, 8, 3}, want: []int64{3, 8, 12}}, + {name: "already sorted", input: []int64{3, 8, 12}, want: []int64{3, 8, 12}}, + {name: "single", input: []int64{3}, want: []int64{3}}, + {name: "empty", input: []int64{}, want: []int64{}}, + {name: "nil", input: nil, want: nil}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, sortedUniqueAccountIDs(tt.input)) + }) + } +} diff --git a/backend/internal/repository/proxy_repo.go b/backend/internal/repository/proxy_repo.go index b34c0cb559..fcb2e53b87 100644 --- a/backend/internal/repository/proxy_repo.go +++ b/backend/internal/repository/proxy_repo.go @@ -489,7 +489,7 @@ func (r *proxyRepository) ListAllForFallback(ctx context.Context) ([]service.Pro // SweepExpiredProxies 扫描到期 active 代理，标记 expired 并按 fallback 策略改写绑定账号的 proxy_id， // 最终触发 scheduler outbox 使 Redis 快照缓存失效。返回受影响的账号行数。 // 原子性边界：每个过期代理的「标记 expired + 改投账号」在各自子事务内原子执行（见 sweepOneExpiredProxy）； -// 全部代理处理完后若有账号被改投，再统一 enqueue 一次 full_rebuild 事件——该 enqueue 在子事务之外 +// 全部代理处理完后若有账号被改投，再统一 enqueue 一次 account_bulk_changed 事件——该 enqueue 在子事务之外 // （走 r.sql、失败仅记日志、由调度器周期性 full rebuild 兜底），故「改投 → 失效」整体并非原子。 func (r *proxyRepository) SweepExpiredProxies(ctx context.Context, now time.Time) (int64, error) { // 快照读（事务前）：允许脏读不影响正确性，事务内已加锁写。 @@ -503,7 +503,7 @@ func (r *proxyRepository) SweepExpiredProxies(ctx context.Context, now time.Time } var totalChanged int64 - accountsTouched := false + allChangedAccountIDs := make([]int64, 0) for _, p := range all { if p.Status != service.StatusActive || !p.IsExpired(now) { @@ -516,79 +516,116 @@ func (r *proxyRepository) SweepExpiredProxies(ctx context.Context, now time.Time logger.LegacyPrintf("repository.proxy", "[ProxyExpiry] proxy %d expired but fallback chain unresolved (cycle/all-expired); accounts kept", p.ID) } - changed, sweepErr := r.sweepOneExpiredProxy(ctx, p.ID, target, change) + changedAccountIDs, sweepErr := r.sweepOneExpiredProxy(ctx, p.ID, target, change) if sweepErr != nil { return totalChanged, sweepErr } - if changed > 0 { - totalChanged += changed - accountsTouched = true - } + totalChanged += int64(len(changedAccountIDs)) + allChangedAccountIDs = append(allChangedAccountIDs, changedAccountIDs...) } - if accountsTouched { - if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventFullRebuild, nil, nil, nil); err != nil { - logger.LegacyPrintf("repository.proxy", "[SchedulerOutbox] enqueue proxy expiry rebuild failed: err=%v", err) + changedAccountIDs := sortedUniqueAccountIDs(allChangedAccountIDs) + if len(changedAccountIDs) > 0 { + // 各代理的改投事务已经提交；这里仅汇总真实被 UPDATE 命中的账号， + // 避免代理到期时用全量重建刷新所有调度分桶。 + payload := map[string]any{"account_ids": changedAccountIDs} + if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil { + logger.LegacyPrintf("repository.proxy", "[SchedulerOutbox] enqueue proxy expiry account changes failed: err=%v", err) } } return totalChanged, nil } +func sortedUniqueAccountIDs(accountIDs []int64) []int64 { + if len(accountIDs) < 2 { + return accountIDs + } + sort.Slice(accountIDs, func(i, j int) bool { return accountIDs[i] < accountIDs[j] }) + write := 1 + for _, accountID := range accountIDs[1:] { + if accountID == accountIDs[write-1] { + continue + } + accountIDs[write] = accountID + write++ + } + return accountIDs[:write] +} + // sweepOneExpiredProxy 在单事务内原子执行：标记代理 expired + 改投绑定账号。 // 若 r.client 已绑定事务（测试注入场景），直接在 r.sql 上执行，由外层事务保证原子性。 -func (r *proxyRepository) sweepOneExpiredProxy(ctx context.Context, proxyID int64, target *int64, change bool) (int64, error) { +func (r *proxyRepository) sweepOneExpiredProxy(ctx context.Context, proxyID int64, target *int64, change bool) ([]int64, error) { // 尝试开启子事务；若 r.client 已是事务 client，则返回 ErrTxStarted，退回使用 r.sql。 tx, txErr := r.client.Tx(ctx) if txErr != nil { if txErr != dbent.ErrTxStarted { - return 0, txErr + return nil, txErr } // 已在外层事务中（集成测试场景），直接用 r.sql 执行 return r.sweepOneExpiredProxyOnExec(ctx, r.sql, proxyID, target, change) } // 使用新事务执行 - var n int64 + var accountIDs []int64 var err error - n, err = r.sweepOneExpiredProxyOnExec(ctx, tx, proxyID, target, change) + accountIDs, err = r.sweepOneExpiredProxyOnExec(ctx, tx, proxyID, target, change) if err != nil { _ = tx.Rollback() - return 0, err + return nil, err } if commitErr := tx.Commit(); commitErr != nil { - return 0, commitErr + return nil, commitErr } - return n, nil + return accountIDs, nil } // sweepOneExpiredProxyOnExec 在给定的 sqlExecutor 上执行：标记 expired + 改投账号。 -func (r *proxyRepository) sweepOneExpiredProxyOnExec(ctx context.Context, exec sqlExecutor, proxyID int64, target *int64, change bool) (int64, error) { +func (r *proxyRepository) sweepOneExpiredProxyOnExec(ctx context.Context, exec sqlExecutor, proxyID int64, target *int64, change bool) ([]int64, error) { if _, err := exec.ExecContext(ctx, `UPDATE proxies SET status=$1, updated_at=NOW() WHERE id=$2 AND deleted_at IS NULL`, service.StatusExpired, proxyID); err != nil { - return 0, err + return nil, err } if !change { - return 0, nil + return nil, nil } var ( - res sql.Result - err error + rows *sql.Rows + err error ) if target == nil { - res, err = exec.ExecContext(ctx, ` + rows, err = exec.QueryContext(ctx, ` UPDATE accounts SET proxy_id=NULL, proxy_fallback_origin_id=$1, updated_at=NOW() - WHERE proxy_id=$1 AND proxy_fallback_origin_id IS NULL AND deleted_at IS NULL`, proxyID) + WHERE proxy_id=$1 AND proxy_fallback_origin_id IS NULL AND deleted_at IS NULL + RETURNING id`, proxyID) } else { - res, err = exec.ExecContext(ctx, ` + rows, err = exec.QueryContext(ctx, ` UPDATE accounts SET proxy_id=$2, proxy_fallback_origin_id=$1, updated_at=NOW() - WHERE proxy_id=$1 AND proxy_fallback_origin_id IS NULL AND deleted_at IS NULL`, proxyID, *target) + WHERE proxy_id=$1 AND proxy_fallback_origin_id IS NULL AND deleted_at IS NULL + RETURNING id`, proxyID, *target) } if err != nil { - return 0, err + return nil, err } - n, _ := res.RowsAffected() - return n, nil + + // 必须在提交子事务前读完并关闭 RETURNING 结果集，否则连接仍可能处于 busy 状态。 + accountIDs := make([]int64, 0) + for rows.Next() { + var accountID int64 + if err := rows.Scan(&accountID); err != nil { + _ = rows.Close() + return nil, err + } + accountIDs = append(accountIDs, accountID) + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return nil, err + } + if err := rows.Close(); err != nil { + return nil, err + } + return accountIDs, nil } // CountExpired 返回已过期（status=expired）的代理数量。 diff --git a/backend/internal/repository/redis.go b/backend/internal/repository/redis.go index 2b4ee4e636..0ead4644c1 100644 --- a/backend/internal/repository/redis.go +++ b/backend/internal/repository/redis.go @@ -21,7 +21,11 @@ import ( // 2. MinIdleConns: 保持最小空闲连接，减少冷启动延迟（默认 10） // 3. DialTimeout/ReadTimeout/WriteTimeout: 精确控制各阶段超时 func InitRedis(cfg *config.Config) *redis.Client { - return redis.NewClient(buildRedisOptions(cfg)) + client := redis.NewClient(buildRedisOptions(cfg)) + if cfg.Server.EnableServerTiming { + client.AddHook(serverTimingRedisHook{}) + } + return client } // buildRedisOptions 构建 Redis 连接选项 diff --git a/backend/internal/repository/req_client_pool.go b/backend/internal/repository/req_client_pool.go index 32501f7b19..95ab27ce32 100644 --- a/backend/internal/repository/req_client_pool.go +++ b/backend/internal/repository/req_client_pool.go @@ -2,11 +2,13 @@ package repository import ( "fmt" + "net/http" "strings" "sync" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/imroc/req/v3" ) @@ -57,6 +59,7 @@ func getSharedReqClient(opts reqClientOptions) (*req.Client, error) { if trimmed != "" { client.SetProxyURL(trimmed) } + client = instrumentReqClient(client) actual, _ := sharedReqClients.LoadOrStore(key, client) if c, ok := actual.(*req.Client); ok { @@ -65,6 +68,17 @@ func getSharedReqClient(opts reqClientOptions) (*req.Client, error) { return client, nil } +func instrumentReqClient(client *req.Client) *req.Client { + if client == nil { + return nil + } + client.GetTransport().WrapRoundTripFunc(func(rt http.RoundTripper) req.HttpRoundTripFunc { + timed := servertiming.WrapRoundTripper(rt) + return timed.RoundTrip + }) + return client +} + func buildReqClientKey(opts reqClientOptions) string { return fmt.Sprintf("%s|%s|%t|%t", strings.TrimSpace(opts.ProxyURL), diff --git a/backend/internal/repository/req_client_pool_test.go b/backend/internal/repository/req_client_pool_test.go index 9067d0129f..3a27841c5a 100644 --- a/backend/internal/repository/req_client_pool_test.go +++ b/backend/internal/repository/req_client_pool_test.go @@ -1,12 +1,17 @@ package repository import ( + "context" + "net/http" + "net/http/httptest" "reflect" + "strings" "sync" "testing" "time" "unsafe" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/imroc/req/v3" "github.com/stretchr/testify/require" ) @@ -118,3 +123,20 @@ func TestCreateGeminiReqClient_ForceHTTP2Disabled(t *testing.T) { require.NoError(t, err) require.Equal(t, "", forceHTTPVersion(t, client)) } + +func TestInstrumentReqClientRecordsDependency(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + defer server.Close() + + collector := servertiming.New(time.Now()) + ctx := servertiming.WithCollector(context.Background(), collector) + client := instrumentReqClient(req.C()) + response, err := client.R().SetContext(ctx).Get(server.URL) + require.NoError(t, err) + require.Equal(t, http.StatusNoContent, response.StatusCode) + + header := collector.HeaderValue(time.Now(), "bypass") + require.True(t, strings.Contains(header, "dep_http;dur="), header) +} diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index c8e1fe14e0..c68bcf96e6 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "log/slog" "strconv" "time" @@ -163,14 +164,15 @@ func (c *schedulerCache) SetSnapshot(ctx context.Context, bucket service.Schedul versionStr := strconv.FormatInt(version, 10) snapshotKey := schedulerSnapshotKey(bucket, versionStr) - if err := c.writeAccounts(ctx, accounts); err != nil { + cacheableAccounts, err := c.writeAccounts(ctx, accounts) + if err != nil { return err } - if len(accounts) > 0 { + if len(cacheableAccounts) > 0 { // 使用序号作为 score，保持数据库返回的排序语义。 - members := make([]redis.Z, 0, len(accounts)) - for idx, account := range accounts { + members := make([]redis.Z, 0, len(cacheableAccounts)) + for idx, account := range cacheableAccounts { members = append(members, redis.Z{ Score: float64(idx), Member: strconv.FormatInt(account.ID, 10), @@ -224,7 +226,14 @@ func (c *schedulerCache) SetAccount(ctx context.Context, account *service.Accoun if account == nil || account.ID <= 0 { return nil } - return c.writeAccounts(ctx, []service.Account{*account}) + cacheableAccounts, err := c.writeAccounts(ctx, []service.Account{*account}) + if err != nil { + return err + } + if len(cacheableAccounts) == 0 { + return c.DeleteAccount(ctx, account.ID) + } + return nil } func (c *schedulerCache) DeleteAccount(ctx context.Context, accountID int64) error { @@ -262,13 +271,14 @@ func (c *schedulerCache) UpdateLastUsed(ctx context.Context, updates map[int64]t return err } account.LastUsedAt = ptrTime(updates[ids[i]]) - updated, err := json.Marshal(account) + updated, metaPayload, err := marshalSchedulerCacheAccount(*account) if err != nil { - return err - } - metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(*account)) - if err != nil { - return err + slog.Warn("scheduler cache removes account with unencodable payload", + "account_id", ids[i], + "error", err, + ) + pipe.Del(ctx, keys[i], schedulerAccountMetaKey(strconv.FormatInt(ids[i], 10))) + continue } pipe.Set(ctx, keys[i], updated, 0) pipe.Set(ctx, schedulerAccountMetaKey(strconv.FormatInt(ids[i], 10)), metaPayload, 0) @@ -359,12 +369,13 @@ func decodeCachedAccount(val any) (*service.Account, error) { return &account, nil } -func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.Account) error { +func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.Account) ([]service.Account, error) { if len(accounts) == 0 { - return nil + return nil, nil } pipe := c.rdb.Pipeline() + cacheableAccounts := make([]service.Account, 0, len(accounts)) pending := 0 flush := func() error { if pending == 0 { @@ -379,27 +390,43 @@ func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.A } for _, account := range accounts { - fullPayload, err := json.Marshal(account) + fullPayload, metaPayload, err := marshalSchedulerCacheAccount(account) if err != nil { - return err - } - metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(account)) - if err != nil { - return err + slog.Warn("scheduler cache skips account with unencodable payload", + "account_id", account.ID, + "error", err, + ) + continue } id := strconv.FormatInt(account.ID, 10) pipe.Set(ctx, schedulerAccountKey(id), fullPayload, 0) pipe.Set(ctx, schedulerAccountMetaKey(id), metaPayload, 0) + cacheableAccounts = append(cacheableAccounts, account) pending++ if pending >= c.writeChunkSize { if err := flush(); err != nil { - return err + return nil, err } } } - return flush() + if err := flush(); err != nil { + return nil, err + } + return cacheableAccounts, nil +} + +func marshalSchedulerCacheAccount(account service.Account) ([]byte, []byte, error) { + fullPayload, err := json.Marshal(account) + if err != nil { + return nil, nil, fmt.Errorf("marshal account: %w", err) + } + metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(account)) + if err != nil { + return nil, nil, fmt.Errorf("marshal account metadata: %w", err) + } + return fullPayload, metaPayload, nil } func (c *schedulerCache) mgetChunked(ctx context.Context, keys []string) ([]any, error) { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index 19c4cc4f36..ecca7f3892 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -3,12 +3,78 @@ package repository import ( + "context" "testing" + "time" "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" "github.com/stretchr/testify/require" ) +func newSchedulerCacheUnit(t *testing.T) *schedulerCache { + t.Helper() + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + cache, ok := newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, defaultSchedulerSnapshotWriteChunkSize).(*schedulerCache) + require.True(t, ok) + return cache +} + +func TestSchedulerCacheWriteAccountsSkipsUnencodableTimes(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + + cacheable, err := cache.writeAccounts(ctx, []service.Account{ + {ID: 111, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}, + {ID: 112, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}, + }) + require.NoError(t, err) + require.Len(t, cacheable, 1) + require.Equal(t, int64(111), cacheable[0].ID) + + cached, err := cache.GetAccount(ctx, 111) + require.NoError(t, err) + require.NotNil(t, cached) + + invalid, err := cache.GetAccount(ctx, 112) + require.NoError(t, err) + require.Nil(t, invalid) +} + +func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + + account := service.Account{ID: 113, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey} + require.NoError(t, cache.SetAccount(ctx, &account)) + + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + account.ExpiresAt = &invalidTime + require.NoError(t, cache.SetAccount(ctx, &account)) + + cached, err := cache.GetAccount(ctx, account.ID) + require.NoError(t, err) + require.Nil(t, cached) +} + +func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + account := service.Account{ID: 114, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey} + require.NoError(t, cache.SetAccount(ctx, &account)) + + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: invalidTime})) + + cached, err := cache.GetAccount(ctx, account.ID) + require.NoError(t, err) + require.Nil(t, cached) +} + func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) { account := service.Account{ ID: 42, diff --git a/backend/internal/repository/scheduler_outbox_repo.go b/backend/internal/repository/scheduler_outbox_repo.go index 59772fbb5a..500f6e850d 100644 --- a/backend/internal/repository/scheduler_outbox_repo.go +++ b/backend/internal/repository/scheduler_outbox_repo.go @@ -93,6 +93,24 @@ func (r *schedulerOutboxRepository) ListAfterAndReleaseDedup(ctx context.Context return events, nil } +func (r *schedulerOutboxRepository) FirstCreatedAtAfter(ctx context.Context, afterID int64) (time.Time, bool, error) { + var createdAt time.Time + err := r.db.QueryRowContext(ctx, ` + SELECT created_at + FROM scheduler_outbox + WHERE id > $1 + ORDER BY id ASC + LIMIT 1 + `, afterID).Scan(&createdAt) + if err == sql.ErrNoRows { + return time.Time{}, false, nil + } + if err != nil { + return time.Time{}, false, err + } + return createdAt, true, nil +} + func (r *schedulerOutboxRepository) MaxID(ctx context.Context) (int64, error) { var maxID int64 if err := r.db.QueryRowContext(ctx, "SELECT COALESCE(MAX(id), 0) FROM scheduler_outbox").Scan(&maxID); err != nil { diff --git a/backend/internal/repository/scheduler_outbox_repo_test.go b/backend/internal/repository/scheduler_outbox_repo_test.go index 619339d207..250014f6ee 100644 --- a/backend/internal/repository/scheduler_outbox_repo_test.go +++ b/backend/internal/repository/scheduler_outbox_repo_test.go @@ -4,11 +4,63 @@ import ( "context" "regexp" "testing" + "time" sqlmock "github.com/DATA-DOG/go-sqlmock" "github.com/stretchr/testify/require" ) +func TestSchedulerOutboxRepositoryFirstCreatedAtAfter(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + repo := &schedulerOutboxRepository{db: db} + createdAt := time.Now().UTC().Truncate(time.Microsecond) + const expectedSQL = ` + SELECT created_at + FROM scheduler_outbox + WHERE id > $1 + ORDER BY id ASC + LIMIT 1 + ` + mock.ExpectQuery(regexp.QuoteMeta(expectedSQL)). + WithArgs(int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(createdAt)) + + got, ok, err := repo.FirstCreatedAtAfter(context.Background(), 42) + + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, createdAt, got) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestSchedulerOutboxRepositoryFirstCreatedAtAfterReturnsNotFound(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + repo := &schedulerOutboxRepository{db: db} + const expectedSQL = ` + SELECT created_at + FROM scheduler_outbox + WHERE id > $1 + ORDER BY id ASC + LIMIT 1 + ` + mock.ExpectQuery(regexp.QuoteMeta(expectedSQL)). + WithArgs(int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"created_at"})) + + got, ok, err := repo.FirstCreatedAtAfter(context.Background(), 42) + + require.NoError(t, err) + require.False(t, ok) + require.True(t, got.IsZero()) + require.NoError(t, mock.ExpectationsWereMet()) +} + func TestSchedulerOutboxRepositoryDeleteConsumedUpToUsesBoundedCTE(t *testing.T) { db, mock, err := sqlmock.New() require.NoError(t, err) diff --git a/backend/internal/repository/server_timing_redis.go b/backend/internal/repository/server_timing_redis.go new file mode 100644 index 0000000000..dba35450de --- /dev/null +++ b/backend/internal/repository/server_timing_redis.go @@ -0,0 +1,39 @@ +package repository + +import ( + "context" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" + "github.com/redis/go-redis/v9" +) + +type serverTimingRedisHook struct{} + +func (serverTimingRedisHook) DialHook(next redis.DialHook) redis.DialHook { + return next +} + +func (serverTimingRedisHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook { + return func(ctx context.Context, cmd redis.Cmder) error { + if !servertiming.Active(ctx) { + return next(ctx, cmd) + } + startedAt := time.Now() + err := next(ctx, cmd) + servertiming.Record(ctx, servertiming.MetricRedis, startedAt, time.Now(), 1) + return err + } +} + +func (serverTimingRedisHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook { + return func(ctx context.Context, cmds []redis.Cmder) error { + if !servertiming.Active(ctx) { + return next(ctx, cmds) + } + startedAt := time.Now() + err := next(ctx, cmds) + servertiming.Record(ctx, servertiming.MetricRedis, startedAt, time.Now(), len(cmds)) + return err + } +} diff --git a/backend/internal/repository/server_timing_redis_test.go b/backend/internal/repository/server_timing_redis_test.go new file mode 100644 index 0000000000..d1ae47e3b0 --- /dev/null +++ b/backend/internal/repository/server_timing_redis_test.go @@ -0,0 +1,63 @@ +package repository + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" + "github.com/redis/go-redis/v9" +) + +func TestServerTimingRedisHookRecordsCommands(t *testing.T) { + collector := servertiming.New(time.Now()) + ctx := servertiming.WithCollector(context.Background(), collector) + hook := serverTimingRedisHook{} + + process := hook.ProcessHook(func(context.Context, redis.Cmder) error { + time.Sleep(time.Millisecond) + return errors.New("redis failure") + }) + if err := process(ctx, redis.NewStringCmd(ctx, "get", "sensitive-key")); err == nil { + t.Fatal("ProcessHook did not return the underlying error") + } + + pipeline := hook.ProcessPipelineHook(func(context.Context, []redis.Cmder) error { + time.Sleep(time.Millisecond) + return nil + }) + commands := []redis.Cmder{ + redis.NewStringCmd(ctx, "get", "first-secret"), + redis.NewStringCmd(ctx, "get", "second-secret"), + redis.NewStatusCmd(ctx, "set", "third-secret", "value"), + } + if err := pipeline(ctx, commands); err != nil { + t.Fatal(err) + } + + header := collector.HeaderValue(time.Now(), "bypass") + if !strings.Contains(header, `commands=4`) { + t.Fatalf("header %q does not report one command and a three-command pipeline", header) + } + if strings.Contains(header, "secret") || strings.Contains(header, "get") { + t.Fatalf("Redis command details leaked into header: %q", header) + } +} + +func TestServerTimingRedisHookSkipsInactiveContext(t *testing.T) { + called := false + hook := serverTimingRedisHook{} + process := hook.ProcessHook(func(context.Context, redis.Cmder) error { + called = true + return nil + }) + ctx := context.Background() + if err := process(ctx, redis.NewStringCmd(ctx, "ping")); err != nil { + t.Fatal(err) + } + if !called { + t.Fatal("inactive Redis command did not reach the next hook") + } +} diff --git a/backend/internal/repository/server_timing_sql.go b/backend/internal/repository/server_timing_sql.go new file mode 100644 index 0000000000..062663f08b --- /dev/null +++ b/backend/internal/repository/server_timing_sql.go @@ -0,0 +1,311 @@ +package repository + +import ( + "context" + "database/sql/driver" + "errors" + "io" + "reflect" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" +) + +type serverTimingConnector struct { + base driver.Connector +} + +func newServerTimingConnector(base driver.Connector) driver.Connector { + return &serverTimingConnector{base: base} +} + +func (c *serverTimingConnector) Connect(ctx context.Context) (driver.Conn, error) { + startedAt := time.Now() + conn, err := c.base.Connect(ctx) + servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now()) + if err != nil { + return nil, err + } + return &serverTimingConn{Conn: conn}, nil +} + +func (c *serverTimingConnector) Driver() driver.Driver { + return c.base.Driver() +} + +type serverTimingConn struct { + driver.Conn +} + +func (c *serverTimingConn) Prepare(query string) (driver.Stmt, error) { + stmt, err := c.Conn.Prepare(query) + if err != nil { + return nil, err + } + return &serverTimingStmt{Stmt: stmt}, nil +} + +func (c *serverTimingConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) { + startedAt := time.Now() + var ( + stmt driver.Stmt + err error + ) + if preparer, ok := c.Conn.(driver.ConnPrepareContext); ok { + stmt, err = preparer.PrepareContext(ctx, query) + } else { + stmt, err = c.Conn.Prepare(query) + } + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + if err != nil { + return nil, err + } + return &serverTimingStmt{Stmt: stmt}, nil +} + +func (c *serverTimingConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) { + execer, ok := c.Conn.(driver.ExecerContext) + if !ok { + return nil, driver.ErrSkip + } + startedAt := time.Now() + result, err := execer.ExecContext(ctx, query, args) + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + return result, err +} + +func (c *serverTimingConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) { + queryer, ok := c.Conn.(driver.QueryerContext) + if !ok { + return nil, driver.ErrSkip + } + startedAt := time.Now() + rows, err := queryer.QueryContext(ctx, query, args) + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + if err != nil || rows == nil { + return rows, err + } + return newServerTimingRows(ctx, rows), nil +} + +func (c *serverTimingConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) { + startedAt := time.Now() + var ( + tx driver.Tx + err error + ) + if beginner, ok := c.Conn.(driver.ConnBeginTx); ok { + tx, err = beginner.BeginTx(ctx, opts) + } else { + if opts.Isolation != driver.IsolationLevel(0) { + return nil, errors.New("driver does not support non-default isolation") + } + if opts.ReadOnly { + return nil, errors.New("driver does not support read-only transactions") + } + // The wrapper exposes ConnBeginTx, so it must retain database/sql's + // legacy fallback for drivers that only implement Conn.Begin. + tx, err = c.Conn.Begin() //nolint:staticcheck // Required driver compatibility fallback. + } + servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now()) + if err != nil || tx == nil { + return tx, err + } + return &serverTimingTx{Tx: tx, ctx: ctx}, nil +} + +func (c *serverTimingConn) Ping(ctx context.Context) error { + if pinger, ok := c.Conn.(driver.Pinger); ok { + startedAt := time.Now() + err := pinger.Ping(ctx) + servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err + } + return nil +} + +func (c *serverTimingConn) ResetSession(ctx context.Context) error { + if resetter, ok := c.Conn.(driver.SessionResetter); ok { + startedAt := time.Now() + err := resetter.ResetSession(ctx) + servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err + } + return nil +} + +func (c *serverTimingConn) IsValid() bool { + if validator, ok := c.Conn.(driver.Validator); ok { + return validator.IsValid() + } + return true +} + +func (c *serverTimingConn) CheckNamedValue(value *driver.NamedValue) error { + if checker, ok := c.Conn.(driver.NamedValueChecker); ok { + return checker.CheckNamedValue(value) + } + return driver.ErrSkip +} + +type serverTimingStmt struct { + driver.Stmt +} + +func (s *serverTimingStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) { + startedAt := time.Now() + var ( + result driver.Result + err error + ) + if execer, ok := s.Stmt.(driver.StmtExecContext); ok { + result, err = execer.ExecContext(ctx, args) + } else { + var values []driver.Value + values, err = namedValues(args) + if err == nil { + // The wrapper exposes StmtExecContext and must preserve the fallback + // database/sql would use for a legacy driver statement. + result, err = s.Stmt.Exec(values) //nolint:staticcheck // Required driver compatibility fallback. + } + } + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + return result, err +} + +func (s *serverTimingStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) { + startedAt := time.Now() + var ( + rows driver.Rows + err error + ) + if queryer, ok := s.Stmt.(driver.StmtQueryContext); ok { + rows, err = queryer.QueryContext(ctx, args) + } else { + var values []driver.Value + values, err = namedValues(args) + if err == nil { + // The wrapper exposes StmtQueryContext and must preserve the fallback + // database/sql would use for a legacy driver statement. + rows, err = s.Stmt.Query(values) //nolint:staticcheck // Required driver compatibility fallback. + } + } + servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1) + if err != nil || rows == nil { + return rows, err + } + return newServerTimingRows(ctx, rows), nil +} + +func (s *serverTimingStmt) CheckNamedValue(value *driver.NamedValue) error { + if checker, ok := s.Stmt.(driver.NamedValueChecker); ok { + return checker.CheckNamedValue(value) + } + return driver.ErrSkip +} + +func namedValues(args []driver.NamedValue) ([]driver.Value, error) { + values := make([]driver.Value, len(args)) + for i, arg := range args { + if arg.Name != "" { + return nil, errors.New("named parameters are not supported") + } + values[i] = arg.Value + } + return values, nil +} + +type serverTimingRows struct { + driver.Rows + ctx context.Context +} + +func newServerTimingRows(ctx context.Context, rows driver.Rows) *serverTimingRows { + return &serverTimingRows{Rows: rows, ctx: ctx} +} + +func (r *serverTimingRows) Close() error { + startedAt := time.Now() + err := r.Rows.Close() + servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} + +func (r *serverTimingRows) Next(dest []driver.Value) error { + startedAt := time.Now() + err := r.Rows.Next(dest) + servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} + +func (r *serverTimingRows) HasNextResultSet() bool { + if rows, ok := r.Rows.(driver.RowsNextResultSet); ok { + return rows.HasNextResultSet() + } + return false +} + +func (r *serverTimingRows) NextResultSet() error { + rows, ok := r.Rows.(driver.RowsNextResultSet) + if !ok { + return io.EOF + } + startedAt := time.Now() + err := rows.NextResultSet() + servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} + +func (r *serverTimingRows) ColumnTypeScanType(index int) reflect.Type { + if rows, ok := r.Rows.(driver.RowsColumnTypeScanType); ok { + return rows.ColumnTypeScanType(index) + } + return reflect.TypeOf(new(any)).Elem() +} + +func (r *serverTimingRows) ColumnTypeDatabaseTypeName(index int) string { + if rows, ok := r.Rows.(driver.RowsColumnTypeDatabaseTypeName); ok { + return rows.ColumnTypeDatabaseTypeName(index) + } + return "" +} + +func (r *serverTimingRows) ColumnTypeLength(index int) (int64, bool) { + if rows, ok := r.Rows.(driver.RowsColumnTypeLength); ok { + return rows.ColumnTypeLength(index) + } + return 0, false +} + +func (r *serverTimingRows) ColumnTypeNullable(index int) (bool, bool) { + if rows, ok := r.Rows.(driver.RowsColumnTypeNullable); ok { + return rows.ColumnTypeNullable(index) + } + return false, false +} + +func (r *serverTimingRows) ColumnTypePrecisionScale(index int) (int64, int64, bool) { + if rows, ok := r.Rows.(driver.RowsColumnTypePrecisionScale); ok { + return rows.ColumnTypePrecisionScale(index) + } + return 0, 0, false +} + +type serverTimingTx struct { + driver.Tx + ctx context.Context +} + +func (t *serverTimingTx) Commit() error { + startedAt := time.Now() + err := t.Tx.Commit() + servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} + +func (t *serverTimingTx) Rollback() error { + startedAt := time.Now() + err := t.Tx.Rollback() + servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now()) + return err +} diff --git a/backend/internal/repository/server_timing_sql_test.go b/backend/internal/repository/server_timing_sql_test.go new file mode 100644 index 0000000000..3a8bbbe03e --- /dev/null +++ b/backend/internal/repository/server_timing_sql_test.go @@ -0,0 +1,258 @@ +package repository + +import ( + "context" + "database/sql/driver" + "io" + "regexp" + "strconv" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" +) + +const fakeDriverDelay = 2 * time.Millisecond + +type timingFakeDriver struct{} + +func (timingFakeDriver) Open(string) (driver.Conn, error) { return newTimingFakeConn(), nil } + +type timingFakeConnector struct { + conn driver.Conn +} + +func (c timingFakeConnector) Connect(context.Context) (driver.Conn, error) { + time.Sleep(fakeDriverDelay) + return c.conn, nil +} + +func (timingFakeConnector) Driver() driver.Driver { return timingFakeDriver{} } + +type timingFakeConn struct{} + +func newTimingFakeConn() *timingFakeConn { return &timingFakeConn{} } + +func (c *timingFakeConn) Prepare(string) (driver.Stmt, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeStmt{}, nil +} + +func (c *timingFakeConn) PrepareContext(context.Context, string) (driver.Stmt, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeStmt{}, nil +} + +func (c *timingFakeConn) Close() error { return nil } + +func (c *timingFakeConn) Begin() (driver.Tx, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeTx{}, nil +} + +func (c *timingFakeConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeTx{}, nil +} + +func (c *timingFakeConn) ExecContext(context.Context, string, []driver.NamedValue) (driver.Result, error) { + time.Sleep(fakeDriverDelay) + return driver.RowsAffected(1), nil +} + +func (c *timingFakeConn) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil +} + +func (c *timingFakeConn) Ping(context.Context) error { + time.Sleep(fakeDriverDelay) + return nil +} + +func (c *timingFakeConn) ResetSession(context.Context) error { + time.Sleep(fakeDriverDelay) + return nil +} + +type timingFakeStmt struct{} + +func (s *timingFakeStmt) Close() error { return nil } +func (s *timingFakeStmt) NumInput() int { return -1 } + +func (s *timingFakeStmt) Exec([]driver.Value) (driver.Result, error) { + time.Sleep(fakeDriverDelay) + return driver.RowsAffected(1), nil +} + +func (s *timingFakeStmt) Query([]driver.Value) (driver.Rows, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil +} + +func (s *timingFakeStmt) ExecContext(context.Context, []driver.NamedValue) (driver.Result, error) { + time.Sleep(fakeDriverDelay) + return driver.RowsAffected(1), nil +} + +func (s *timingFakeStmt) QueryContext(context.Context, []driver.NamedValue) (driver.Rows, error) { + time.Sleep(fakeDriverDelay) + return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil +} + +type timingFakeRows struct { + values [][]driver.Value + index int +} + +func (r *timingFakeRows) Columns() []string { return []string{"value"} } + +func (r *timingFakeRows) Close() error { + time.Sleep(fakeDriverDelay) + return nil +} + +func (r *timingFakeRows) Next(dest []driver.Value) error { + time.Sleep(fakeDriverDelay) + if r.index >= len(r.values) { + return io.EOF + } + copy(dest, r.values[r.index]) + r.index++ + return nil +} + +type timingFakeTx struct{} + +func (t *timingFakeTx) Commit() error { + time.Sleep(fakeDriverDelay) + return nil +} + +func (t *timingFakeTx) Rollback() error { + time.Sleep(fakeDriverDelay) + return nil +} + +func metricDuration(t *testing.T, header, metric string) float64 { + t.Helper() + re := regexp.MustCompile(`(?:^|, )` + regexp.QuoteMeta(metric) + `;dur=([0-9]+(?:\.[0-9]+)?)`) + match := re.FindStringSubmatch(header) + if len(match) != 2 { + t.Fatalf("metric %q missing from header %q", metric, header) + } + value, err := strconv.ParseFloat(match[1], 64) + if err != nil { + t.Fatalf("parse %s duration: %v", metric, err) + } + return value +} + +func TestServerTimingConnectorRecordsDriverCallsWithoutRowLifetime(t *testing.T) { + startedAt := time.Now() + collector := servertiming.New(startedAt) + ctx := servertiming.WithCollector(context.Background(), collector) + + wrapped := newServerTimingConnector(timingFakeConnector{conn: newTimingFakeConn()}) + rawConn, err := wrapped.Connect(ctx) + if err != nil { + t.Fatal(err) + } + conn, ok := rawConn.(*serverTimingConn) + if !ok { + t.Fatalf("Connect() returned %T, want *serverTimingConn", rawConn) + } + + if _, err := conn.ExecContext(ctx, "sensitive update", nil); err != nil { + t.Fatal(err) + } + rows, err := conn.QueryContext(ctx, "sensitive select", nil) + if err != nil { + t.Fatal(err) + } + values := make([]driver.Value, 1) + if err := rows.Next(values); err != nil { + t.Fatal(err) + } + + // Application work between row reads must remain app time. + time.Sleep(30 * time.Millisecond) + if err := rows.Next(values); err != io.EOF { + t.Fatalf("rows.Next() = %v, want EOF", err) + } + if err := rows.Close(); err != nil { + t.Fatal(err) + } + + header := collector.HeaderValue(time.Now(), "bypass") + if !strings.Contains(header, `queries=2`) { + t.Fatalf("header %q does not report two SQL operations", header) + } + if strings.Contains(header, "sensitive") { + t.Fatalf("SQL text leaked into header: %q", header) + } + if app, db := metricDuration(t, header, "app"), metricDuration(t, header, "db"); app <= db { + t.Fatalf("row processing gap was counted as DB time: app=%.1fms db=%.1fms header=%q", app, db, header) + } +} + +func TestServerTimingPreparedStatementsAndTransactions(t *testing.T) { + collector := servertiming.New(time.Now()) + ctx := servertiming.WithCollector(context.Background(), collector) + conn := &serverTimingConn{Conn: newTimingFakeConn()} + + stmt, err := conn.PrepareContext(ctx, "prepare sensitive statement") + if err != nil { + t.Fatal(err) + } + timedStmt, ok := stmt.(*serverTimingStmt) + if !ok { + t.Fatalf("PrepareContext() returned %T, want *serverTimingStmt", stmt) + } + if _, err := timedStmt.ExecContext(ctx, nil); err != nil { + t.Fatal(err) + } + rows, err := timedStmt.QueryContext(ctx, nil) + if err != nil { + t.Fatal(err) + } + if err := rows.Close(); err != nil { + t.Fatal(err) + } + + tx, err := conn.BeginTx(ctx, driver.TxOptions{}) + if err != nil { + t.Fatal(err) + } + if err := tx.Commit(); err != nil { + t.Fatal(err) + } + if err := conn.Ping(ctx); err != nil { + t.Fatal(err) + } + if err := conn.ResetSession(ctx); err != nil { + t.Fatal(err) + } + + header := collector.HeaderValue(time.Now(), "bypass") + if !strings.Contains(header, `queries=3`) { + t.Fatalf("header %q does not report prepare, exec, and query operations", header) + } + if metricDuration(t, header, "db") <= 0 { + t.Fatalf("DB duration was not recorded: %q", header) + } +} + +func TestNamedValuesRejectNamedParameters(t *testing.T) { + if _, err := namedValues([]driver.NamedValue{{Name: "secret", Value: 1}}); err == nil { + t.Fatal("namedValues accepted a named parameter") + } + values, err := namedValues([]driver.NamedValue{{Ordinal: 1, Value: "value"}}) + if err != nil { + t.Fatal(err) + } + if len(values) != 1 || values[0] != "value" { + t.Fatalf("namedValues() = %#v", values) + } +} diff --git a/backend/internal/repository/usage_log_repo_insert.go b/backend/internal/repository/usage_log_repo_insert.go index dfd8969512..ec09b308a0 100644 --- a/backend/internal/repository/usage_log_repo_insert.go +++ b/backend/internal/repository/usage_log_repo_insert.go @@ -71,6 +71,7 @@ var usageLogInsertArgTypes = [...]string{ "text", // inbound_endpoint "text", // upstream_endpoint "boolean", // cache_ttl_overridden + "boolean", // long_context_billing_applied "bigint", // channel_id "text", // model_mapping_chain "text", // billing_tier @@ -263,6 +264,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, + long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, @@ -275,7 +277,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53 + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54 ) ON CONFLICT (request_id, api_key_id) DO NOTHING RETURNING id, created_at @@ -714,6 +716,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage inbound_endpoint, upstream_endpoint, cache_ttl_overridden, + long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, @@ -722,7 +725,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage created_at ) AS (VALUES `) - args := make([]any, 0, len(keys)*53) + args := make([]any, 0, len(keys)*54) argPos := 1 for idx, key := range keys { if idx > 0 { @@ -798,6 +801,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage inbound_endpoint, upstream_endpoint, cache_ttl_overridden, + long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, @@ -853,6 +857,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage inbound_endpoint, upstream_endpoint, cache_ttl_overridden, + long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, @@ -948,6 +953,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( inbound_endpoint, upstream_endpoint, cache_ttl_overridden, + long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, @@ -956,7 +962,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( created_at ) AS (VALUES `) - args := make([]any, 0, len(preparedList)*53) + args := make([]any, 0, len(preparedList)*54) argPos := 1 for idx, prepared := range preparedList { if idx > 0 { @@ -1029,6 +1035,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( inbound_endpoint, upstream_endpoint, cache_ttl_overridden, + long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, @@ -1084,6 +1091,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( inbound_endpoint, upstream_endpoint, cache_ttl_overridden, + long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, @@ -1147,6 +1155,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared inbound_endpoint, upstream_endpoint, cache_ttl_overridden, + long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, @@ -1159,7 +1168,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53 + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54 ) ON CONFLICT (request_id, api_key_id) DO NOTHING `, prepared.args...) @@ -1264,6 +1273,7 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared { inboundEndpoint, upstreamEndpoint, log.CacheTTLOverridden, + log.LongContextBillingApplied, channelID, modelMappingChain, billingTier, diff --git a/backend/internal/repository/usage_log_repo_query.go b/backend/internal/repository/usage_log_repo_query.go index c178429bab..1fdedd8665 100644 --- a/backend/internal/repository/usage_log_repo_query.go +++ b/backend/internal/repository/usage_log_repo_query.go @@ -19,7 +19,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/service" ) -const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" +const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" func (r *usageLogRepository) GetByID(ctx context.Context, id int64) (log *service.UsageLog, err error) { query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE id = $1" @@ -425,60 +425,61 @@ func (r *usageLogRepository) loadSubscriptions(ctx context.Context, ids []int64) func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, error) { var ( - id int64 - userID int64 - apiKeyID int64 - accountID int64 - requestID sql.NullString - model string - requestedModel sql.NullString - upstreamModel sql.NullString - groupID sql.NullInt64 - subscriptionID sql.NullInt64 - inputTokens int - outputTokens int - cacheCreationTokens int - cacheReadTokens int - cacheCreation5m int - cacheCreation1h int - imageOutputTokens int - imageOutputCost float64 - inputCost float64 - outputCost float64 - cacheCreationCost float64 - cacheReadCost float64 - totalCost float64 - actualCost float64 - rateMultiplier float64 - accountRateMultiplier sql.NullFloat64 - billingType int16 - requestTypeRaw int16 - stream bool - openaiWSMode bool - durationMs sql.NullInt64 - firstTokenMs sql.NullInt64 - userAgent sql.NullString - ipAddress sql.NullString - imageCount int - imageSize sql.NullString - imageInputSize sql.NullString - imageOutputSize sql.NullString - imageSizeSource sql.NullString - imageSizeBreakdown sql.NullString - videoCount int - videoResolution sql.NullString - videoDurationSeconds sql.NullInt64 - serviceTier sql.NullString - reasoningEffort sql.NullString - inboundEndpoint sql.NullString - upstreamEndpoint sql.NullString - cacheTTLOverridden bool - channelID sql.NullInt64 - modelMappingChain sql.NullString - billingTier sql.NullString - billingMode sql.NullString - accountStatsCost sql.NullFloat64 - createdAt time.Time + id int64 + userID int64 + apiKeyID int64 + accountID int64 + requestID sql.NullString + model string + requestedModel sql.NullString + upstreamModel sql.NullString + groupID sql.NullInt64 + subscriptionID sql.NullInt64 + inputTokens int + outputTokens int + cacheCreationTokens int + cacheReadTokens int + cacheCreation5m int + cacheCreation1h int + imageOutputTokens int + imageOutputCost float64 + inputCost float64 + outputCost float64 + cacheCreationCost float64 + cacheReadCost float64 + totalCost float64 + actualCost float64 + rateMultiplier float64 + accountRateMultiplier sql.NullFloat64 + billingType int16 + requestTypeRaw int16 + stream bool + openaiWSMode bool + durationMs sql.NullInt64 + firstTokenMs sql.NullInt64 + userAgent sql.NullString + ipAddress sql.NullString + imageCount int + imageSize sql.NullString + imageInputSize sql.NullString + imageOutputSize sql.NullString + imageSizeSource sql.NullString + imageSizeBreakdown sql.NullString + videoCount int + videoResolution sql.NullString + videoDurationSeconds sql.NullInt64 + serviceTier sql.NullString + reasoningEffort sql.NullString + inboundEndpoint sql.NullString + upstreamEndpoint sql.NullString + cacheTTLOverridden bool + longContextBillingApplied bool + channelID sql.NullInt64 + modelMappingChain sql.NullString + billingTier sql.NullString + billingMode sql.NullString + accountStatsCost sql.NullFloat64 + createdAt time.Time ) if err := scanner.Scan( @@ -530,6 +531,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e &inboundEndpoint, &upstreamEndpoint, &cacheTTLOverridden, + &longContextBillingApplied, &channelID, &modelMappingChain, &billingTier, @@ -541,34 +543,35 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e } log := &service.UsageLog{ - ID: id, - UserID: userID, - APIKeyID: apiKeyID, - AccountID: accountID, - Model: model, - RequestedModel: coalesceTrimmedString(requestedModel, model), - InputTokens: inputTokens, - OutputTokens: outputTokens, - CacheCreationTokens: cacheCreationTokens, - CacheReadTokens: cacheReadTokens, - CacheCreation5mTokens: cacheCreation5m, - CacheCreation1hTokens: cacheCreation1h, - ImageOutputTokens: imageOutputTokens, - ImageOutputCost: imageOutputCost, - InputCost: inputCost, - OutputCost: outputCost, - CacheCreationCost: cacheCreationCost, - CacheReadCost: cacheReadCost, - TotalCost: totalCost, - ActualCost: actualCost, - RateMultiplier: rateMultiplier, - AccountRateMultiplier: nullFloat64Ptr(accountRateMultiplier), - BillingType: int8(billingType), - RequestType: service.RequestTypeFromInt16(requestTypeRaw), - ImageCount: imageCount, - VideoCount: videoCount, - CacheTTLOverridden: cacheTTLOverridden, - CreatedAt: createdAt, + ID: id, + UserID: userID, + APIKeyID: apiKeyID, + AccountID: accountID, + Model: model, + RequestedModel: coalesceTrimmedString(requestedModel, model), + InputTokens: inputTokens, + OutputTokens: outputTokens, + CacheCreationTokens: cacheCreationTokens, + CacheReadTokens: cacheReadTokens, + CacheCreation5mTokens: cacheCreation5m, + CacheCreation1hTokens: cacheCreation1h, + ImageOutputTokens: imageOutputTokens, + ImageOutputCost: imageOutputCost, + InputCost: inputCost, + OutputCost: outputCost, + CacheCreationCost: cacheCreationCost, + CacheReadCost: cacheReadCost, + TotalCost: totalCost, + ActualCost: actualCost, + RateMultiplier: rateMultiplier, + AccountRateMultiplier: nullFloat64Ptr(accountRateMultiplier), + BillingType: int8(billingType), + RequestType: service.RequestTypeFromInt16(requestTypeRaw), + ImageCount: imageCount, + VideoCount: videoCount, + CacheTTLOverridden: cacheTTLOverridden, + LongContextBillingApplied: longContextBillingApplied, + CreatedAt: createdAt, } // 先回填 legacy 字段，再基于 legacy + request_type 计算最终请求类型，保证历史数据兼容。 log.Stream = stream diff --git a/backend/internal/repository/usage_log_repo_request_type_test.go b/backend/internal/repository/usage_log_repo_request_type_test.go index c32ad2b63f..052c319183 100644 --- a/backend/internal/repository/usage_log_repo_request_type_test.go +++ b/backend/internal/repository/usage_log_repo_request_type_test.go @@ -88,6 +88,7 @@ func TestUsageLogRepositoryCreateSyncRequestTypeAndLegacyFields(t *testing.T) { sqlmock.AnyArg(), // inbound_endpoint sqlmock.AnyArg(), // upstream_endpoint log.CacheTTLOverridden, + log.LongContextBillingApplied, sqlmock.AnyArg(), // channel_id sqlmock.AnyArg(), // model_mapping_chain sqlmock.AnyArg(), // billing_tier @@ -174,6 +175,7 @@ func TestUsageLogRepositoryCreate_PersistsServiceTier(t *testing.T) { sqlmock.AnyArg(), sqlmock.AnyArg(), log.CacheTTLOverridden, + log.LongContextBillingApplied, sqlmock.AnyArg(), // channel_id sqlmock.AnyArg(), // model_mapping_chain sqlmock.AnyArg(), // billing_tier @@ -813,6 +815,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, sql.NullString{}, false, + false, sql.NullInt64{}, sql.NullString{}, sql.NullString{}, @@ -884,6 +887,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, sql.NullString{}, false, + false, sql.NullInt64{}, // channel_id sql.NullString{}, // model_mapping_chain sql.NullString{}, // billing_tier @@ -939,6 +943,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, sql.NullString{}, false, + false, sql.NullInt64{}, // channel_id sql.NullString{}, // model_mapping_chain sql.NullString{}, // billing_tier @@ -994,6 +999,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, sql.NullString{}, false, + false, sql.NullInt64{}, // channel_id sql.NullString{}, // model_mapping_chain sql.NullString{}, // billing_tier diff --git a/backend/internal/server/api_contract_test.go b/backend/internal/server/api_contract_test.go index 372cc46bbf..a5e3fde155 100644 --- a/backend/internal/server/api_contract_test.go +++ b/backend/internal/server/api_contract_test.go @@ -594,6 +594,7 @@ func TestAPIContracts(t *testing.T) { "total_cost": 0.5, "actual_cost": 0.5, "rate_multiplier": 1, + "long_context_billing_applied": false, "billing_type": 0, "stream": true, "duration_ms": 100, diff --git a/backend/internal/server/middleware/cors.go b/backend/internal/server/middleware/cors.go index 03d5d025de..0283d53115 100644 --- a/backend/internal/server/middleware/cors.go +++ b/backend/internal/server/middleware/cors.go @@ -52,7 +52,7 @@ func CORS(cfg config.CORSConfig) gin.HandlerFunc { } allowHeaders := []string{ "Content-Type", "Content-Length", "Accept-Encoding", "X-CSRF-Token", "Authorization", - "accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key", + "accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key", "X-Admin-UI-Request", } // OpenAI Node SDK 会发送 x-stainless-* 请求头，需在 CORS 中显式放行。 openAIProperties := []string{ @@ -83,7 +83,7 @@ func CORS(cfg config.CORSConfig) gin.HandlerFunc { } c.Writer.Header().Set("Access-Control-Allow-Headers", allowHeadersValue) c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE, PATCH") - c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag") + c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag, Server-Timing") c.Writer.Header().Set("Access-Control-Max-Age", "86400") } // 处理预检请求 diff --git a/backend/internal/server/middleware/cors_test.go b/backend/internal/server/middleware/cors_test.go index 6d0bea3608..6a61f696df 100644 --- a/backend/internal/server/middleware/cors_test.go +++ b/backend/internal/server/middleware/cors_test.go @@ -103,8 +103,10 @@ func TestCORS_AllowedOrigin_HasAllowHeaders(t *testing.T) { // 应设置 Allow-Headers、Allow-Methods 和 Max-Age assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"), "允许的 origin 应收到 Allow-Headers") + assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-Admin-UI-Request") assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"), "允许的 origin 应收到 Allow-Methods") + assert.Contains(t, w.Header().Get("Access-Control-Expose-Headers"), "Server-Timing") assert.Equal(t, "86400", w.Header().Get("Access-Control-Max-Age"), "允许的 origin 应收到 Max-Age=86400") assert.Equal(t, "https://allowed.example.com", w.Header().Get("Access-Control-Allow-Origin"), diff --git a/backend/internal/server/middleware/server_timing.go b/backend/internal/server/middleware/server_timing.go new file mode 100644 index 0000000000..2bb21071e0 --- /dev/null +++ b/backend/internal/server/middleware/server_timing.go @@ -0,0 +1,132 @@ +package middleware + +import ( + "net/http" + "strings" + "sync" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" + "github.com/gin-gonic/gin" +) + +const ( + snapshotCacheHeader = "X-Snapshot-Cache" + usageCacheHeader = "X-Usage-Stats-Cache" +) + +type serverTimingResponseWriter struct { + gin.ResponseWriter + context *gin.Context + once sync.Once +} + +func (w *serverTimingResponseWriter) Unwrap() http.ResponseWriter { + return w.ResponseWriter +} + +// ServerTiming collects timing only for requests made by the Admin web UI. +func ServerTiming(enabled bool) gin.HandlerFunc { + if !enabled { + return func(c *gin.Context) { + c.Next() + } + } + return func(c *gin.Context) { + if !isAdminUIRequest(c) || c.Request == nil { + c.Next() + return + } + + collector := servertiming.New(time.Now()) + c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector)) + writer := &serverTimingResponseWriter{ + ResponseWriter: c.Writer, + context: c, + } + c.Writer = writer + c.Next() + writer.finalize() + } +} + +func (w *serverTimingResponseWriter) WriteHeader(statusCode int) { + w.ResponseWriter.WriteHeader(statusCode) +} + +func (w *serverTimingResponseWriter) WriteHeaderNow() { + w.finalize() + w.ResponseWriter.WriteHeaderNow() +} + +func (w *serverTimingResponseWriter) Write(data []byte) (int, error) { + w.finalize() + return w.ResponseWriter.Write(data) +} + +func (w *serverTimingResponseWriter) WriteString(data string) (int, error) { + w.finalize() + return w.ResponseWriter.WriteString(data) +} + +func (w *serverTimingResponseWriter) Flush() { + w.finalize() + w.ResponseWriter.Flush() +} + +func (w *serverTimingResponseWriter) finalize() { + if w == nil { + return + } + w.once.Do(func() { + if value := ServerTimingHeaderValue(w.context); value != "" { + w.ResponseWriter.Header().Set(servertiming.HeaderName, value) + } + }) +} + +// ServerTimingHeaderValue returns a timing value only for an authenticated admin. +func ServerTimingHeaderValue(c *gin.Context) string { + if c == nil || c.Request == nil { + return "" + } + role, ok := GetUserRoleFromContext(c) + if !ok || role != "admin" { + return "" + } + return servertiming.HeaderValue(c.Request.Context(), time.Now(), responseCacheStatus(c.Writer.Header())) +} + +// ServerTimingResponseHeader builds the extra header map required by WebSocket upgrades. +func ServerTimingResponseHeader(c *gin.Context) http.Header { + value := ServerTimingHeaderValue(c) + if value == "" { + return nil + } + return http.Header{servertiming.HeaderName: []string{value}} +} + +func isAdminUIRequest(c *gin.Context) bool { + if c == nil || c.Request == nil || c.Request.URL == nil { + return false + } + if strings.TrimSpace(c.GetHeader(servertiming.AdminUIHeader)) == "1" { + return true + } + path := strings.TrimSpace(c.Request.URL.Path) + return path == "/api/v1/admin" || strings.HasPrefix(path, "/api/v1/admin/") +} + +func responseCacheStatus(header http.Header) string { + for _, name := range []string{snapshotCacheHeader, usageCacheHeader} { + switch strings.ToLower(strings.TrimSpace(header.Get(name))) { + case "hit": + return "hit" + case "miss": + return "miss" + case "bypass": + return "bypass" + } + } + return "bypass" +} diff --git a/backend/internal/server/middleware/server_timing_test.go b/backend/internal/server/middleware/server_timing_test.go new file mode 100644 index 0000000000..c064840ece --- /dev/null +++ b/backend/internal/server/middleware/server_timing_test.go @@ -0,0 +1,188 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" + "github.com/gin-gonic/gin" +) + +func runServerTimingRequest( + t *testing.T, + enabled bool, + path string, + marker string, + role string, + handler gin.HandlerFunc, +) *httptest.ResponseRecorder { + t.Helper() + gin.SetMode(gin.TestMode) + engine := gin.New() + engine.Use(ServerTiming(enabled)) + engine.Any("/*path", func(c *gin.Context) { + if role != "" { + c.Set(string(ContextKeyUserRole), role) + } + handler(c) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, path, nil) + if marker != "" { + request.Header.Set(servertiming.AdminUIHeader, marker) + } + engine.ServeHTTP(recorder, request) + return recorder +} + +func TestServerTimingScopesAndRoleGate(t *testing.T) { + tests := []struct { + name string + enabled bool + path string + marker string + role string + wantHeader bool + }{ + {name: "disabled", enabled: false, path: "/api/v1/admin/users", role: "admin"}, + {name: "admin API path", enabled: true, path: "/api/v1/admin/users", role: "admin", wantHeader: true}, + {name: "shared API marked by admin UI", enabled: true, path: "/api/v1/groups/available", marker: "1", role: "admin", wantHeader: true}, + {name: "non admin role", enabled: true, path: "/api/v1/groups/available", marker: "1", role: "user"}, + {name: "unauthenticated public request", enabled: true, path: "/api/v1/settings/public", marker: "1"}, + {name: "unmarked shared API", enabled: true, path: "/api/v1/groups/available", role: "admin"}, + {name: "invalid marker", enabled: true, path: "/api/v1/groups/available", marker: "true", role: "admin"}, + {name: "admin prefix boundary", enabled: true, path: "/api/v1/administrator", role: "admin"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := runServerTimingRequest(t, tt.enabled, tt.path, tt.marker, tt.role, func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + header := recorder.Header().Get(servertiming.HeaderName) + if tt.wantHeader && header == "" { + t.Fatalf("%s header missing", servertiming.HeaderName) + } + if !tt.wantHeader && header != "" { + t.Fatalf("unexpected %s header: %q", servertiming.HeaderName, header) + } + if header != "" && (!strings.Contains(header, "total;dur=") || !strings.Contains(header, `cache;desc="bypass"`)) { + t.Fatalf("incomplete timing header: %q", header) + } + }) + } +} + +func TestServerTimingCollectorIsRequestScoped(t *testing.T) { + active := false + recorder := runServerTimingRequest(t, true, "/api/v1/keys", "1", "admin", func(c *gin.Context) { + active = servertiming.Active(c.Request.Context()) + c.Status(http.StatusNoContent) + }) + if !active { + t.Fatal("collector was not attached to marked request context") + } + if recorder.Header().Get(servertiming.HeaderName) == "" { + t.Fatal("timing header missing from status-only response") + } +} + +func TestServerTimingFinalizesBeforeEarlyCommit(t *testing.T) { + recorder := runServerTimingRequest(t, true, "/api/v1/admin/stream", "", "admin", func(c *gin.Context) { + c.Status(http.StatusAccepted) + c.Writer.WriteHeaderNow() + }) + if got := recorder.Header().Get(servertiming.HeaderName); got == "" { + t.Fatal("timing header was not written before response commit") + } +} + +func TestServerTimingFinalizesOnFlush(t *testing.T) { + recorder := runServerTimingRequest(t, true, "/api/v1/admin/export", "", "admin", func(c *gin.Context) { + c.Writer.Flush() + }) + if got := recorder.Header().Get(servertiming.HeaderName); got == "" { + t.Fatal("timing header was not written before stream flush") + } +} + +func TestServerTimingStatusResponses(t *testing.T) { + tests := []struct { + name string + status int + }{ + {name: "not modified", status: http.StatusNotModified}, + {name: "internal error", status: http.StatusInternalServerError}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := runServerTimingRequest(t, true, "/api/v1/admin/test", "", "admin", func(c *gin.Context) { + c.Status(tt.status) + }) + if recorder.Code != tt.status { + t.Fatalf("status = %d, want %d", recorder.Code, tt.status) + } + if got := recorder.Header().Get(servertiming.HeaderName); got == "" { + t.Fatalf("timing header missing from status %d response", tt.status) + } + }) + } +} + +func TestServerTimingResponseWriterUnwraps(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + baseWriter := c.Writer + writer := &serverTimingResponseWriter{ResponseWriter: baseWriter} + if got := writer.Unwrap(); got != baseWriter { + t.Fatalf("Unwrap() = %T, want original Gin writer", got) + } +} + +func TestServerTimingCacheOutcome(t *testing.T) { + tests := []struct { + name string + headerName string + value string + want string + }{ + {name: "snapshot hit", headerName: snapshotCacheHeader, value: "hit", want: "hit"}, + {name: "usage miss", headerName: usageCacheHeader, value: "MISS", want: "miss"}, + {name: "invalid", headerName: snapshotCacheHeader, value: "stale", want: "bypass"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := runServerTimingRequest(t, true, "/api/v1/admin/dashboard", "", "admin", func(c *gin.Context) { + c.Header(tt.headerName, tt.value) + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + want := `cache;desc="` + tt.want + `"` + if got := recorder.Header().Get(servertiming.HeaderName); !strings.Contains(got, want) { + t.Fatalf("timing header %q does not contain %q", got, want) + } + }) + } +} + +func TestServerTimingResponseHeaderForWebSocket(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ws/qps", nil) + collector := servertiming.New(time.Now()) + c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector)) + c.Set(string(ContextKeyUserRole), "admin") + + header := ServerTimingResponseHeader(c) + if header.Get(servertiming.HeaderName) == "" { + t.Fatal("WebSocket response header missing timing value") + } + + c.Set(string(ContextKeyUserRole), "user") + if got := ServerTimingResponseHeader(c); got != nil { + t.Fatalf("non-admin WebSocket received timing header: %#v", got) + } +} diff --git a/backend/internal/server/router.go b/backend/internal/server/router.go index 3d86373779..5fc70149fe 100644 --- a/backend/internal/server/router.go +++ b/backend/internal/server/router.go @@ -60,6 +60,7 @@ func SetupRouter( } return nil })) + r.Use(middleware2.ServerTiming(cfg.Server.EnableServerTiming)) // Serve embedded frontend with settings injection if available if web.HasEmbeddedFrontend() { diff --git a/backend/internal/server/routes/admin.go b/backend/internal/server/routes/admin.go index 0d7e2a505a..5132022d4a 100644 --- a/backend/internal/server/routes/admin.go +++ b/backend/internal/server/routes/admin.go @@ -399,6 +399,7 @@ func registerGrokOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) { grok.POST("/oauth/exchange-code", h.Admin.GrokOAuth.ExchangeCode) grok.POST("/oauth/refresh-token", h.Admin.GrokOAuth.RefreshToken) grok.POST("/oauth/create-from-oauth", h.Admin.GrokOAuth.CreateAccountFromOAuth) + grok.POST("/sso-to-oauth", h.Admin.GrokOAuth.CreateAccountsFromSSO) grok.POST("/accounts/:id/refresh", h.Admin.GrokOAuth.RefreshAccountToken) grok.GET("/accounts/:id/quota", h.Admin.GrokOAuth.QueryQuota) grok.POST("/accounts/:id/reset-quota", h.Admin.GrokOAuth.ResetQuota) diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index 7960137604..45db227e58 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -84,6 +84,22 @@ func RegisterGatewayRoutes( }, }) } + videoEditHandler := func(c *gin.Context) { + if getGroupPlatform(c) == service.PlatformGrok { + h.OpenAIGateway.GrokVideoEdit(c) + return + } + service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) + c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Videos API is not supported for this platform"}}) + } + videoExtensionHandler := func(c *gin.Context) { + if getGroupPlatform(c) == service.PlatformGrok { + h.OpenAIGateway.GrokVideoExtension(c) + return + } + service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) + c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Videos API is not supported for this platform"}}) + } // API网关（Claude API兼容） gateway := r.Group("/v1") gateway.Use(bodyLimit) @@ -185,6 +201,8 @@ func RegisterGatewayRoutes( gateway.DELETE("/images/batches/:id", h.BatchImage.DeleteRecord) gateway.DELETE("/images/batches/:id/outputs", h.BatchImage.DeleteOutputs) gateway.POST("/videos/generations", videoGenerationHandler) + gateway.POST("/videos/edits", videoEditHandler) + gateway.POST("/videos/extensions", videoExtensionHandler) gateway.GET("/videos/:request_id", videoStatusHandler) } @@ -252,6 +270,8 @@ func RegisterGatewayRoutes( r.POST("/images/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler) r.POST("/images/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler) r.POST("/videos/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoGenerationHandler) + r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoEditHandler) + r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoExtensionHandler) r.GET("/videos/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoStatusHandler) // Antigravity 模型列表 diff --git a/backend/internal/server/routes/gateway_test.go b/backend/internal/server/routes/gateway_test.go index 6b15fbfa9b..65c6824440 100644 --- a/backend/internal/server/routes/gateway_test.go +++ b/backend/internal/server/routes/gateway_test.go @@ -123,6 +123,10 @@ func TestGatewayRoutesGrokImagesAndVideosPathsAreRegistered(t *testing.T) { "/images/edits", "/v1/videos/generations", "/videos/generations", + "/v1/videos/edits", + "/videos/edits", + "/v1/videos/extensions", + "/videos/extensions", } { req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"grok-imagine","prompt":"draw a cat"}`)) req.Header.Set("Content-Type", "application/json") @@ -156,6 +160,10 @@ func TestGatewayRoutesNonGrokVideosAreRejectedAtPlatformGate(t *testing.T) { }{ {http.MethodPost, "/v1/videos/generations", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`}, {http.MethodPost, "/videos/generations", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`}, + {http.MethodPost, "/v1/videos/edits", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`}, + {http.MethodPost, "/videos/edits", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`}, + {http.MethodPost, "/v1/videos/extensions", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`}, + {http.MethodPost, "/videos/extensions", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`}, {http.MethodGet, "/v1/videos/request-123", ""}, {http.MethodGet, "/videos/request-123", ""}, } { diff --git a/backend/internal/server/routes/payment.go b/backend/internal/server/routes/payment.go index 7c54770028..9434a0b76d 100644 --- a/backend/internal/server/routes/payment.go +++ b/backend/internal/server/routes/payment.go @@ -28,7 +28,6 @@ func RegisterPaymentRoutes( authenticated.GET("/config", paymentHandler.GetPaymentConfig) authenticated.GET("/checkout-info", paymentHandler.GetCheckoutInfo) authenticated.GET("/plans", paymentHandler.GetPlans) - authenticated.GET("/channels", paymentHandler.GetChannels) authenticated.GET("/limits", paymentHandler.GetLimits) orders := authenticated.Group("/orders") diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 62cb5a2e66..3e67fa5982 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -83,6 +83,8 @@ type Account struct { type OpenAIEndpointCapability string +const openAILongContextBillingEnabledKey = "openai_long_context_billing_enabled" + const ( OpenAIEndpointCapabilityChatCompletions OpenAIEndpointCapability = "chat_completions" OpenAIEndpointCapabilityEmbeddings OpenAIEndpointCapability = "embeddings" @@ -1192,6 +1194,14 @@ func (a *Account) IsOpenAI() bool { return a.Platform == PlatformOpenAI } +func (a *Account) IsOpenAILongContextBillingEnabled() bool { + if a == nil || !a.IsOpenAI() || a.Extra == nil { + return false + } + enabled, ok := a.Extra[openAILongContextBillingEnabledKey].(bool) + return ok && enabled +} + func (a *Account) IsAnthropic() bool { return a.Platform == PlatformAnthropic } @@ -1251,6 +1261,9 @@ func (a *Account) GetOpenAIRefreshToken() string { return a.GetCredential("refresh_token") } +// GetGrokBaseURL selects the upstream used by Grok text and Responses traffic. +// Grok media traffic has a different transport contract and must use +// GetGrokMediaBaseURL instead. func (a *Account) GetGrokBaseURL() string { if !a.IsGrok() { return "" @@ -1260,6 +1273,10 @@ func (a *Account) GetGrokBaseURL() string { if strings.TrimSpace(baseURL) == "" || isOfficialGrokAPIBaseURL(baseURL) { return xai.DefaultCLIBaseURL } + if _, err := xai.ValidateTrustedBaseURL(baseURL); err == nil { + return baseURL + } + return xai.DefaultCLIBaseURL } if baseURL != "" { return baseURL @@ -1267,12 +1284,45 @@ func (a *Account) GetGrokBaseURL() string { return xai.DefaultBaseURL } +// GetGrokMediaBaseURL selects the upstream used by Grok Imagine APIs. +// +// OAuth text requests need the CLI subscription proxy, but that proxy has a +// smaller request-body limit than the official Imagine API. Media requests can +// contain large base64 inputs, so default OAuth accounts must use api.x.ai. +// API-key accounts and explicit unsafe development overrides retain their +// configured base URL. +func (a *Account) GetGrokMediaBaseURL() string { + if !a.IsGrok() { + return "" + } + if !a.IsGrokOAuth() { + return a.GetGrokBaseURL() + } + + baseURL := a.GetCredential("base_url") + if strings.TrimSpace(baseURL) == "" || isOfficialGrokAPIBaseURL(baseURL) || isOfficialGrokCLIBaseURL(baseURL) { + return xai.DefaultBaseURL + } + if _, err := xai.ValidateTrustedBaseURL(baseURL); err == nil { + return baseURL + } + return xai.DefaultBaseURL +} + func isOfficialGrokAPIBaseURL(raw string) bool { + return isOfficialGrokBaseURL(raw, xai.DefaultBaseURL) +} + +func isOfficialGrokCLIBaseURL(raw string) bool { + return isOfficialGrokBaseURL(raw, xai.DefaultCLIBaseURL) +} + +func isOfficialGrokBaseURL(raw, expected string) bool { parsed, err := url.Parse(strings.TrimSpace(raw)) if err != nil || parsed == nil || parsed.Opaque != "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { return false } - defaultURL, err := url.Parse(xai.DefaultBaseURL) + defaultURL, err := url.Parse(expected) if err != nil { return false } diff --git a/backend/internal/service/account_base_url_test.go b/backend/internal/service/account_base_url_test.go index 0ffaa21ae4..59f53db3db 100644 --- a/backend/internal/service/account_base_url_test.go +++ b/backend/internal/service/account_base_url_test.go @@ -266,7 +266,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) { expected: "https://api.x.ai:8443/v1", }, { - name: "oauth explicit custom base_url remains supported", + name: "oauth explicit custom base_url stays pinned to CLI proxy by default", account: Account{ Type: AccountTypeOAuth, Platform: PlatformGrok, @@ -274,7 +274,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) { "base_url": "https://custom.example.com/v1", }, }, - expected: "https://custom.example.com/v1", + expected: xai.DefaultCLIBaseURL, }, { name: "API key without base_url uses official credit-backed API", @@ -293,3 +293,117 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) { }) } } + +func TestGetGrokBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + account := Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://custom.example.com/v1", + }, + } + + require.Equal(t, "https://custom.example.com/v1", account.GetGrokBaseURL()) +} + +func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) { + tests := []struct { + name string + account Account + expected string + }{ + { + name: "oauth without base_url uses official media API", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{}, + }, + expected: xai.DefaultBaseURL, + }, + { + name: "oauth stored CLI proxy uses official media API", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": xai.DefaultCLIBaseURL, + }, + }, + expected: xai.DefaultBaseURL, + }, + { + name: "oauth stored CLI proxy variant uses official media API", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "HTTPS://CLI-CHAT-PROXY.GROK.COM:443/%76%31/", + }, + }, + expected: xai.DefaultBaseURL, + }, + { + name: "oauth legacy official API remains on official media API", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": xai.DefaultBaseURL, + }, + }, + expected: xai.DefaultBaseURL, + }, + { + name: "oauth untrusted custom base_url is pinned to official media API", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://custom.example.com/v1", + }, + }, + expected: xai.DefaultBaseURL, + }, + { + name: "API key retains its configured media API", + account: Account{ + Type: AccountTypeAPIKey, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://grok.example.com/v1", + }, + }, + expected: "https://grok.example.com/v1", + }, + { + name: "non-Grok account has no Grok media base URL", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformOpenAI, + Credentials: map[string]any{}, + }, + expected: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.expected, tt.account.GetGrokMediaBaseURL()) + }) + } +} + +func TestGetGrokMediaBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + account := Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://custom.example.com/v1", + }, + } + + require.Equal(t, "https://custom.example.com/v1", account.GetGrokMediaBaseURL()) +} diff --git a/backend/internal/service/account_long_context_billing_test.go b/backend/internal/service/account_long_context_billing_test.go new file mode 100644 index 0000000000..709559d932 --- /dev/null +++ b/backend/internal/service/account_long_context_billing_test.go @@ -0,0 +1,290 @@ +//go:build unit + +package service + +import ( + "context" + "net/http" + "testing" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/stretchr/testify/require" +) + +func TestAccountIsOpenAILongContextBillingEnabled(t *testing.T) { + tests := []struct { + name string + account *Account + want bool + }{ + {name: "nil account is disabled", account: nil, want: false}, + {name: "non OpenAI account is disabled", account: &Account{Platform: PlatformGrok}, want: false}, + {name: "missing extra defaults disabled", account: &Account{Platform: PlatformOpenAI}, want: false}, + {name: "missing key defaults disabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{}}, want: false}, + {name: "explicit true is enabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_long_context_billing_enabled": true}}, want: true}, + {name: "explicit false is disabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_long_context_billing_enabled": false}}, want: false}, + {name: "malformed value is disabled", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_long_context_billing_enabled": "false"}}, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, tt.account.IsOpenAILongContextBillingEnabled()) + }) + } +} + +func TestNormalizeOpenAILongContextBillingExtra(t *testing.T) { + t.Run("OpenAI missing key persists disabled default", func(t *testing.T) { + extra, err := normalizeOpenAILongContextBillingExtra(PlatformOpenAI, nil) + + require.NoError(t, err) + require.Equal(t, false, extra["openai_long_context_billing_enabled"]) + }) + + t.Run("OpenAI explicit false is preserved", func(t *testing.T) { + extra, err := normalizeOpenAILongContextBillingExtra(PlatformOpenAI, map[string]any{"openai_long_context_billing_enabled": false}) + + require.NoError(t, err) + require.Equal(t, false, extra["openai_long_context_billing_enabled"]) + }) + + t.Run("OpenAI malformed value is rejected", func(t *testing.T) { + _, err := normalizeOpenAILongContextBillingExtra(PlatformOpenAI, map[string]any{"openai_long_context_billing_enabled": "false"}) + + require.Error(t, err) + require.Equal(t, http.StatusBadRequest, infraerrors.Code(err)) + }) + + t.Run("non OpenAI extra is unchanged", func(t *testing.T) { + extra, err := normalizeOpenAILongContextBillingExtra(PlatformGrok, nil) + + require.NoError(t, err) + require.Nil(t, extra) + }) + + t.Run("non OpenAI malformed value is ignored", func(t *testing.T) { + extra := map[string]any{openAILongContextBillingEnabledKey: "provider-owned"} + normalized, err := normalizeOpenAILongContextBillingExtra(PlatformAnthropic, extra) + + require.NoError(t, err) + require.Equal(t, extra, normalized) + }) +} + +type longContextBillingRepoStub struct { + accountRepoStub + account *Account + accounts []*Account + createdAccount *Account + updateExtraCalls int + bulkUpdateCalls int +} + +func (r *longContextBillingRepoStub) Create(_ context.Context, account *Account) error { + account.ID = 1 + r.account = account + r.createdAccount = account + return nil +} + +func (r *longContextBillingRepoStub) GetByID(_ context.Context, _ int64) (*Account, error) { + return r.account, nil +} + +func (r *longContextBillingRepoStub) GetByIDs(_ context.Context, _ []int64) ([]*Account, error) { + if r.accounts != nil { + return r.accounts, nil + } + if r.account == nil { + return nil, nil + } + return []*Account{r.account}, nil +} + +func (r *longContextBillingRepoStub) Update(_ context.Context, account *Account) error { + r.account = account + return nil +} + +func (r *longContextBillingRepoStub) UpdateExtra(_ context.Context, _ int64, _ map[string]any) error { + r.updateExtraCalls++ + return nil +} + +func (r *longContextBillingRepoStub) BulkUpdate(_ context.Context, _ []int64, _ AccountBulkUpdate) (int64, error) { + r.bulkUpdateCalls++ + return 1, nil +} + +func TestAdminServiceCreateAccountDefaultsOpenAILongContextBillingDisabled(t *testing.T) { + repo := &longContextBillingRepoStub{} + svc := &adminServiceImpl{accountRepo: repo} + + account, err := svc.CreateAccount(context.Background(), &CreateAccountInput{ + Name: "openai-account", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "test"}, + SkipDefaultGroupBind: true, + }) + + require.NoError(t, err) + require.Same(t, account, repo.createdAccount) + require.Equal(t, false, account.Extra[openAILongContextBillingEnabledKey]) +} + +func TestAdminServiceCreateAccountRejectsMalformedOpenAILongContextBillingValue(t *testing.T) { + repo := &longContextBillingRepoStub{} + svc := &adminServiceImpl{accountRepo: repo} + + account, err := svc.CreateAccount(context.Background(), &CreateAccountInput{ + Platform: PlatformOpenAI, + Extra: map[string]any{openAILongContextBillingEnabledKey: "false"}, + }) + + require.Nil(t, account) + require.Equal(t, http.StatusBadRequest, infraerrors.Code(err)) + require.Nil(t, repo.createdAccount) +} + +func TestAdminServiceUpdateAccountPreservesOpenAILongContextBillingOptOutWhenOmitted(t *testing.T) { + repo := &longContextBillingRepoStub{account: &Account{ + ID: 1, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Extra: map[string]any{openAILongContextBillingEnabledKey: false}, + }} + svc := &adminServiceImpl{accountRepo: repo} + + account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: map[string]any{}}) + + require.NoError(t, err) + require.Equal(t, false, account.Extra[openAILongContextBillingEnabledKey]) +} + +func TestAdminServiceUpdateAccountAllowsExplicitCodexImportOptIn(t *testing.T) { + repo := &longContextBillingRepoStub{account: &Account{ + ID: 1, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{"access_token": "old-token"}, + Extra: map[string]any{ + openAILongContextBillingEnabledKey: false, + "import_source": "codex_session", + }, + }} + svc := &adminServiceImpl{accountRepo: repo} + + account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{ + Credentials: map[string]any{"access_token": "new-token"}, + Extra: map[string]any{ + openAILongContextBillingEnabledKey: true, + "import_source": "codex_session", + }, + }) + + require.NoError(t, err) + require.Equal(t, true, account.Extra[openAILongContextBillingEnabledKey]) +} + +func TestAdminServiceUpdateAccountAllowsExplicitOptInOutsideCodexImport(t *testing.T) { + repo := &longContextBillingRepoStub{account: &Account{ + ID: 1, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Extra: map[string]any{ + openAILongContextBillingEnabledKey: false, + "import_source": "codex_session", + }, + }} + svc := &adminServiceImpl{accountRepo: repo} + + account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: map[string]any{ + openAILongContextBillingEnabledKey: true, + "import_source": "codex_session", + }}) + + require.NoError(t, err) + require.Equal(t, true, account.Extra[openAILongContextBillingEnabledKey]) +} + +func TestAdminServiceUpdateAccountRejectsMalformedOpenAILongContextBillingValue(t *testing.T) { + repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformOpenAI}} + svc := &adminServiceImpl{accountRepo: repo} + + account, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: map[string]any{ + openAILongContextBillingEnabledKey: 1, + }}) + + require.Nil(t, account) + require.Equal(t, http.StatusBadRequest, infraerrors.Code(err)) +} + +func TestAdminServiceUpdateAccountExtraRejectsMalformedOpenAILongContextBillingValue(t *testing.T) { + repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformOpenAI}} + svc := &adminServiceImpl{accountRepo: repo} + + err := svc.UpdateAccountExtra(context.Background(), 1, map[string]any{ + openAILongContextBillingEnabledKey: "true", + }) + + require.Equal(t, http.StatusBadRequest, infraerrors.Code(err)) + require.Zero(t, repo.updateExtraCalls) +} + +func TestAdminServiceUpdateAccountExtraAllowsProviderOwnedValueForNonOpenAIAccount(t *testing.T) { + repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformAnthropic}} + svc := &adminServiceImpl{accountRepo: repo} + + err := svc.UpdateAccountExtra(context.Background(), 1, map[string]any{ + openAILongContextBillingEnabledKey: "provider-owned", + }) + + require.NoError(t, err) + require.Equal(t, 1, repo.updateExtraCalls) +} + +func TestAdminServiceBulkUpdateAccountsRejectsMalformedOpenAILongContextBillingValue(t *testing.T) { + repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformOpenAI}} + svc := &adminServiceImpl{accountRepo: repo} + + result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1}, + Extra: map[string]any{openAILongContextBillingEnabledKey: []bool{true}}, + }) + + require.Nil(t, result) + require.Equal(t, http.StatusBadRequest, infraerrors.Code(err)) + require.Zero(t, repo.bulkUpdateCalls) +} + +func TestAdminServiceBulkUpdateAccountsAllowsProviderOwnedValueForNonOpenAIAccounts(t *testing.T) { + repo := &longContextBillingRepoStub{account: &Account{ID: 1, Platform: PlatformGrok}} + svc := &adminServiceImpl{accountRepo: repo} + + result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1}, + Extra: map[string]any{openAILongContextBillingEnabledKey: []string{"provider-owned"}}, + }) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, repo.bulkUpdateCalls) +} + +func TestAdminServiceBulkUpdateAccountsRejectsMalformedValueForMixedTargetsIncludingOpenAI(t *testing.T) { + repo := &longContextBillingRepoStub{accounts: []*Account{ + {ID: 1, Platform: PlatformGrok}, + {ID: 2, Platform: PlatformOpenAI}, + }} + svc := &adminServiceImpl{accountRepo: repo} + + result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + Extra: map[string]any{openAILongContextBillingEnabledKey: "malformed"}, + }) + + require.Nil(t, result) + require.Equal(t, http.StatusBadRequest, infraerrors.Code(err)) + require.Zero(t, repo.bulkUpdateCalls) +} diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index e62f459e08..222b3f8a4d 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -582,8 +582,13 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account c.Writer.Header().Set("X-Accel-Buffering", "no") c.Writer.Flush() - // Create OpenAI Responses API payload - payload := createOpenAITestPayload(testModelID, isOAuth) + // Create OpenAI Responses API payload. OAuth accounts use ChatGPT Codex + // upstream and must apply the same model normalization as real forwarding. + upstreamTestModelID := testModelID + if isOAuth { + upstreamTestModelID = normalizeOpenAIModelForUpstream(credentialAccount, testModelID) + } + payload := createOpenAITestPayload(upstreamTestModelID, isOAuth) payloadBytes, _ := json.Marshal(payload) // Send test_start event @@ -664,7 +669,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * testModelID := strings.TrimSpace(modelID) if testModelID == "" { - testModelID = "grok-4.3" + testModelID = grokDefaultResponsesModel } if mapped := strings.TrimSpace(account.GetMappedModel(testModelID)); mapped != "" { testModelID = mapped diff --git a/backend/internal/service/account_test_service_grok_test.go b/backend/internal/service/account_test_service_grok_test.go index 497224b713..4b0890ff44 100644 --- a/backend/internal/service/account_test_service_grok_test.go +++ b/backend/internal/service/account_test_service_grok_test.go @@ -80,6 +80,47 @@ func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testin require.Contains(t, rec.Body.String(), `"type":"test_complete"`) } +func TestAccountTestService_TestAccountConnection_GrokDefaultsEmptyModelTo45(t *testing.T) { + gin.SetMode(gin.TestMode) + + account := &Account{ + ID: 16, + Name: "grok-oauth-default-model", + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "grok-access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" + + "data: {\"type\":\"response.completed\"}\n\n", + )), + }} + svc := &AccountTestService{ + accountRepo: repo, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + httpUpstream: upstream, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/16/test", nil) + + err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeDefault) + + require.NoError(t, err) + require.Equal(t, grokDefaultResponsesModel, gjson.GetBytes(upstream.lastBody, "model").String()) + require.Contains(t, recorder.Body.String(), `"model":"grok-4.5"`) +} + func TestAccountTestService_Grok429PersistsRateLimitReset(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/account_test_service_openai_test.go b/backend/internal/service/account_test_service_openai_test.go index af28085123..083d882ea7 100644 --- a/backend/internal/service/account_test_service_openai_test.go +++ b/backend/internal/service/account_test_service_openai_test.go @@ -137,6 +137,34 @@ func TestAccountTestService_OpenAISuccessPersistsSnapshotFromHeaders(t *testing. require.Contains(t, recorder.Body.String(), "test_complete") } +func TestAccountTestService_OpenAIOAuthTestNormalizesGPT56Alias(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := newTestContext() + + resp := newJSONResponse(http.StatusOK, "") + resp.Body = io.NopCloser(strings.NewReader(`data: {"type":"response.completed"} + +`)) + + upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}} + svc := &AccountTestService{httpUpstream: upstream} + account := &Account{ + ID: 90, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{"access_token": "test-token"}, + } + + err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.6", "", "") + require.NoError(t, err) + require.Len(t, upstream.requests, 1) + + body, err := io.ReadAll(upstream.requests[0].Body) + require.NoError(t, err) + require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(body, "model").String()) +} + func TestAccountTestService_OpenAIShadowUsesParentCredentialsAndShadowModel(t *testing.T) { gin.SetMode(gin.TestMode) ctx, recorder := newTestContext() diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index 281122d4f0..7b42a1d1c5 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -111,6 +111,8 @@ const ( apiQueryMaxJitter = 800 * time.Millisecond // 用量查询最大随机延迟 windowStatsCacheTTL = 1 * time.Minute openAIProbeCacheTTL = 10 * time.Minute + grokProbeRetryTTL = 1 * time.Minute + grokFreeQuotaWindow = 24 * time.Hour openAICodexProbeVersion = "0.144.1" ) @@ -122,6 +124,7 @@ type UsageCache struct { apiFlight singleflight.Group // 防止同一账号的并发请求击穿缓存（Anthropic） antigravityFlight singleflight.Group // 防止同一 Antigravity 账号的并发请求击穿缓存 openAIProbeCache sync.Map // accountID -> time.Time + grokProbeCache sync.Map // accountID -> last billing probe attempt } // NewUsageCache 创建 UsageCache 实例 @@ -196,15 +199,19 @@ type UsageInfo struct { AntigravityQuota map[string]*AntigravityModelQuota `json:"antigravity_quota,omitempty"` // Grok / xAI 被动额度快照 - GrokRequestQuota *xai.QuotaWindow `json:"grok_request_quota,omitempty"` - GrokTokenQuota *xai.QuotaWindow `json:"grok_token_quota,omitempty"` - GrokRetryAfterSeconds *int `json:"grok_retry_after_seconds,omitempty"` - GrokEntitlementStatus string `json:"grok_entitlement_status,omitempty"` - GrokQuotaSnapshotState string `json:"grok_quota_snapshot_state,omitempty"` - GrokLastQuotaProbeAt string `json:"grok_last_quota_probe_at,omitempty"` - GrokLastHeadersSeenAt string `json:"grok_last_headers_seen_at,omitempty"` - GrokLastStatusCode int `json:"grok_last_status_code,omitempty"` - GrokLocalUsage *WindowStats `json:"grok_local_usage,omitempty"` + GrokRequestQuota *xai.QuotaWindow `json:"grok_request_quota,omitempty"` + GrokTokenQuota *xai.QuotaWindow `json:"grok_token_quota,omitempty"` + GrokRetryAfterSeconds *int `json:"grok_retry_after_seconds,omitempty"` + GrokEntitlementStatus string `json:"grok_entitlement_status,omitempty"` + GrokQuotaSnapshotState string `json:"grok_quota_snapshot_state,omitempty"` + GrokLastQuotaProbeAt string `json:"grok_last_quota_probe_at,omitempty"` + GrokLastHeadersSeenAt string `json:"grok_last_headers_seen_at,omitempty"` + GrokLastStatusCode int `json:"grok_last_status_code,omitempty"` + GrokLocalUsage *WindowStats `json:"grok_local_usage,omitempty"` + GrokLocalUsage24h *WindowStats `json:"grok_local_usage_24h,omitempty"` + GrokLocalUsage7d *WindowStats `json:"grok_local_usage_7d,omitempty"` + GrokLocalUsageMonthly *WindowStats `json:"grok_local_usage_monthly,omitempty"` + GrokBilling *xai.BillingSummary `json:"grok_billing,omitempty"` // Antigravity 账号级信息 SubscriptionTier string `json:"subscription_tier,omitempty"` // 归一化订阅等级: FREE/PRO/ULTRA/UNKNOWN @@ -287,6 +294,7 @@ type AccountUsageService struct { geminiQuotaService *GeminiQuotaService antigravityQuotaFetcher *AntigravityQuotaFetcher grokQuotaFetcher *GrokQuotaFetcher + grokQuotaService *GrokQuotaService openAIQuotaService *OpenAIQuotaService cache *UsageCache identityCache IdentityCache @@ -301,6 +309,7 @@ func NewAccountUsageService( geminiQuotaService *GeminiQuotaService, antigravityQuotaFetcher *AntigravityQuotaFetcher, grokQuotaFetcher *GrokQuotaFetcher, + grokQuotaService *GrokQuotaService, openAIQuotaService *OpenAIQuotaService, cache *UsageCache, identityCache IdentityCache, @@ -313,6 +322,7 @@ func NewAccountUsageService( geminiQuotaService: geminiQuotaService, antigravityQuotaFetcher: antigravityQuotaFetcher, grokQuotaFetcher: grokQuotaFetcher, + grokQuotaService: grokQuotaService, openAIQuotaService: openAIQuotaService, cache: cache, identityCache: identityCache, @@ -358,8 +368,8 @@ func (s *AccountUsageService) GetUsage(ctx context.Context, accountID int64, for } if account.Platform == PlatformGrok { - usage, err := s.getGrokUsage(ctx, account) - if err == nil { + usage, err := s.getGrokUsage(ctx, account, forceProbe) + if err == nil && usage != nil && usage.Error == "" { s.tryClearRecoverableAccountError(ctx, account) } return usage, err @@ -930,11 +940,21 @@ func (s *AccountUsageService) getAntigravityUsage(ctx context.Context, account * return usage, nil } -func (s *AccountUsageService) getGrokUsage(ctx context.Context, account *Account) (*UsageInfo, error) { +func (s *AccountUsageService) getGrokUsage(ctx context.Context, account *Account, force bool) (*UsageInfo, error) { if s.grokQuotaFetcher == nil { now := time.Now() return &UsageInfo{UpdatedAt: &now}, nil } + var billingProbeResult *GrokQuotaProbeResult + if account != nil && account.IsGrokOAuth() && s.grokQuotaService != nil && (force || grokBillingSnapshotNeedsRefresh(account, time.Now())) && s.shouldProbeGrokBilling(account.ID, time.Now(), force) { + result, err := s.grokQuotaService.ProbeBilling(ctx, account.ID) + if err == nil && result != nil && result.Billing != nil { + billingProbeResult = result + mergeAccountExtra(account, map[string]any{grokBillingExtraKey: result.Billing}) + } else if err != nil && force { + return nil, err + } + } usage := s.grokQuotaFetcher.BuildUsageInfo(account) if usage.GrokQuotaSnapshotState == "" { if usage.ErrorCode == "quota_unknown" { @@ -944,9 +964,20 @@ func (s *AccountUsageService) getGrokUsage(ctx context.Context, account *Account } } - if s.usageLogRepo != nil && account != nil { - if stats, err := s.usageLogRepo.GetAccountTodayStats(ctx, account.ID); err == nil && stats != nil { - usage.GrokLocalUsage = windowStatsFromAccountStats(stats) + if account != nil { + if s.usageLogRepo != nil { + if stats, err := s.usageLogRepo.GetAccountTodayStats(ctx, account.ID); err == nil && stats != nil { + usage.GrokLocalUsage = windowStatsFromAccountStats(stats) + } + } + if billingProbeResult != nil { + usage.GrokLocalUsage24h = billingProbeResult.LocalUsage24h + usage.GrokLocalUsage7d = billingProbeResult.LocalUsage7d + usage.GrokLocalUsageMonthly = billingProbeResult.LocalUsageMonthly + } else if s.usageLogRepo != nil { + usage.GrokLocalUsage24h, usage.GrokLocalUsage7d, usage.GrokLocalUsageMonthly = grokLocalUsageForQuota( + ctx, s.usageLogRepo, account.ID, usage.GrokBilling, time.Now().UTC(), + ) } } @@ -954,6 +985,110 @@ func (s *AccountUsageService) getGrokUsage(ctx context.Context, account *Account return usage, nil } +func grokLocalUsageForQuota( + ctx context.Context, + repo UsageLogRepository, + accountID int64, + billing *xai.BillingSummary, + now time.Time, +) (*WindowStats, *WindowStats, *WindowStats) { + if grokBillingHasAuthoritativeQuota(billing) { + weekly, monthly := grokLocalUsageForBilling(ctx, repo, accountID, billing, now) + return nil, weekly, monthly + } + return grokLocalUsage24h(ctx, repo, accountID, now), nil, nil +} + +func grokLocalUsage24h(ctx context.Context, repo UsageLogRepository, accountID int64, now time.Time) *WindowStats { + if repo == nil || accountID <= 0 { + return nil + } + start := now.UTC().Add(-grokFreeQuotaWindow) + stats, err := repo.GetAccountWindowStats(ctx, accountID, start) + if err != nil { + slog.Warn("grok_rolling_24h_usage_query_failed", "account_id", accountID, "window_start", start, "error", err) + return nil + } + return windowStatsFromAccountStats(stats) +} + +func grokLocalUsageForBilling( + ctx context.Context, + repo UsageLogRepository, + accountID int64, + billing *xai.BillingSummary, + now time.Time, +) (*WindowStats, *WindowStats) { + var weekly *WindowStats + var monthly *WindowStats + if repo == nil || accountID <= 0 { + return weekly, monthly + } + if start, ok := currentGrokBillingWindow(billing, true, now); ok { + if stats, err := repo.GetAccountWindowStats(ctx, accountID, start); err == nil { + weekly = windowStatsFromAccountStats(stats) + } else { + slog.Warn("grok_window_usage_query_failed", "account_id", accountID, "window_start", start, "error", err) + } + } + if start, ok := currentGrokBillingWindow(billing, false, now); ok { + if stats, err := repo.GetAccountWindowStats(ctx, accountID, start); err == nil { + monthly = windowStatsFromAccountStats(stats) + } else { + slog.Warn("grok_monthly_usage_query_failed", "account_id", accountID, "window_start", start, "error", err) + } + } + return weekly, monthly +} + +func currentGrokBillingWindow(billing *xai.BillingSummary, weekly bool, now time.Time) (time.Time, bool) { + if billing == nil { + return time.Time{}, false + } + startRaw, endRaw := billing.BillingPeriodStart, billing.BillingPeriodEnd + if weekly { + if billing.PeriodType != "weekly" { + return time.Time{}, false + } + startRaw, endRaw = billing.PeriodStart, billing.PeriodEnd + } + start, startErr := parseTime(strings.TrimSpace(startRaw)) + end, endErr := parseTime(strings.TrimSpace(endRaw)) + if startErr != nil || endErr != nil || now.Before(start) || !now.Before(end) { + return time.Time{}, false + } + return start, true +} + +func grokBillingSnapshotNeedsRefresh(account *Account, now time.Time) bool { + if account == nil { + return false + } + billing, err := grokBillingSnapshotFromExtra(account.Extra) + if err != nil || billing == nil || billing.Partial || len(billing.FailedWindows) > 0 { + return true + } + stamp := strings.TrimSpace(billing.UpdatedAt) + if stamp == "" { + stamp = strings.TrimSpace(billing.FetchedAt) + } + updatedAt, err := parseTime(stamp) + return err != nil || now.Sub(updatedAt) >= openAIProbeCacheTTL +} + +func (s *AccountUsageService) shouldProbeGrokBilling(accountID int64, now time.Time, force bool) bool { + if force || s == nil || s.cache == nil || accountID <= 0 { + return true + } + if cached, ok := s.cache.grokProbeCache.Load(accountID); ok { + if ts, ok := cached.(time.Time); ok && now.Sub(ts) < grokProbeRetryTTL { + return false + } + } + s.cache.grokProbeCache.Store(accountID, now) + return true +} + // recalcAntigravityRemainingSeconds 重新计算 Antigravity UsageInfo 中各窗口的 RemainingSeconds // 用于从缓存取出时更新倒计时，避免返回过时的剩余秒数 func recalcAntigravityRemainingSeconds(info *UsageInfo) { diff --git a/backend/internal/service/admin_account.go b/backend/internal/service/admin_account.go index 52e5ce719b..8cb6d8e63b 100644 --- a/backend/internal/service/admin_account.go +++ b/backend/internal/service/admin_account.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "log/slog" + "maps" "net/http" "strconv" "strings" @@ -68,7 +69,65 @@ func normalizeAccountConcurrency(platform, accountType string, concurrency int) return concurrency } +// ValidateOpenAILongContextBillingExtra validates the OpenAI account billing flag when present. +func ValidateOpenAILongContextBillingExtra(platform string, extra map[string]any) error { + if platform != PlatformOpenAI { + return nil + } + raw, exists := extra[openAILongContextBillingEnabledKey] + if !exists { + return nil + } + if _, ok := raw.(bool); !ok { + return infraerrors.BadRequest( + "OPENAI_LONG_CONTEXT_BILLING_INVALID", + "openai_long_context_billing_enabled must be a boolean", + ) + } + return nil +} + +func normalizeOpenAILongContextBillingExtra(platform string, extra map[string]any) (map[string]any, error) { + if platform != PlatformOpenAI { + return extra, nil + } + if err := ValidateOpenAILongContextBillingExtra(platform, extra); err != nil { + return nil, err + } + + normalized := maps.Clone(extra) + if normalized == nil { + normalized = make(map[string]any, 1) + } + _, exists := normalized[openAILongContextBillingEnabledKey] + if !exists { + normalized[openAILongContextBillingEnabledKey] = false + } + return normalized, nil +} + +func normalizeOpenAILongContextBillingUpdateExtra(account *Account, input *UpdateAccountInput) (map[string]any, error) { + normalized, err := normalizeOpenAILongContextBillingExtra(account.Platform, input.Extra) + if err != nil || account.Platform != PlatformOpenAI { + return normalized, err + } + + _, provided := input.Extra[openAILongContextBillingEnabledKey] + current, hasCurrent := account.Extra[openAILongContextBillingEnabledKey].(bool) + if !provided { + if hasCurrent { + normalized[openAILongContextBillingEnabledKey] = current + } + } + return normalized, nil +} + func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccountInput) (*Account, error) { + accountExtra, err := normalizeOpenAILongContextBillingExtra(input.Platform, input.Extra) + if err != nil { + return nil, err + } + // 绑定分组 groupIDs := input.GroupIDs // 如果没有指定分组,自动绑定对应平台的默认分组 @@ -103,7 +162,7 @@ func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccou Platform: input.Platform, Type: input.Type, Credentials: input.Credentials, - Extra: input.Extra, + Extra: accountExtra, ProxyID: input.ProxyID, Concurrency: normalizeAccountConcurrency(input.Platform, input.Type, input.Concurrency), Priority: input.Priority, @@ -183,6 +242,13 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U if err != nil { return nil, err } + var normalizedExtra map[string]any + if input.Extra != nil { + normalizedExtra, err = normalizeOpenAILongContextBillingUpdateExtra(account, input) + if err != nil { + return nil, err + } + } // 安全/身份不变量(影子账号):通用更新路径被 edit/re-auth/refresh/batch 共用, // 必须在此守住,否则仅在创建时的保证可被这些路径绕过。 if account.IsCredentialShadow() { @@ -238,10 +304,10 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U // 保留配额用量字段，防止编辑账号时意外重置 for _, key := range []string{"quota_used", "quota_daily_used", "quota_daily_start", "quota_weekly_used", "quota_weekly_start"} { if v, ok := account.Extra[key]; ok { - input.Extra[key] = v + normalizedExtra[key] = v } } - account.Extra = input.Extra + account.Extra = normalizedExtra if account.Platform == PlatformAntigravity && wasOveragesEnabled && !account.IsOveragesEnabled() { delete(account.Extra, "antigravity_credits_overages") // 清理旧版 overages 运行态 // 清除 AICredits 限流 key @@ -353,6 +419,15 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U // UpdateAccountExtra 仅对 Extra JSONB 做 key 级合并，避免覆盖其它运行态键 // （如 model_rate_limits / passive_usage_* 等）。 func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, updates map[string]any) error { + if _, exists := updates[openAILongContextBillingEnabledKey]; exists { + account, err := s.accountRepo.GetByID(ctx, id) + if err != nil { + return err + } + if err := ValidateOpenAILongContextBillingExtra(account.Platform, updates); err != nil { + return err + } + } if len(updates) == 0 { return nil } @@ -386,16 +461,28 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp } needMixedChannelCheck := input.GroupIDs != nil && !input.SkipMixedChannelCheck + _, hasLongContextBillingUpdate := input.Extra[openAILongContextBillingEnabledKey] // 预取所有目标账号，供凭据守卫/代理守卫/混合渠道检查共用，避免多次 DB 查询。 var cachedTargets []*Account - if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck { + if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate { loaded, err := s.accountRepo.GetByIDs(ctx, input.AccountIDs) if err != nil { return nil, err } cachedTargets = loaded } + if hasLongContextBillingUpdate { + for _, account := range cachedTargets { + if account == nil || account.Platform != PlatformOpenAI { + continue + } + if err := ValidateOpenAILongContextBillingExtra(account.Platform, input.Extra); err != nil { + return nil, err + } + break + } + } // 影子账号绝不持有凭据:批量更新携带凭据时,目标中不得含影子(外审 G5,与单账号 // UpdateAccount 守卫对齐)。覆盖显式 IDs 与 filter 解析出的 IDs(此处 AccountIDs 已解析完成)。 @@ -745,6 +832,9 @@ func (s *adminServiceImpl) CreateShadow(ctx context.Context, parentID int64, opt Priority: priority, Concurrency: concurrency, Schedulable: true, + Extra: map[string]any{ + openAILongContextBillingEnabledKey: parent.IsOpenAILongContextBillingEnabled(), + }, } // 5. 持久化（Create 填充 shadow.ID）。并发竞态:预查(步骤2)放行后另一请求抢先建成,本次会撞 diff --git a/backend/internal/service/admin_service_spark_shadow_test.go b/backend/internal/service/admin_service_spark_shadow_test.go index 6b4017207a..0eda0d93c7 100644 --- a/backend/internal/service/admin_service_spark_shadow_test.go +++ b/backend/internal/service/admin_service_spark_shadow_test.go @@ -157,6 +157,39 @@ func TestCreateShadow(t *testing.T) { require.Error(t, err) } +func TestCreateShadowInheritsParentEffectiveOpenAILongContextBillingValue(t *testing.T) { + tests := []struct { + name string + parentExtra map[string]any + want bool + }{ + {name: "missing parent value defaults disabled", want: false}, + {name: "explicit parent opt-out is inherited", parentExtra: map[string]any{openAILongContextBillingEnabledKey: false}, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := newSparkShadowRepoStub() + svc := &adminServiceImpl{accountRepo: repo} + parent := &Account{ + Name: "parent", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{"access_token": "token"}, + Extra: tt.parentExtra, + } + require.NoError(t, repo.Create(context.Background(), parent)) + + shadow, err := svc.CreateShadow(context.Background(), parent.ID, ShadowOptions{Name: "shadow"}) + + require.NoError(t, err) + require.Equal(t, tt.want, shadow.Extra[openAILongContextBillingEnabledKey]) + require.Equal(t, tt.want, shadow.IsOpenAILongContextBillingEnabled()) + }) + } +} + // TestCreateShadow_BindGroups は BindGroups の後置呼び出しを検証する。 // 影子账号が指定グループに属し、ListSchedulableByGroupID で取得可能であること。 func TestCreateShadow_BindGroups(t *testing.T) { diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 43d7f22d78..7fa69d41ea 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -153,14 +153,15 @@ type UsageTokens struct { // CostBreakdown 费用明细 type CostBreakdown struct { - InputCost float64 - OutputCost float64 - ImageOutputCost float64 - CacheCreationCost float64 - CacheReadCost float64 - TotalCost float64 - ActualCost float64 // 应用倍率后的实际费用 - BillingMode string // 计费模式（"token"/"per_request"/"image"），由 CalculateCostUnified 填充 + InputCost float64 + OutputCost float64 + ImageOutputCost float64 + CacheCreationCost float64 + CacheReadCost float64 + TotalCost float64 + ActualCost float64 // 应用倍率后的实际费用 + BillingMode string // 计费模式（"token"/"per_request"/"image"），由 CalculateCostUnified 填充 + LongContextBillingApplied bool } // ErrModelPricingUnavailable indicates that none of the configured pricing @@ -865,16 +866,17 @@ func (s *BillingService) GetModelPricingWithChannel(model string, channelPricing // CostInput 统一计费输入 type CostInput struct { - Ctx context.Context - Model string - GroupID *int64 // 用于渠道定价查找 - Tokens UsageTokens - RequestCount int // 按次计费时使用 - SizeTier string // 按次/图片模式的层级标签（"1K","2K","4K","HD" 等） - RateMultiplier float64 - ServiceTier string // "priority","flex","" 等 - Resolver *ModelPricingResolver // 定价解析器 - Resolved *ResolvedPricing // 可选：预解析的定价结果（避免重复 Resolve 调用） + Ctx context.Context + Model string + GroupID *int64 // 用于渠道定价查找 + Tokens UsageTokens + RequestCount int // 按次计费时使用 + SizeTier string // 按次/图片模式的层级标签（"1K","2K","4K","HD" 等） + RateMultiplier float64 + ServiceTier string // "priority","flex","" 等 + Resolver *ModelPricingResolver // 定价解析器 + Resolved *ResolvedPricing // 可选：预解析的定价结果（避免重复 Resolve 调用） + LongContextBillingEnabled *bool } // CalculateCostUnified 统一计费入口，支持三种计费模式。 @@ -882,7 +884,18 @@ type CostInput struct { func (s *BillingService) CalculateCostUnified(input CostInput) (*CostBreakdown, error) { if input.Resolver == nil { // 无 Resolver，回退到旧路径 - return s.calculateCostInternal(input.Model, input.Tokens, input.RateMultiplier, input.ServiceTier, nil) + applyLongContextBilling := true + if input.LongContextBillingEnabled != nil { + applyLongContextBilling = *input.LongContextBillingEnabled + } + return s.calculateCostInternalWithPolicy( + input.Model, + input.Tokens, + input.RateMultiplier, + input.ServiceTier, + nil, + applyLongContextBilling, + ) } // 优先使用预解析结果，避免重复 Resolve 调用 @@ -929,6 +942,9 @@ func (s *BillingService) calculateTokenCost(resolved *ResolvedPricing, input Cos // 长上下文定价仅在无区间定价时应用（区间定价已包含上下文分层） applyLongCtx := len(resolved.Intervals) == 0 + if input.LongContextBillingEnabled != nil { + applyLongCtx = applyLongCtx && *input.LongContextBillingEnabled + } return s.computeTokenBreakdown(pricing, input.Tokens, input.RateMultiplier, input.ServiceTier, applyLongCtx), nil } @@ -969,7 +985,10 @@ func (s *BillingService) computeTokenBreakdown( tierMultiplier = serviceTierCostMultiplier(serviceTier) } - if applyLongCtx && s.shouldApplySessionLongContextPricing(tokens, pricing) { + longContextPricingEligible := applyLongCtx && s.shouldApplySessionLongContextPricing(tokens, pricing) + var baselineCost *CostBreakdown + if longContextPricingEligible { + baselineCost = s.computeTokenBreakdown(pricing, tokens, rateMultiplier, serviceTier, false) inputPrice *= pricing.LongContextInputMultiplier outputPrice *= pricing.LongContextOutputMultiplier // 缓存读取本质上是输入侧的复用，应与 input 一同应用长上下文倍率； @@ -1033,6 +1052,7 @@ func (s *BillingService) computeTokenBreakdown( bd.TotalCost = bd.InputCost + bd.OutputCost + bd.ImageOutputCost + bd.CacheCreationCost + bd.CacheReadCost bd.ActualCost = bd.TotalCost * rateMultiplier + bd.LongContextBillingApplied = baselineCost != nil && bd.ActualCost > baselineCost.ActualCost return bd } @@ -1092,7 +1112,28 @@ func (s *BillingService) CalculateCostWithServiceTier(model string, tokens Usage return s.calculateCostInternal(model, tokens, rateMultiplier, serviceTier, nil) } +func (s *BillingService) calculateCostWithServiceTierPolicy( + model string, + tokens UsageTokens, + rateMultiplier float64, + serviceTier string, + longContextBillingEnabled bool, +) (*CostBreakdown, error) { + return s.calculateCostInternalWithPolicy(model, tokens, rateMultiplier, serviceTier, nil, longContextBillingEnabled) +} + func (s *BillingService) calculateCostInternal(model string, tokens UsageTokens, rateMultiplier float64, serviceTier string, channelPricing *ChannelModelPricing) (*CostBreakdown, error) { + return s.calculateCostInternalWithPolicy(model, tokens, rateMultiplier, serviceTier, channelPricing, true) +} + +func (s *BillingService) calculateCostInternalWithPolicy( + model string, + tokens UsageTokens, + rateMultiplier float64, + serviceTier string, + channelPricing *ChannelModelPricing, + longContextBillingEnabled bool, +) (*CostBreakdown, error) { var pricing *ModelPricing var err error if channelPricing != nil { @@ -1104,8 +1145,7 @@ func (s *BillingService) calculateCostInternal(model string, tokens UsageTokens, return nil, err } - // 旧路径始终检查长上下文定价（无区间定价概念） - return s.computeTokenBreakdown(pricing, tokens, rateMultiplier, serviceTier, true), nil + return s.computeTokenBreakdown(pricing, tokens, rateMultiplier, serviceTier, longContextBillingEnabled), nil } func (s *BillingService) applyModelSpecificPricingPolicy(model string, pricing *ModelPricing) *ModelPricing { @@ -1236,13 +1276,14 @@ func (s *BillingService) CalculateCostWithLongContext(model string, tokens Usage // 合并成本 return &CostBreakdown{ - InputCost: inRangeCost.InputCost + outRangeCost.InputCost, - OutputCost: inRangeCost.OutputCost, - ImageOutputCost: inRangeCost.ImageOutputCost, - CacheCreationCost: inRangeCost.CacheCreationCost, - CacheReadCost: inRangeCost.CacheReadCost + outRangeCost.CacheReadCost, - TotalCost: inRangeCost.TotalCost + outRangeCost.TotalCost, - ActualCost: inRangeCost.ActualCost + outRangeCost.ActualCost, + InputCost: inRangeCost.InputCost + outRangeCost.InputCost, + OutputCost: inRangeCost.OutputCost, + ImageOutputCost: inRangeCost.ImageOutputCost, + CacheCreationCost: inRangeCost.CacheCreationCost, + CacheReadCost: inRangeCost.CacheReadCost + outRangeCost.CacheReadCost, + TotalCost: inRangeCost.TotalCost + outRangeCost.TotalCost, + ActualCost: inRangeCost.ActualCost + outRangeCost.ActualCost, + LongContextBillingApplied: outRangeCost.ActualCost > 0, }, nil } diff --git a/backend/internal/service/billing_service_test.go b/backend/internal/service/billing_service_test.go index 53014412fd..885da194e3 100644 --- a/backend/internal/service/billing_service_test.go +++ b/backend/internal/service/billing_service_test.go @@ -261,6 +261,23 @@ func TestCalculateCost_OpenAIGPT54LongContextAppliesWholeSessionMultipliers(t *t require.InDelta(t, expectedOutput, cost.OutputCost, 1e-10) require.InDelta(t, expectedInput+expectedOutput, cost.TotalCost, 1e-10) require.InDelta(t, expectedInput+expectedOutput, cost.ActualCost, 1e-10) + require.True(t, cost.LongContextBillingApplied) +} + +func TestCalculateCost_OpenAIGPT54LongContextMarkerRequiresActualCostIncrease(t *testing.T) { + svc := newTestBillingService() + + cost, err := svc.calculateCostWithServiceTierPolicy( + "gpt-5.4-2026-03-05", + UsageTokens{InputTokens: 300000}, + 0, + "", + true, + ) + + require.NoError(t, err) + require.Zero(t, cost.ActualCost) + require.False(t, cost.LongContextBillingApplied) } func TestCalculateCost_OpenAIGPT55ProUsesGPT55PricingPolicy(t *testing.T) { @@ -831,6 +848,17 @@ func TestCalculateCostWithLongContext_AboveThreshold_CacheBelowThreshold(t *test require.True(t, cost.ActualCost > normalCost.ActualCost, "长上下文费用应高于正常费用") } +func TestCalculateCostWithLongContext_MarkerRequiresActualCostIncrease(t *testing.T) { + svc := newTestBillingService() + tokens := UsageTokens{InputTokens: 300000} + + cost, err := svc.CalculateCostWithLongContext("claude-sonnet-4", tokens, 0, 200000, 2.0) + + require.NoError(t, err) + require.Zero(t, cost.ActualCost) + require.False(t, cost.LongContextBillingApplied) +} + func TestCalculateCostWithLongContext_DisabledThreshold(t *testing.T) { svc := newTestBillingService() diff --git a/backend/internal/service/channel_monitor_checker.go b/backend/internal/service/channel_monitor_checker.go index 7fb829a3cb..ad4058f9e6 100644 --- a/backend/internal/service/channel_monitor_checker.go +++ b/backend/internal/service/channel_monitor_checker.go @@ -13,6 +13,7 @@ import ( "strings" "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/tidwall/gjson" ) @@ -34,7 +35,7 @@ func newSSRFSafeHTTPClient(timeout time.Duration) *http.Client { TLSHandshakeTimeout: monitorTLSHandshakeTimeout, ResponseHeaderTimeout: monitorResponseHeaderTimeout, } - return &http.Client{Timeout: timeout, Transport: tr} + return &http.Client{Timeout: timeout, Transport: servertiming.WrapRoundTripper(tr)} } // CheckOptions 承载一次检测的自定义入参。 @@ -167,6 +168,7 @@ type providerAdapter struct { //nolint:gochecknoglobals // 适配器表是只读静态数据，初始化后不变更。 var providerAdapters = map[string]providerAdapter{ MonitorProviderOpenAI: providerOpenAIChatAdapter, + MonitorProviderGrok: providerGrokChatAdapter, MonitorProviderAnthropic: { buildPath: func(string) string { return providerAnthropicPath }, buildBody: func(model, prompt string) ([]byte, error) { @@ -204,20 +206,27 @@ var providerAdapters = map[string]providerAdapter{ } //nolint:gochecknoglobals // 适配器表是只读静态数据，初始化后不变更。 -var providerOpenAIChatAdapter = providerAdapter{ - buildPath: func(string) string { return providerOpenAIPath }, - buildBody: func(model, prompt string) ([]byte, error) { - return json.Marshal(map[string]any{ - "model": model, - "messages": []map[string]string{{"role": "user", "content": prompt}}, - "max_tokens": monitorChallengeMaxTokens, - "stream": false, - }) - }, - buildHeaders: func(apiKey string) map[string]string { - return map[string]string{"Authorization": "Bearer " + apiKey} - }, - textPath: "choices.0.message.content", +var providerOpenAIChatAdapter = newOpenAICompatibleChatAdapter(providerOpenAIPath) + +//nolint:gochecknoglobals // 适配器表是只读静态数据，初始化后不变更。 +var providerGrokChatAdapter = newOpenAICompatibleChatAdapter(providerGrokPath) + +func newOpenAICompatibleChatAdapter(path string) providerAdapter { + return providerAdapter{ + buildPath: func(string) string { return path }, + buildBody: func(model, prompt string) ([]byte, error) { + return json.Marshal(map[string]any{ + "model": model, + "messages": []map[string]string{{"role": "user", "content": prompt}}, + "max_tokens": monitorChallengeMaxTokens, + "stream": false, + }) + }, + buildHeaders: func(apiKey string) map[string]string { + return map[string]string{"Authorization": "Bearer " + apiKey} + }, + textPath: "choices.0.message.content", + } } //nolint:gochecknoglobals // 适配器表是只读静态数据，初始化后不变更。 @@ -407,8 +416,9 @@ func buildRequestBody(adapter providerAdapter, provider, apiMode, model, prompt var bodyMergeKeyDenyList = map[string]map[string]bool{ MonitorProviderOpenAI + ":" + MonitorAPIModeChatCompletions: {"model": true, "messages": true, "stream": true}, MonitorProviderOpenAI + ":" + MonitorAPIModeResponses: {"model": true, "instructions": true, "input": true, "stream": true}, - MonitorProviderAnthropic: {"model": true, "messages": true}, - MonitorProviderGemini: {"contents": true}, + MonitorProviderGrok: {"model": true, "messages": true, "stream": true}, + MonitorProviderAnthropic: {"model": true, "messages": true}, + MonitorProviderGemini: {"contents": true}, } func checkAPIMode(opts *CheckOptions) string { @@ -426,7 +436,7 @@ func bodyMergeDenyKey(provider, apiMode string) string { } func validateReplaceRequestBody(provider, apiMode string, body map[string]any) error { - if provider != MonitorProviderOpenAI { + if provider != MonitorProviderOpenAI && provider != MonitorProviderGrok { return nil } switch defaultAPIMode(apiMode) { @@ -527,6 +537,8 @@ var monitorAPIKeyPatterns = []struct { {regexp.MustCompile(`sk-ant-[A-Za-z0-9_-]{20,}`), "sk-ant-***REDACTED***"}, // OpenAI / Anthropic 通用 sk-: sk-xxxxxxx {regexp.MustCompile(`sk-[A-Za-z0-9-]{20,}`), "sk-***REDACTED***"}, + // xAI API Key：xai-xxxxxxx + {regexp.MustCompile(`xai-[A-Za-z0-9_-]{6,}`), "xai-***REDACTED***"}, // Gemini / Google API Key：固定前缀 + 35 位 {regexp.MustCompile(`AIza[A-Za-z0-9_-]{35}`), "AIza***REDACTED***"}, // JWT 三段式（Bearer 后常出现）：eyJxxx.eyJxxx.signature @@ -536,7 +548,7 @@ var monitorAPIKeyPatterns = []struct { // sanitizeErrorMessage 擦除错误/响应文本中可能泄露的 API key。 // 处理两类来源： // 1. URL query 中的 ?key= / ?api_key= 等（Go *url.Error 会回填完整 URL） -// 2. 上游 HTTP body 文本里直接出现的 sk-* / AIza* / JWT 等密钥碎片 +// 2. 上游 HTTP body 文本里直接出现的 sk-* / xai-* / AIza* / JWT 等密钥碎片 // // 注意：与 gemini_messages_compat_service.go 的 sanitizeUpstreamErrorMessage 关注点类似但参数集更广， // 监控模块独立维护，避免互相耦合。 diff --git a/backend/internal/service/channel_monitor_checker_body_test.go b/backend/internal/service/channel_monitor_checker_body_test.go index bba3d7dfb7..bcf7af0b98 100644 --- a/backend/internal/service/channel_monitor_checker_body_test.go +++ b/backend/internal/service/channel_monitor_checker_body_test.go @@ -64,6 +64,7 @@ type openAICaptureHandler struct { lastHeaders http.Header lastPath string status int + rawResponse string responsesLeadingReasoning bool } @@ -80,6 +81,10 @@ func (h *openAICaptureHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) } w.Header().Set("Content-Type", "application/json") w.WriteHeader(h.status) + if h.rawResponse != "" { + _, _ = w.Write([]byte(h.rawResponse)) + return + } answer := answerFromOpenAIRequest(parsed) if h.lastPath == providerOpenAIResponsesPath { @@ -190,6 +195,90 @@ func TestRunCheckForModel_OpenAI_DefaultChatRequest(t *testing.T) { } } +func TestGrokMonitorConfiguration(t *testing.T) { + if err := validateProvider(MonitorProviderGrok); err != nil { + t.Fatalf("grok provider should be supported: %v", err) + } + if got := normalizeMonitorPrimaryModel(MonitorProviderGrok, ""); got != MonitorDefaultGrokModel { + t.Fatalf("expected default Grok model %q, got %q", MonitorDefaultGrokModel, got) + } + if err := validateAPIMode(MonitorProviderGrok, MonitorAPIModeChatCompletions); err != nil { + t.Fatalf("grok chat_completions mode should be valid: %v", err) + } + if err := validateAPIMode(MonitorProviderGrok, MonitorAPIModeResponses); err == nil { + t.Fatal("grok responses mode should be rejected by channel monitoring") + } + if err := validateReplaceRequestBody(MonitorProviderGrok, MonitorAPIModeChatCompletions, map[string]any{}); err == nil { + t.Fatal("grok replace-mode body should require messages") + } +} + +func TestRunCheckForModel_Grok_DefaultChatRequest(t *testing.T) { + h := &openAICaptureHandler{} + endpoint := setupFakeOpenAI(t, h) + + res := runCheckForModel(context.Background(), MonitorProviderGrok, endpoint, "xai-key", MonitorDefaultGrokModel, nil) + + if res.Status != MonitorStatusOperational { + t.Fatalf("Grok request should pass challenge, got status=%s message=%q", res.Status, res.Message) + } + if res.LatencyMs == nil { + t.Fatal("Grok request should record latency") + } + if h.lastPath != providerGrokPath { + t.Fatalf("expected Grok chat completions path %q, got %q", providerGrokPath, h.lastPath) + } + if h.lastBody["model"] != MonitorDefaultGrokModel { + t.Errorf("Grok body should contain model=%s, got %v", MonitorDefaultGrokModel, h.lastBody["model"]) + } + if _, ok := h.lastBody["messages"]; !ok { + t.Error("Grok body should contain messages") + } + if h.lastBody["stream"] != false { + t.Errorf("Grok body should set stream=false, got %v", h.lastBody["stream"]) + } + if h.lastHeaders.Get("Authorization") != "Bearer xai-key" { + t.Errorf("expected Grok bearer auth header, got %q", h.lastHeaders.Get("Authorization")) + } +} + +func TestRunCheckForModel_Grok_UpstreamFailure(t *testing.T) { + h := &openAICaptureHandler{status: http.StatusTooManyRequests} + endpoint := setupFakeOpenAI(t, h) + + res := runCheckForModel(context.Background(), MonitorProviderGrok, endpoint, "xai-key", MonitorDefaultGrokModel, nil) + + if res.Status != MonitorStatusError { + t.Fatalf("Grok 429 should be recorded as error, got status=%s message=%q", res.Status, res.Message) + } + if !strings.Contains(res.Message, "upstream HTTP 429") { + t.Fatalf("Grok failure should preserve upstream status, got %q", res.Message) + } + if res.LatencyMs == nil { + t.Fatal("Grok failure should still record latency") + } +} + +func TestRunCheckForModel_Grok_RedactsXAIKeyFromUpstreamBody(t *testing.T) { + h := &openAICaptureHandler{ + status: http.StatusUnauthorized, + rawResponse: `{"error":{"message":"invalid API key xai-secret"}}`, + } + endpoint := setupFakeOpenAI(t, h) + + res := runCheckForModel(context.Background(), MonitorProviderGrok, endpoint, "request-key", MonitorDefaultGrokModel, nil) + + if res.Status != MonitorStatusError { + t.Fatalf("Grok upstream failure should be recorded as error, got %s", res.Status) + } + if strings.Contains(res.Message, "xai-secret") { + t.Fatalf("Grok error message leaked xAI key: %q", res.Message) + } + if !strings.Contains(res.Message, "xai-***REDACTED***") { + t.Fatalf("Grok error message should contain redaction marker, got %q", res.Message) + } +} + func TestRunCheckForModel_OpenAIResponses_DefaultRequest(t *testing.T) { h := &openAICaptureHandler{} endpoint := setupFakeOpenAI(t, h) diff --git a/backend/internal/service/channel_monitor_const.go b/backend/internal/service/channel_monitor_const.go index 61f9f79894..2ee8eabde8 100644 --- a/backend/internal/service/channel_monitor_const.go +++ b/backend/internal/service/channel_monitor_const.go @@ -47,6 +47,8 @@ const ( // providerOpenAIPath OpenAI Chat Completions 路径。 providerOpenAIPath = "/v1/chat/completions" + // providerGrokPath Grok OpenAI-compatible Chat Completions 路径。 + providerGrokPath = "/v1/chat/completions" // providerOpenAIResponsesPath OpenAI Responses API 路径。 providerOpenAIResponsesPath = "/v1/responses" // providerAnthropicPath Anthropic Messages 路径。 @@ -54,10 +56,14 @@ const ( // providerGeminiPathTemplate Gemini generateContent 路径模板（含 model 占位）。 providerGeminiPathTemplate = "/v1beta/models/%s:generateContent" - // MonitorProviderOpenAI / Anthropic / Gemini provider 字符串常量（也是 ent enum 的实际值）。 + // MonitorProviderOpenAI / Anthropic / Gemini / Grok provider 字符串常量（也是 ent enum 的实际值）。 MonitorProviderOpenAI = "openai" MonitorProviderAnthropic = "anthropic" MonitorProviderGemini = "gemini" + MonitorProviderGrok = "grok" + + // MonitorDefaultGrokModel 是新增 Grok 监控未显式指定模型时使用的轻量测活模型。 + MonitorDefaultGrokModel = "grok-4.5" // MonitorStatusOperational 等监控状态字符串常量（与 ent enum 一致）。 MonitorStatusOperational = "operational" @@ -112,13 +118,13 @@ var ( "CHANNEL_MONITOR_NOT_FOUND", "channel monitor not found", ) ErrChannelMonitorInvalidProvider = infraerrors.BadRequest( - "CHANNEL_MONITOR_INVALID_PROVIDER", "provider must be one of openai/anthropic/gemini", + "CHANNEL_MONITOR_INVALID_PROVIDER", "provider must be one of openai/anthropic/gemini/grok", ) ErrChannelMonitorInvalidAPIMode = infraerrors.BadRequest( "CHANNEL_MONITOR_INVALID_API_MODE", "api_mode must be chat_completions or responses; responses is only supported for openai", ) ErrChannelMonitorInvalidRequestBody = infraerrors.BadRequest( - "CHANNEL_MONITOR_INVALID_REQUEST_BODY", "openai replace-mode body_override must include non-empty messages for chat_completions or non-empty instructions and input for responses", + "CHANNEL_MONITOR_INVALID_REQUEST_BODY", "openai-compatible replace-mode body_override must include non-empty messages for chat_completions or non-empty instructions and input for responses", ) ErrChannelMonitorInvalidInterval = infraerrors.BadRequest( "CHANNEL_MONITOR_INVALID_INTERVAL", "interval_seconds must be in [15, 3600]", diff --git a/backend/internal/service/channel_monitor_service.go b/backend/internal/service/channel_monitor_service.go index 7b53bb20b0..b5dea22589 100644 --- a/backend/internal/service/channel_monitor_service.go +++ b/backend/internal/service/channel_monitor_service.go @@ -123,7 +123,7 @@ func (s *ChannelMonitorService) Create(ctx context.Context, p ChannelMonitorCrea APIMode: defaultAPIMode(p.APIMode), Endpoint: normalizeEndpoint(p.Endpoint), APIKey: encrypted, // 注意：传入 repository 时该字段为密文 - PrimaryModel: strings.TrimSpace(p.PrimaryModel), + PrimaryModel: normalizeMonitorPrimaryModel(p.Provider, p.PrimaryModel), ExtraModels: normalizeModels(p.ExtraModels), GroupName: strings.TrimSpace(p.GroupName), Enabled: p.Enabled, @@ -167,7 +167,7 @@ func validateCreateParams(p ChannelMonitorCreateParams) error { if strings.TrimSpace(p.APIKey) == "" { return ErrChannelMonitorMissingAPIKey } - if strings.TrimSpace(p.PrimaryModel) == "" { + if normalizeMonitorPrimaryModel(p.Provider, p.PrimaryModel) == "" { return ErrChannelMonitorMissingPrimaryModel } return nil @@ -486,8 +486,8 @@ func applyMonitorUpdate(existing *ChannelMonitor, p ChannelMonitorUpdateParams) if err := validateProvider(*p.Provider); err != nil { return err } + providerChanged = existing.Provider != *p.Provider existing.Provider = *p.Provider - providerChanged = true } if p.Endpoint != nil { if err := validateEndpoint(*p.Endpoint); err != nil { @@ -496,7 +496,13 @@ func applyMonitorUpdate(existing *ChannelMonitor, p ChannelMonitorUpdateParams) existing.Endpoint = normalizeEndpoint(*p.Endpoint) } if p.PrimaryModel != nil { - existing.PrimaryModel = strings.TrimSpace(*p.PrimaryModel) + primaryModel := normalizeMonitorPrimaryModel(existing.Provider, *p.PrimaryModel) + if primaryModel == "" { + return ErrChannelMonitorMissingPrimaryModel + } + existing.PrimaryModel = primaryModel + } else if providerChanged && existing.Provider == MonitorProviderGrok { + existing.PrimaryModel = MonitorDefaultGrokModel } if p.ExtraModels != nil { existing.ExtraModels = normalizeModels(*p.ExtraModels) diff --git a/backend/internal/service/channel_monitor_service_grok_test.go b/backend/internal/service/channel_monitor_service_grok_test.go new file mode 100644 index 0000000000..20c9db2666 --- /dev/null +++ b/backend/internal/service/channel_monitor_service_grok_test.go @@ -0,0 +1,85 @@ +//go:build unit + +package service + +import "testing" + +func TestApplyMonitorUpdate_ProviderOnlySwitchToGrokUsesDefaultModel(t *testing.T) { + grok := MonitorProviderGrok + existing := &ChannelMonitor{ + Provider: MonitorProviderOpenAI, + APIMode: MonitorAPIModeResponses, + PrimaryModel: "gpt-5", + IntervalSeconds: 60, + } + + err := applyMonitorUpdate(existing, ChannelMonitorUpdateParams{Provider: &grok}) + if err != nil { + t.Fatalf("provider-only switch to Grok failed: %v", err) + } + if existing.PrimaryModel != MonitorDefaultGrokModel { + t.Fatalf("expected Grok default model %q, got %q", MonitorDefaultGrokModel, existing.PrimaryModel) + } + if existing.APIMode != MonitorAPIModeChatCompletions { + t.Fatalf("expected Grok API mode %q, got %q", MonitorAPIModeChatCompletions, existing.APIMode) + } +} + +func TestApplyMonitorUpdate_SwitchToGrokPreservesExplicitModel(t *testing.T) { + grok := MonitorProviderGrok + explicitModel := "grok-4.3" + existing := &ChannelMonitor{ + Provider: MonitorProviderOpenAI, + APIMode: MonitorAPIModeChatCompletions, + PrimaryModel: "gpt-5", + IntervalSeconds: 60, + } + + err := applyMonitorUpdate(existing, ChannelMonitorUpdateParams{ + Provider: &grok, + PrimaryModel: &explicitModel, + }) + if err != nil { + t.Fatalf("switch to Grok with explicit model failed: %v", err) + } + if existing.PrimaryModel != explicitModel { + t.Fatalf("expected explicit model %q, got %q", explicitModel, existing.PrimaryModel) + } +} + +func TestApplyMonitorUpdate_SameGrokProviderDoesNotResetExistingModel(t *testing.T) { + grok := MonitorProviderGrok + existing := &ChannelMonitor{ + Provider: MonitorProviderGrok, + APIMode: MonitorAPIModeChatCompletions, + PrimaryModel: "grok-4.3", + IntervalSeconds: 60, + } + + err := applyMonitorUpdate(existing, ChannelMonitorUpdateParams{Provider: &grok}) + if err != nil { + t.Fatalf("same-provider Grok update failed: %v", err) + } + if existing.PrimaryModel != "grok-4.3" { + t.Fatalf("same-provider update reset existing model to %q", existing.PrimaryModel) + } +} + +func TestApplyMonitorUpdate_SwitchToGrokRejectsResponsesMode(t *testing.T) { + grok := MonitorProviderGrok + responses := MonitorAPIModeResponses + existing := &ChannelMonitor{ + Provider: MonitorProviderOpenAI, + APIMode: MonitorAPIModeChatCompletions, + PrimaryModel: "gpt-5", + IntervalSeconds: 60, + } + + err := applyMonitorUpdate(existing, ChannelMonitorUpdateParams{ + Provider: &grok, + APIMode: &responses, + }) + if err == nil { + t.Fatal("Grok responses mode should remain unsupported") + } +} diff --git a/backend/internal/service/channel_monitor_template_types.go b/backend/internal/service/channel_monitor_template_types.go index 03cd518d28..0b824d577d 100644 --- a/backend/internal/service/channel_monitor_template_types.go +++ b/backend/internal/service/channel_monitor_template_types.go @@ -55,7 +55,7 @@ var ( "CHANNEL_MONITOR_TEMPLATE_NOT_FOUND", "channel monitor request template not found", ) ErrChannelMonitorTemplateInvalidProvider = infraerrors.BadRequest( - "CHANNEL_MONITOR_TEMPLATE_INVALID_PROVIDER", "template provider must be one of openai/anthropic/gemini", + "CHANNEL_MONITOR_TEMPLATE_INVALID_PROVIDER", "template provider must be one of openai/anthropic/gemini/grok", ) ErrChannelMonitorTemplateInvalidAPIMode = infraerrors.BadRequest( "CHANNEL_MONITOR_TEMPLATE_INVALID_API_MODE", "template api_mode must be chat_completions or responses; responses is only supported for openai", diff --git a/backend/internal/service/channel_monitor_validate.go b/backend/internal/service/channel_monitor_validate.go index c5a4783b91..7740dc83b7 100644 --- a/backend/internal/service/channel_monitor_validate.go +++ b/backend/internal/service/channel_monitor_validate.go @@ -124,6 +124,16 @@ func normalizeModels(in []string) []string { return out } +// normalizeMonitorPrimaryModel applies the Grok health-check default while +// preserving the existing required-model behavior for every other provider. +func normalizeMonitorPrimaryModel(provider, model string) string { + model = strings.TrimSpace(model) + if model == "" && provider == MonitorProviderGrok { + return MonitorDefaultGrokModel + } + return model +} + // defaultAPIMode 空串归一为 chat_completions，保证历史数据与旧客户端兼容。 func defaultAPIMode(apiMode string) string { if strings.TrimSpace(apiMode) == "" { diff --git a/backend/internal/service/concurrency_service.go b/backend/internal/service/concurrency_service.go index f2f2aade89..df18379704 100644 --- a/backend/internal/service/concurrency_service.go +++ b/backend/internal/service/concurrency_service.go @@ -6,6 +6,7 @@ import ( "crypto/sha256" "encoding/binary" "encoding/hex" + "errors" "os" "strconv" "sync" @@ -13,6 +14,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "go.uber.org/zap" "golang.org/x/sync/singleflight" ) @@ -59,6 +61,131 @@ type APIKeyConcurrencyCache interface { GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) } +// OpenAIWSIngressLeaseCache owns the short-lived distributed lease used to +// bound live client WebSocket sessions. It is deliberately independent of the +// request-slot namespace: idle ingress connections do not occupy turn slots. +type OpenAIWSIngressLeaseCache interface { + AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error) + RefreshOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) (bool, error) + ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error +} + +const ( + openAIWSIngressLeaseTTL = 60 * time.Second + openAIWSIngressLeaseRefreshInterval = 20 * time.Second + openAIWSIngressLeaseOperationTO = 2 * time.Second +) + +var ErrOpenAIWSIngressLeaseLost = errors.New("openai websocket ingress lease lost") + +// OpenAIWSIngressLease keeps a Redis-backed ingress lease alive and cancels +// its context if Redis cannot confirm ownership for a full lease lifetime. +// Call Release on every handler exit to reclaim capacity immediately. +type OpenAIWSIngressLease struct { + ctx context.Context + cancel context.CancelCauseFunc + cache OpenAIWSIngressLeaseCache + apiKeyID int64 + leaseID string + + stopOnce sync.Once + stopCh chan struct{} + refreshDone chan struct{} +} + +func (l *OpenAIWSIngressLease) Context() context.Context { + if l == nil || l.ctx == nil { + return context.Background() + } + return l.ctx +} + +func (l *OpenAIWSIngressLease) Release() { + if l == nil { + return + } + l.stopOnce.Do(func() { + if l.stopCh != nil { + close(l.stopCh) + } + if l.cancel != nil { + l.cancel(nil) + } + if l.refreshDone != nil { + <-l.refreshDone + } + if l.cache == nil || l.apiKeyID <= 0 || l.leaseID == "" { + return + } + releaseCtx, releaseCancel := context.WithTimeout(context.Background(), openAIWSIngressLeaseOperationTO) + defer releaseCancel() + if err := l.cache.ReleaseOpenAIWSIngressLease(releaseCtx, l.apiKeyID, l.leaseID); err != nil { + logger.L().Warn("openai_ws_ingress_lease_release_failed", + zap.Int64("api_key_id", l.apiKeyID), + zap.Error(err), + ) + } + }) +} + +func (l *OpenAIWSIngressLease) refreshLoop() { + defer func() { + if l != nil && l.refreshDone != nil { + close(l.refreshDone) + } + }() + if l == nil || l.cache == nil { + return + } + ticker := time.NewTicker(openAIWSIngressLeaseRefreshInterval) + defer ticker.Stop() + lastConfirmedAt := time.Now() + for { + select { + case <-l.ctx.Done(): + return + case <-l.stopCh: + return + case <-ticker.C: + var lost bool + lastConfirmedAt, lost = l.refresh(lastConfirmedAt) + if lost { + l.cancel(ErrOpenAIWSIngressLeaseLost) + return + } + } + } +} + +// refresh confirms the lease is still owned. A missing member is an immediate +// lease loss; transient Redis errors are tolerated only for one full lease TTL. +func (l *OpenAIWSIngressLease) refresh(lastConfirmedAt time.Time) (time.Time, bool) { + refreshCtx, refreshCancel := context.WithTimeout(context.Background(), openAIWSIngressLeaseOperationTO) + owned, err := l.cache.RefreshOpenAIWSIngressLease(refreshCtx, l.apiKeyID, l.leaseID) + refreshCancel() + if err == nil && owned { + return time.Now(), false + } + if err == nil { + err = ErrOpenAIWSIngressLeaseLost + } + elapsed := time.Since(lastConfirmedAt) + logger.L().Warn("openai_ws_ingress_lease_refresh_failed", + zap.Int64("api_key_id", l.apiKeyID), + zap.Duration("unconfirmed_for", elapsed), + zap.Error(err), + ) + if errors.Is(err, ErrOpenAIWSIngressLeaseLost) || elapsed >= openAIWSIngressLeaseTTL { + logger.L().Error("openai_ws_ingress_lease_lost", + zap.Int64("api_key_id", l.apiKeyID), + zap.Duration("unconfirmed_for", elapsed), + zap.Error(err), + ) + return lastConfirmedAt, true + } + return lastConfirmedAt, false +} + var ( requestIDPrefix = initRequestIDPrefix() requestIDCounter atomic.Uint64 @@ -125,6 +252,47 @@ func NewConcurrencyService(cache ConcurrencyCache) *ConcurrencyService { return svc } +// AcquireOpenAIWSIngressLease atomically reserves one live ingress connection +// for an API key. A non-positive limit explicitly disables this protection. +func (s *ConcurrencyService) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int) (*OpenAIWSIngressLease, bool, error) { + if maxConnections <= 0 { + return nil, true, nil + } + if s == nil || s.cache == nil || apiKeyID <= 0 { + return nil, false, errors.New("openai websocket ingress lease cache is unavailable") + } + cache, ok := s.cache.(OpenAIWSIngressLeaseCache) + if !ok { + return nil, false, errors.New("openai websocket ingress lease cache is unsupported") + } + leaseID := generateRequestID() + baseCtx := context.Background() + if ctx != nil { + baseCtx = context.WithoutCancel(ctx) + } + acquireCtx, acquireCancel := context.WithTimeout(baseCtx, openAIWSIngressLeaseOperationTO) + acquired, err := cache.AcquireOpenAIWSIngressLease(acquireCtx, apiKeyID, maxConnections, leaseID) + acquireCancel() + if err != nil || !acquired { + return nil, acquired, err + } + if ctx == nil { + ctx = context.Background() + } + leaseCtx, leaseCancel := context.WithCancelCause(ctx) + lease := &OpenAIWSIngressLease{ + ctx: leaseCtx, + cancel: leaseCancel, + cache: cache, + apiKeyID: apiKeyID, + leaseID: leaseID, + stopCh: make(chan struct{}), + refreshDone: make(chan struct{}), + } + go lease.refreshLoop() + return lease, true, nil +} + // SetAccountLoadBatchCacheTTL 设置账号负载批量读取的极短 TTL 缓存；非正数表示禁用缓存。 func (s *ConcurrencyService) SetAccountLoadBatchCacheTTL(ttl time.Duration) { if s == nil { diff --git a/backend/internal/service/concurrency_service_test.go b/backend/internal/service/concurrency_service_test.go index 3f358bbe6a..d079c7e60a 100644 --- a/backend/internal/service/concurrency_service_test.go +++ b/backend/internal/service/concurrency_service_test.go @@ -45,7 +45,47 @@ type stubConcurrencyCacheForTest struct { releasedAPIKeyRequestIDs []string } +type ingressLeaseCacheForTest struct { + stubConcurrencyCacheForTest + acquireIngressResult bool + acquireIngressErr error + acquireIngressFn func(context.Context, int64, int, string) (bool, error) + refreshIngressResult bool + refreshIngressErr error + refreshIngressFn func(context.Context, int64, string) (bool, error) + releaseIngressErr error + releaseIngressFn func(context.Context, int64, string) error + acquireIngressCalls int + refreshIngressCalls int + releaseIngressCalls int +} + +func (c *ingressLeaseCacheForTest) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error) { + c.acquireIngressCalls++ + if c.acquireIngressFn != nil { + return c.acquireIngressFn(ctx, apiKeyID, maxConnections, leaseID) + } + return c.acquireIngressResult, c.acquireIngressErr +} + +func (c *ingressLeaseCacheForTest) RefreshOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) (bool, error) { + c.refreshIngressCalls++ + if c.refreshIngressFn != nil { + return c.refreshIngressFn(ctx, apiKeyID, leaseID) + } + return c.refreshIngressResult, c.refreshIngressErr +} + +func (c *ingressLeaseCacheForTest) ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error { + c.releaseIngressCalls++ + if c.releaseIngressFn != nil { + return c.releaseIngressFn(ctx, apiKeyID, leaseID) + } + return c.releaseIngressErr +} + var _ ConcurrencyCache = (*stubConcurrencyCacheForTest)(nil) +var _ OpenAIWSIngressLeaseCache = (*ingressLeaseCacheForTest)(nil) func (c *stubConcurrencyCacheForTest) AcquireAccountSlot(_ context.Context, _ int64, _ int, _ string) (bool, error) { return c.acquireResult, c.acquireErr @@ -285,6 +325,114 @@ func TestGetAPIKeyConcurrencyBatch_Fallbacks(t *testing.T) { }) } +func TestAcquireOpenAIWSIngressLease(t *testing.T) { + t.Run("zero value release is safe", func(t *testing.T) { + var lease OpenAIWSIngressLease + require.NotPanics(t, lease.Release) + }) + + t.Run("disabled", func(t *testing.T) { + cache := &ingressLeaseCacheForTest{} + lease, acquired, err := NewConcurrencyService(cache).AcquireOpenAIWSIngressLease(nil, 1, 0) + require.NoError(t, err) + require.True(t, acquired) + require.Nil(t, lease) + require.Zero(t, cache.acquireIngressCalls) + }) + + t.Run("unsupported cache fails closed", func(t *testing.T) { + lease, acquired, err := NewConcurrencyService(&stubConcurrencyCacheForTest{}).AcquireOpenAIWSIngressLease(context.Background(), 1, 1) + require.Error(t, err) + require.False(t, acquired) + require.Nil(t, lease) + }) + + t.Run("capacity rejected", func(t *testing.T) { + cache := &ingressLeaseCacheForTest{acquireIngressResult: false} + lease, acquired, err := NewConcurrencyService(cache).AcquireOpenAIWSIngressLease(context.Background(), 1, 1) + require.NoError(t, err) + require.False(t, acquired) + require.Nil(t, lease) + }) + + t.Run("release returns capacity", func(t *testing.T) { + cache := &ingressLeaseCacheForTest{acquireIngressResult: true, refreshIngressResult: true} + lease, acquired, err := NewConcurrencyService(cache).AcquireOpenAIWSIngressLease(nil, 1, 1) + require.NoError(t, err) + require.True(t, acquired) + require.NotNil(t, lease) + lease.Release() + lease.Release() + require.Equal(t, 1, cache.releaseIngressCalls) + }) +} + +func TestOpenAIWSIngressLeaseRefreshLoss(t *testing.T) { + t.Run("missing lease is lost immediately", func(t *testing.T) { + cache := &ingressLeaseCacheForTest{refreshIngressResult: false} + lease := &OpenAIWSIngressLease{cache: cache, apiKeyID: 1, leaseID: "missing"} + _, lost := lease.refresh(time.Now()) + require.True(t, lost) + require.Equal(t, 1, cache.refreshIngressCalls) + }) + + t.Run("persistent redis errors lose lease after ttl", func(t *testing.T) { + cache := &ingressLeaseCacheForTest{refreshIngressErr: errors.New("redis unavailable")} + lease := &OpenAIWSIngressLease{cache: cache, apiKeyID: 1, leaseID: "unconfirmed"} + _, lost := lease.refresh(time.Now().Add(-openAIWSIngressLeaseTTL)) + require.True(t, lost) + require.Equal(t, 1, cache.refreshIngressCalls) + }) +} + +func TestOpenAIWSIngressLeaseReleaseWaitsForInFlightRefresh(t *testing.T) { + refreshStarted := make(chan struct{}) + allowRefresh := make(chan struct{}) + cache := &ingressLeaseCacheForTest{ + refreshIngressFn: func(context.Context, int64, string) (bool, error) { + close(refreshStarted) + <-allowRefresh + return true, nil + }, + } + ctx, cancel := context.WithCancelCause(context.Background()) + lease := &OpenAIWSIngressLease{ + ctx: ctx, + cancel: cancel, + cache: cache, + apiKeyID: 1, + leaseID: "in-flight-refresh", + stopCh: make(chan struct{}), + refreshDone: make(chan struct{}), + } + go func() { + defer close(lease.refreshDone) + _, _ = lease.refresh(time.Now()) + }() + <-refreshStarted + + released := make(chan struct{}) + go func() { + lease.Release() + close(released) + }() + + select { + case <-released: + t.Fatal("release returned before the in-flight refresh completed") + case <-time.After(20 * time.Millisecond): + } + require.Zero(t, cache.releaseIngressCalls) + + close(allowRefresh) + select { + case <-released: + case <-time.After(time.Second): + t.Fatal("release did not complete after the refresh returned") + } + require.Equal(t, 1, cache.releaseIngressCalls) +} + func TestGenerateRequestID_UsesStablePrefixAndMonotonicCounter(t *testing.T) { id1 := generateRequestID() id2 := generateRequestID() diff --git a/backend/internal/service/content_moderation.go b/backend/internal/service/content_moderation.go index 6d3b91d205..f633c8ad17 100644 --- a/backend/internal/service/content_moderation.go +++ b/backend/internal/service/content_moderation.go @@ -22,6 +22,7 @@ import ( infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" ) const ( @@ -561,7 +562,7 @@ func NewContentModerationService( userRepo: userRepo, authCacheInvalidator: authCacheInvalidator, emailService: emailService, - httpClient: &http.Client{}, + httpClient: servertiming.InstrumentClient(nil), workerCount: maxContentModerationWorkerCount, asyncQueue: make(chan contentModerationTask, maxContentModerationQueueSize), keyHealth: make(map[string]*contentModerationKeyHealth), diff --git a/backend/internal/service/crs_sync_long_context_billing_test.go b/backend/internal/service/crs_sync_long_context_billing_test.go new file mode 100644 index 0000000000..6439f08190 --- /dev/null +++ b/backend/internal/service/crs_sync_long_context_billing_test.go @@ -0,0 +1,169 @@ +//go:build unit + +package service + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type crsLongContextAccountRepo struct { + AccountRepository + accounts map[string]*Account + nextID int64 +} + +type crsOpenAILongContextSource struct { + collection string + credentials map[string]any + extra map[string]any +} + +func newCRSLongContextAccountRepo(existing ...*Account) *crsLongContextAccountRepo { + repo := &crsLongContextAccountRepo{accounts: make(map[string]*Account)} + for _, account := range existing { + if account == nil { + continue + } + crsID, _ := account.Extra["crs_account_id"].(string) + repo.accounts[crsID] = account + if account.ID > repo.nextID { + repo.nextID = account.ID + } + } + return repo +} + +func (r *crsLongContextAccountRepo) Create(_ context.Context, account *Account) error { + r.nextID++ + account.ID = r.nextID + crsID, _ := account.Extra["crs_account_id"].(string) + r.accounts[crsID] = account + return nil +} + +func (r *crsLongContextAccountRepo) Update(_ context.Context, account *Account) error { + crsID, _ := account.Extra["crs_account_id"].(string) + r.accounts[crsID] = account + return nil +} + +func (r *crsLongContextAccountRepo) GetByCRSAccountID(_ context.Context, crsID string) (*Account, error) { + return r.accounts[crsID], nil +} + +func (r *crsLongContextAccountRepo) ListShadowsByParent(_ context.Context, _ int64) ([]*Account, error) { + return nil, nil +} + +func TestCRSSyncOpenAILongContextBilling(t *testing.T) { + tests := []struct { + name string + collection string + credentials map[string]any + sourceExtra map[string]any + existingExtra map[string]any + wantAction string + wantEnabled bool + }{ + {name: "OAuth create defaults missing value disabled", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, wantAction: "created"}, + {name: "OAuth create preserves source true", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "created", wantEnabled: true}, + {name: "OAuth create preserves source false", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "created"}, + {name: "OAuth update defaults missing value disabled", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, existingExtra: map[string]any{"existing": true}, wantAction: "updated"}, + {name: "OAuth update preserves existing true when source omits value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "updated", wantEnabled: true}, + {name: "OAuth update preserves existing false when source omits value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "updated"}, + {name: "OAuth update preserves source true over existing false", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: true}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "updated", wantEnabled: true}, + {name: "OAuth update preserves source false over existing true", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: false}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "updated"}, + {name: "OAuth rejects malformed source value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, wantAction: "failed"}, + {name: "OAuth rejects malformed existing value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, wantAction: "failed"}, + {name: "OAuth update rejects malformed source value", collection: "openaiOAuthAccounts", credentials: map[string]any{"access_token": "oauth-token"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "failed"}, + {name: "API key create defaults missing value disabled", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, wantAction: "created"}, + {name: "API key create preserves source true", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "created", wantEnabled: true}, + {name: "API key create preserves source false", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "created"}, + {name: "API key update defaults missing value disabled", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, existingExtra: map[string]any{"existing": true}, wantAction: "updated"}, + {name: "API key update preserves existing true when source omits value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "updated", wantEnabled: true}, + {name: "API key update preserves existing false when source omits value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "updated"}, + {name: "API key update preserves source true over existing false", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: true}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: false}, wantAction: "updated", wantEnabled: true}, + {name: "API key update preserves source false over existing true", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: false}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "updated"}, + {name: "API key rejects malformed source value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, wantAction: "failed"}, + {name: "API key rejects malformed existing value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, wantAction: "failed"}, + {name: "API key update rejects malformed source value", collection: "openaiResponsesAccounts", credentials: map[string]any{"api_key": "sk-test"}, sourceExtra: map[string]any{openAILongContextBillingEnabledKey: "false"}, existingExtra: map[string]any{openAILongContextBillingEnabledKey: true}, wantAction: "failed"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + const crsID = "crs-openai-1" + var existing *Account + if tt.existingExtra != nil { + existingExtra := mergeMap(tt.existingExtra, map[string]any{"crs_account_id": crsID}) + accountType := AccountTypeOAuth + if tt.collection == "openaiResponsesAccounts" { + accountType = AccountTypeAPIKey + } + existing = &Account{ID: 41, Platform: PlatformOpenAI, Type: accountType, Extra: existingExtra} + } + repo := newCRSLongContextAccountRepo(existing) + result := runCRSOpenAILongContextSync(t, repo, crsOpenAILongContextSource{ + collection: tt.collection, + credentials: tt.credentials, + extra: tt.sourceExtra, + }) + + require.Len(t, result.Items, 1) + require.Equal(t, tt.wantAction, result.Items[0].Action) + if tt.wantAction == "failed" { + require.Contains(t, result.Items[0].Error, "openai_long_context_billing_enabled must be a boolean") + return + } + stored, ok := repo.accounts[crsID].Extra[openAILongContextBillingEnabledKey] + require.True(t, ok) + require.Equal(t, tt.wantEnabled, stored) + }) + } +} + +func runCRSOpenAILongContextSync(t *testing.T, repo AccountRepository, source crsOpenAILongContextSource) *SyncFromCRSResult { + t.Helper() + account := map[string]any{ + "kind": "openai", + "id": "crs-openai-1", + "name": "OpenAI CRS", + "isActive": true, + "schedulable": true, + "credentials": source.credentials, + } + if source.extra != nil { + account["extra"] = source.extra + } + + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + response.Header().Set("Content-Type", "application/json") + if request.URL.Path == "/web/auth/login" { + _, _ = response.Write([]byte(`{"success":true,"token":"admin-token"}`)) + return + } + require.Equal(t, "/admin/sync/export-accounts", request.URL.Path) + require.NoError(t, json.NewEncoder(response).Encode(map[string]any{ + "success": true, + "data": map[string]any{source.collection: []any{account}}, + })) + })) + t.Cleanup(server.Close) + + cfg := &config.Config{} + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + service := NewCRSSyncService(repo, nil, nil, nil, nil, cfg) + result, err := service.SyncFromCRS(context.Background(), SyncFromCRSInput{ + BaseURL: server.URL, + Username: "admin", + Password: "password", + }) + require.NoError(t, err) + return result +} diff --git a/backend/internal/service/crs_sync_service.go b/backend/internal/service/crs_sync_service.go index edf3cd43d2..d0abc74038 100644 --- a/backend/internal/service/crs_sync_service.go +++ b/backend/internal/service/crs_sync_service.go @@ -168,6 +168,7 @@ type crsOpenAIResponsesAccount struct { Status string `json:"status"` Proxy *crsProxy `json:"proxy"` Credentials map[string]any `json:"credentials"` + Extra map[string]any `json:"extra"` } type crsOpenAIOAuthAccount struct { @@ -632,6 +633,18 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput result.Items = append(result.Items, item) continue } + var existingExtra map[string]any + if existing != nil { + existingExtra = existing.Extra + } + extra, err = mergeCRSOpenAILongContextBillingExtra(existingExtra, extra) + if err != nil { + item.Action = "failed" + item.Error = err.Error() + result.Failed++ + result.Items = append(result.Items, item) + continue + } if existing == nil { if !shouldCreateAccount(src.ID, selectedSet) { @@ -670,7 +683,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput continue } - existing.Extra = mergeMap(existing.Extra, extra) + existing.Extra = extra existing.Name = defaultName(src.Name, src.ID) existing.Platform = PlatformOpenAI existing.Type = AccountTypeOAuth @@ -751,11 +764,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput concurrency := 3 status := mapCRSStatus(src.IsActive, src.Status) - extra := map[string]any{ - "crs_account_id": src.ID, - "crs_kind": src.Kind, - "crs_synced_at": now, + extra := make(map[string]any, len(src.Extra)+3) + for key, value := range src.Extra { + extra[key] = value } + extra["crs_account_id"] = src.ID + extra["crs_kind"] = src.Kind + extra["crs_synced_at"] = now existing, err := s.accountRepo.GetByCRSAccountID(ctx, src.ID) if err != nil { @@ -765,6 +780,18 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput result.Items = append(result.Items, item) continue } + var existingExtra map[string]any + if existing != nil { + existingExtra = existing.Extra + } + extra, err = mergeCRSOpenAILongContextBillingExtra(existingExtra, extra) + if err != nil { + item.Action = "failed" + item.Error = err.Error() + result.Failed++ + result.Items = append(result.Items, item) + continue + } if existing == nil { if !shouldCreateAccount(src.ID, selectedSet) { @@ -809,7 +836,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput continue } - existing.Extra = mergeMap(existing.Extra, extra) + existing.Extra = extra existing.Name = defaultName(src.Name, src.ID) existing.Platform = PlatformOpenAI existing.Type = AccountTypeAPIKey @@ -1098,6 +1125,10 @@ func mergeMap(existing map[string]any, updates map[string]any) map[string]any { return out } +func mergeCRSOpenAILongContextBillingExtra(existing, updates map[string]any) (map[string]any, error) { + return normalizeOpenAILongContextBillingExtra(PlatformOpenAI, mergeMap(existing, updates)) +} + func (s *CRSSyncService) mapOrCreateProxy(ctx context.Context, enabled bool, cached *[]Proxy, src *crsProxy, defaultName string) (*int64, error) { if !enabled || src == nil { return nil, nil diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 8a95915981..61ab3abd2f 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -947,6 +947,7 @@ func (s *GatewayService) buildRecordUsageLog( usageLog.CacheReadCost = cost.CacheReadCost usageLog.TotalCost = cost.TotalCost usageLog.ActualCost = cost.ActualCost + usageLog.LongContextBillingApplied = cost.LongContextBillingApplied } return usageLog diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 01101ef6b5..100d720659 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -26,6 +26,8 @@ const ( GrokMediaEndpointImagesGenerations GrokMediaEndpoint = "images_generations" GrokMediaEndpointImagesEdits GrokMediaEndpoint = "images_edits" GrokMediaEndpointVideosGenerations GrokMediaEndpoint = "videos_generations" + GrokMediaEndpointVideosEdits GrokMediaEndpoint = "videos_edits" + GrokMediaEndpointVideosExtensions GrokMediaEndpoint = "videos_extensions" GrokMediaEndpointVideoStatus GrokMediaEndpoint = "video_status" ) @@ -35,7 +37,7 @@ func (e GrokMediaEndpoint) RequiresRequestBody() bool { func (e GrokMediaEndpoint) IsGenerationRequest() bool { switch e { - case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits, GrokMediaEndpointVideosGenerations: + case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits, GrokMediaEndpointVideosGenerations, GrokMediaEndpointVideosEdits, GrokMediaEndpointVideosExtensions: return true default: return false @@ -274,6 +276,10 @@ func (e GrokMediaEndpoint) upstreamURL(baseURL, requestID string) (string, error return xai.BuildImagesEditsURL(baseURL) case GrokMediaEndpointVideosGenerations: return xai.BuildVideosGenerationsURL(baseURL) + case GrokMediaEndpointVideosEdits: + return xai.BuildVideosEditsURL(baseURL) + case GrokMediaEndpointVideosExtensions: + return xai.BuildVideosExtensionsURL(baseURL) case GrokMediaEndpointVideoStatus: return xai.BuildVideoURL(baseURL, requestID) default: @@ -302,7 +308,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( if err != nil { return nil, err } - targetURL, err := endpoint.upstreamURL(account.GetGrokBaseURL(), requestID) + targetURL, err := endpoint.upstreamURL(account.GetGrokMediaBaseURL(), requestID) if err != nil { return nil, err } @@ -531,7 +537,7 @@ func grokMediaUsageFromResponse(endpoint GrokMediaEndpoint, requestInfo GrokMedi meta.ImageSize = requestInfo.SizeTier meta.ImageInputSize = requestInfo.Size meta.ImageOutputSizes = collectOpenAIResponseImageOutputSizesFromJSONBytes(responseBody) - case GrokMediaEndpointVideosGenerations: + case GrokMediaEndpointVideosGenerations, GrokMediaEndpointVideosEdits, GrokMediaEndpointVideosExtensions: meta.ResponseID = extractGrokMediaVideoRequestID(responseBody) meta.VideoCount = 1 meta.VideoResolution = requestInfo.Resolution diff --git a/backend/internal/service/grok_oauth_service.go b/backend/internal/service/grok_oauth_service.go index 0a610ef811..136e5e3b21 100644 --- a/backend/internal/service/grok_oauth_service.go +++ b/backend/internal/service/grok_oauth_service.go @@ -3,8 +3,6 @@ package service import ( "context" "crypto/subtle" - "encoding/base64" - "encoding/json" "net/http" "strings" "time" @@ -101,6 +99,8 @@ type GrokTokenInfo struct { ClientID string `json:"client_id,omitempty"` Scope string `json:"scope,omitempty"` Email string `json:"email,omitempty"` + Subject string `json:"sub,omitempty"` + TeamID string `json:"team_id,omitempty"` SubscriptionTier string `json:"subscription_tier,omitempty"` EntitlementStatus string `json:"entitlement_status,omitempty"` } @@ -175,6 +175,18 @@ func (s *GrokOAuthService) ValidateRefreshToken(ctx context.Context, refreshToke return s.RefreshToken(ctx, refreshToken, proxyURL, xai.EffectiveClientID()) } +func (s *GrokOAuthService) ConvertFromSSO(ctx context.Context, ssoToken string, proxyID *int64) (*GrokTokenInfo, error) { + proxyURL, err := s.proxyURL(ctx, proxyID) + if err != nil { + return nil, err + } + tokenResp, err := s.oauthClient.ConvertSSOToBuild(ctx, ssoToken, proxyURL) + if err != nil { + return nil, err + } + return s.tokenInfoFromResponse(tokenResp, xai.DefaultClientID, nil), nil +} + func (s *GrokOAuthService) RefreshAccountToken(ctx context.Context, account *Account) (*GrokTokenInfo, error) { if account == nil || account.Platform != PlatformGrok { return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_INVALID_ACCOUNT", "account is not a Grok account") @@ -229,6 +241,12 @@ func (s *GrokOAuthService) BuildAccountCredentials(tokenInfo *GrokTokenInfo) map if tokenInfo.Email != "" { creds["email"] = tokenInfo.Email } + if tokenInfo.Subject != "" { + creds["sub"] = tokenInfo.Subject + } + if tokenInfo.TeamID != "" { + creds["team_id"] = tokenInfo.TeamID + } if tokenInfo.SubscriptionTier != "" { creds["subscription_tier"] = tokenInfo.SubscriptionTier } @@ -265,12 +283,23 @@ func (s *GrokOAuthService) tokenInfoFromResponse(tokenResp *xai.TokenResponse, c if info.TokenType == "" { info.TokenType = "Bearer" } - if email := parseJWTEmailClaim(tokenResp.IDToken); email != "" { - info.Email = email - } - if info.Email == "" && existing != nil { - if email, _ := existing["email"].(string); email != "" { - info.Email = email + applyGrokTokenClaims(info, tokenResp.IDToken) + applyGrokTokenClaims(info, tokenResp.AccessToken) + if existing != nil { + if info.Email == "" { + if email, _ := existing["email"].(string); email != "" { + info.Email = email + } + } + if info.Subject == "" { + if subject, _ := existing["sub"].(string); subject != "" { + info.Subject = subject + } + } + if info.TeamID == "" { + if teamID, _ := existing["team_id"].(string); teamID != "" { + info.TeamID = teamID + } } } return info @@ -293,20 +322,21 @@ func (s *GrokOAuthService) proxyURL(ctx context.Context, proxyID *int64) (string return proxy.URL(), nil } -func parseJWTEmailClaim(token string) string { - parts := strings.Split(token, ".") - if len(parts) < 2 { - return "" +func applyGrokTokenClaims(info *GrokTokenInfo, token string) { + if info == nil || strings.TrimSpace(token) == "" { + return } - payload, err := base64.RawURLEncoding.DecodeString(parts[1]) - if err != nil { - return "" + claims := xai.DecodeJWTClaims(token) + if claims == nil { + return } - var claims struct { - Email string `json:"email"` + if info.Email == "" { + info.Email = xai.JWTClaimString(claims, "email") } - if err := json.Unmarshal(payload, &claims); err != nil { - return "" + if info.Subject == "" { + info.Subject = xai.JWTClaimString(claims, "sub") + } + if info.TeamID == "" { + info.TeamID = xai.JWTClaimString(claims, "team_id") } - return strings.TrimSpace(claims.Email) } diff --git a/backend/internal/service/grok_oauth_service_test.go b/backend/internal/service/grok_oauth_service_test.go index f042c96337..54baef03a2 100644 --- a/backend/internal/service/grok_oauth_service_test.go +++ b/backend/internal/service/grok_oauth_service_test.go @@ -4,6 +4,8 @@ package service import ( "context" + "encoding/base64" + "encoding/json" "testing" "time" @@ -13,6 +15,7 @@ import ( type grokOAuthClientStub struct { refreshResponse *xai.TokenResponse + ssoResponse *xai.TokenResponse exchangeCalls int } @@ -25,6 +28,10 @@ func (s *grokOAuthClientStub) RefreshToken(context.Context, string, string, stri return s.refreshResponse, nil } +func (s *grokOAuthClientStub) ConvertSSOToBuild(context.Context, string, string) (*xai.TokenResponse, error) { + return s.ssoResponse, nil +} + func TestGrokOAuthServiceRefreshTokenPreservesOriginalRefreshTokenWhenNotRotated(t *testing.T) { svc := NewGrokOAuthService(nil, &grokOAuthClientStub{ refreshResponse: &xai.TokenResponse{ @@ -79,3 +86,31 @@ func TestGrokOAuthServiceBuildAccountCredentialsDefaultsToSubscriptionProxy(t *t require.Equal(t, xai.DefaultCLIBaseURL, credentials["base_url"]) } + +func TestGrokOAuthServiceConvertFromSSOExtractsBuildClaims(t *testing.T) { + svc := NewGrokOAuthService(nil, &grokOAuthClientStub{ + ssoResponse: &xai.TokenResponse{ + AccessToken: makeGrokOAuthJWT(map[string]any{"sub": "user-sub", "team_id": "team-1"}), + RefreshToken: "refresh-token", + IDToken: makeGrokOAuthJWT(map[string]any{"email": "user@example.com"}), + ExpiresIn: 3600, + }, + }) + defer svc.Stop() + + info, err := svc.ConvertFromSSO(context.Background(), "sso-token", nil) + require.NoError(t, err) + require.Equal(t, "user@example.com", info.Email) + require.Equal(t, "user-sub", info.Subject) + require.Equal(t, "team-1", info.TeamID) + + credentials := svc.BuildAccountCredentials(info) + require.Equal(t, "user@example.com", credentials["email"]) + require.Equal(t, "user-sub", credentials["sub"]) + require.Equal(t, "team-1", credentials["team_id"]) +} + +func makeGrokOAuthJWT(claims map[string]any) string { + payload, _ := json.Marshal(claims) + return "header." + base64.RawURLEncoding.EncodeToString(payload) + ".signature" +} diff --git a/backend/internal/service/grok_quota_fetcher.go b/backend/internal/service/grok_quota_fetcher.go index 0939b78e20..f220fe33b9 100644 --- a/backend/internal/service/grok_quota_fetcher.go +++ b/backend/internal/service/grok_quota_fetcher.go @@ -3,6 +3,8 @@ package service import ( "encoding/json" "fmt" + "net/http" + "strings" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" @@ -24,54 +26,150 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo { } if account == nil { usage.ErrorCode = "quota_unknown" - usage.Error = "Grok quota is unknown until the first upstream response includes xAI rate-limit headers" + usage.Error = "Grok quota is unknown until billing is probed or an upstream response includes xAI rate-limit headers" return usage } + billing, _ := grokBillingSnapshotFromExtra(account.Extra) snapshot, err := grokQuotaSnapshotFromExtra(account.Extra) + if billing != nil { + usage.GrokBilling = billing + if billing.Plan != "" { + usage.SubscriptionTier = billing.Plan + usage.SubscriptionTierRaw = billing.Plan + } + if parsedAt, parseErr := time.Parse(time.RFC3339, billing.UpdatedAt); parseErr == nil { + usage.UpdatedAt = &parsedAt + } + if billing.FetchedAt != "" { + usage.GrokLastQuotaProbeAt = billing.FetchedAt + } + usage.GrokQuotaSnapshotState = "billing_observed" + usage.GrokLastStatusCode = billing.StatusCode + switch billing.StatusCode { + case 401: + usage.NeedsReauth = true + usage.ErrorCode = "unauthenticated" + case 403: + usage.IsForbidden = true + usage.ForbiddenType = "forbidden" + usage.ErrorCode = "forbidden" + case 429: + usage.ErrorCode = "rate_limited" + } + } + if err != nil || snapshot == nil { - usage.ErrorCode = "quota_unknown" - usage.Error = "Grok quota is unknown until the first upstream response includes xAI rate-limit headers" + applyGrokCredentialUsageFallback(usage, account) + if billing == nil { + usage.ErrorCode = "quota_unknown" + usage.Error = "Grok quota is unknown until billing is probed or an upstream response includes xAI rate-limit headers" + } return usage } - if parsedAt, err := time.Parse(time.RFC3339, snapshot.UpdatedAt); err == nil { - usage.UpdatedAt = &parsedAt + if parsedAt, parseErr := time.Parse(time.RFC3339, snapshot.UpdatedAt); parseErr == nil { + if billing == nil || usage.UpdatedAt == nil || parsedAt.After(*usage.UpdatedAt) { + usage.UpdatedAt = &parsedAt + } } usage.GrokRequestQuota = snapshot.Requests usage.GrokTokenQuota = snapshot.Tokens usage.GrokRetryAfterSeconds = snapshot.RetryAfterSeconds - usage.SubscriptionTier = snapshot.SubscriptionTier - usage.SubscriptionTierRaw = snapshot.SubscriptionTier - usage.GrokEntitlementStatus = snapshot.EntitlementStatus - usage.GrokLastQuotaProbeAt = snapshot.LastProbeAt + if usage.SubscriptionTier == "" { + usage.SubscriptionTier = snapshot.SubscriptionTier + usage.SubscriptionTierRaw = snapshot.SubscriptionTier + } + if usage.GrokEntitlementStatus == "" { + usage.GrokEntitlementStatus = snapshot.EntitlementStatus + } + if usage.GrokLastQuotaProbeAt == "" { + usage.GrokLastQuotaProbeAt = snapshot.LastProbeAt + } usage.GrokLastHeadersSeenAt = snapshot.LastHeadersSeenAt - usage.GrokLastStatusCode = snapshot.StatusCode + if snapshot.StatusCode >= http.StatusBadRequest || usage.GrokLastStatusCode == 0 { + usage.GrokLastStatusCode = snapshot.StatusCode + } if snapshot.HasObservedHeaders() { - usage.GrokQuotaSnapshotState = "observed" - } else { + if usage.GrokQuotaSnapshotState == "" { + usage.GrokQuotaSnapshotState = "observed" + } + } else if billing == nil { usage.GrokQuotaSnapshotState = "no_headers" usage.ErrorCode = "quota_unknown" usage.Error = "No xAI quota headers observed on the latest Grok probe" } - switch snapshot.StatusCode { - case 401: - usage.NeedsReauth = true - usage.ErrorCode = "unauthenticated" - case 403: - usage.IsForbidden = true - usage.ForbiddenType = "forbidden" - usage.ErrorCode = "forbidden" - if usage.GrokEntitlementStatus == "" { - usage.GrokEntitlementStatus = "forbidden" + if usage.ErrorCode == "" { + switch snapshot.StatusCode { + case 401: + usage.NeedsReauth = true + usage.ErrorCode = "unauthenticated" + case 403: + usage.IsForbidden = true + usage.ForbiddenType = "forbidden" + usage.ErrorCode = "forbidden" + if usage.GrokEntitlementStatus == "" { + usage.GrokEntitlementStatus = "forbidden" + } + case 429: + usage.ErrorCode = "rate_limited" } - case 429: - usage.ErrorCode = "rate_limited" } + applyGrokCredentialUsageFallback(usage, account) return usage } +func applyGrokCredentialUsageFallback(usage *UsageInfo, account *Account) { + if usage == nil || account == nil { + return + } + if usage.SubscriptionTier == "" { + tier := strings.TrimSpace(account.GetCredential("subscription_tier")) + usage.SubscriptionTier = tier + usage.SubscriptionTierRaw = tier + } + if usage.GrokEntitlementStatus == "" { + usage.GrokEntitlementStatus = strings.TrimSpace(account.GetCredential("entitlement_status")) + } +} + +func grokBillingSnapshotFromExtra(extra map[string]any) (*xai.BillingSummary, error) { + if extra == nil { + return nil, nil + } + raw, ok := extra[grokBillingExtraKey] + if !ok || raw == nil { + return nil, nil + } + switch snapshot := raw.(type) { + case *xai.BillingSummary: + return snapshot, nil + case xai.BillingSummary: + return &snapshot, nil + case map[string]any: + data, err := json.Marshal(snapshot) + if err != nil { + return nil, err + } + var out xai.BillingSummary + if err := json.Unmarshal(data, &out); err != nil { + return nil, err + } + return &out, nil + default: + data, err := json.Marshal(raw) + if err != nil { + return nil, fmt.Errorf("marshal grok billing snapshot: %w", err) + } + var out xai.BillingSummary + if err := json.Unmarshal(data, &out); err != nil { + return nil, err + } + return &out, nil + } +} + func grokQuotaSnapshotFromExtra(extra map[string]any) (*xai.QuotaSnapshot, error) { if extra == nil { return nil, nil diff --git a/backend/internal/service/grok_quota_fetcher_test.go b/backend/internal/service/grok_quota_fetcher_test.go index d2d9c14993..1de9b51c9e 100644 --- a/backend/internal/service/grok_quota_fetcher_test.go +++ b/backend/internal/service/grok_quota_fetcher_test.go @@ -20,7 +20,34 @@ func TestGrokQuotaFetcherBuildUsageInfoUnknownUntilFirstSnapshot(t *testing.T) { usage := NewGrokQuotaFetcher().BuildUsageInfo(&Account{Platform: PlatformGrok, Type: AccountTypeOAuth}) require.Equal(t, "passive", usage.Source) require.Equal(t, "quota_unknown", usage.ErrorCode) - require.Contains(t, usage.Error, "unknown until the first upstream response") + require.Contains(t, usage.Error, "unknown until billing is probed") +} + +func TestGrokQuotaFetcherUsesCredentialTierWhenBillingHasNoPlan(t *testing.T) { + t.Parallel() + + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "subscription_tier": " FREE ", + "entitlement_status": " active ", + }, + Extra: map[string]any{ + grokBillingExtraKey: &xai.BillingSummary{ + PeriodType: "weekly", + StatusCode: http.StatusOK, + UpdatedAt: "2030-01-01T00:00:00Z", + }, + }, + } + + usage := NewGrokQuotaFetcher().BuildUsageInfo(account) + + require.NotNil(t, usage.GrokBilling) + require.Equal(t, "FREE", usage.SubscriptionTier) + require.Equal(t, "FREE", usage.SubscriptionTierRaw) + require.Equal(t, "active", usage.GrokEntitlementStatus) } func TestGrokQuotaFetcherBuildUsageInfoFromSnapshot(t *testing.T) { @@ -68,6 +95,32 @@ func TestGrokQuotaFetcherBuildUsageInfoFromSnapshot(t *testing.T) { require.True(t, usage.UpdatedAt.Equal(time.Date(2030, 1, 1, 0, 0, 0, 0, time.UTC))) } +func TestGrokQuotaFetcherSnapshotErrorOverridesSuccessfulBillingStatus(t *testing.T) { + t.Parallel() + + updatedAt := "2030-01-01T00:00:00Z" + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Extra: map[string]any{ + grokBillingExtraKey: &xai.BillingSummary{ + PeriodType: "weekly", + StatusCode: http.StatusOK, + UpdatedAt: updatedAt, + }, + grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + UpdatedAt: updatedAt, + }, + }, + } + + usage := NewGrokQuotaFetcher().BuildUsageInfo(account) + + require.Equal(t, "rate_limited", usage.ErrorCode) + require.Equal(t, http.StatusTooManyRequests, usage.GrokLastStatusCode) +} + func TestGrokQuotaFetcherBuildUsageInfoFromNoHeadersProbe(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 17e91dac6f..2219a75c15 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -7,27 +7,37 @@ import ( "io" "log/slog" "net/http" + "strconv" "strings" + "sync" "time" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "golang.org/x/sync/singleflight" ) const ( grokQuotaUpstreamTimeout = 20 * time.Second grokQuotaProbeInput = "." - grokQuotaDefaultModel = "grok-4.3" + grokQuotaDefaultModel = grokDefaultResponsesModel + grokBillingExtraKey = "grok_billing_snapshot" ) type GrokQuotaProbeResult struct { - Source string `json:"source"` - Model string `json:"model"` - Snapshot *xai.QuotaSnapshot `json:"snapshot,omitempty"` - StatusCode int `json:"status_code,omitempty"` - HeadersObserved bool `json:"headers_observed"` - ResetSupported bool `json:"reset_supported"` - FetchedAt int64 `json:"fetched_at"` + Source string `json:"source"` + Model string `json:"model,omitempty"` + Billing *xai.BillingSummary `json:"billing,omitempty"` + Snapshot *xai.QuotaSnapshot `json:"snapshot,omitempty"` + LocalUsage24h *WindowStats `json:"local_usage_24h,omitempty"` + LocalUsage7d *WindowStats `json:"local_usage_7d,omitempty"` + LocalUsageMonthly *WindowStats `json:"local_usage_monthly,omitempty"` + StatusCode int `json:"status_code,omitempty"` + HeadersObserved bool `json:"headers_observed"` + ResetSupported bool `json:"reset_supported"` + FetchedAt int64 `json:"fetched_at"` + Persisted bool `json:"persisted"` + ProbeError string `json:"probe_error,omitempty"` } type GrokQuotaResetResult struct { @@ -41,6 +51,8 @@ type GrokQuotaService struct { proxyRepo ProxyRepository tokenProvider *GrokTokenProvider httpUpstream HTTPUpstream + usageLogRepo UsageLogRepository + probeFlight singleflight.Group } func NewGrokQuotaService( @@ -48,16 +60,71 @@ func NewGrokQuotaService( proxyRepo ProxyRepository, tokenProvider *GrokTokenProvider, httpUpstream HTTPUpstream, + usageLogRepos ...UsageLogRepository, ) *GrokQuotaService { + var usageLogRepo UsageLogRepository + if len(usageLogRepos) > 0 { + usageLogRepo = usageLogRepos[0] + } return &GrokQuotaService{ accountRepo: accountRepo, proxyRepo: proxyRepo, tokenProvider: tokenProvider, httpUpstream: httpUpstream, + usageLogRepo: usageLogRepo, } } +// QueryQuota combines xAI billing data with an active quota-header probe for +// Free accounts, whose billing response does not include usage_percent. +func (s *GrokQuotaService) QueryQuota(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { + billingResult, billingErr := s.ProbeBilling(ctx, accountID) + if billingErr == nil && billingResult != nil && grokBillingHasAuthoritativeQuota(billingResult.Billing) { + return billingResult, nil + } + + probeResult, probeErr := s.ProbeUsage(ctx, accountID) + if probeErr != nil { + if billingResult != nil && billingResult.Billing != nil { + billingResult.ProbeError = probeErr.Error() + return billingResult, nil + } + return nil, probeErr + } + if probeResult == nil { + if billingErr != nil { + return nil, billingErr + } + return nil, infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_PROBE_EMPTY", "Grok quota probe returned no result") + } + if billingResult != nil { + probeResult.Source = "hybrid_probe" + probeResult.Billing = billingResult.Billing + probeResult.LocalUsage24h = billingResult.LocalUsage24h + probeResult.LocalUsage7d = billingResult.LocalUsage7d + probeResult.LocalUsageMonthly = billingResult.LocalUsageMonthly + probeResult.Persisted = probeResult.Persisted || billingResult.Persisted + } + return probeResult, nil +} + +func grokBillingHasAuthoritativeQuota(billing *xai.BillingSummary) bool { + if billing == nil { + return false + } + return billing.UsagePercent != nil || + billing.UsedPercent != nil || + (billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0) || + strings.TrimSpace(billing.Plan) != "" +} + func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { + return s.runProbeFlight(ctx, "active:"+strconv.FormatInt(accountID, 10), func(sharedCtx context.Context) (*GrokQuotaProbeResult, error) { + return s.probeUsage(sharedCtx, accountID) + }) +} + +func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { account, token, proxyURL, err := s.prepareProbe(ctx, accountID) if err != nil { return nil, err @@ -95,7 +162,7 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr if limited { normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now()) } - _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ + persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ grokQuotaSnapshotExtraKey: snapshot, }) if limited { @@ -110,19 +177,201 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr HeadersObserved: snapshot.HeadersObserved, ResetSupported: false, FetchedAt: time.Now().Unix(), + Persisted: persistErr == nil, } if resp.StatusCode == http.StatusTooManyRequests { return result, nil } if resp.StatusCode >= 400 { - bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, 240)) - bodyText := truncate(strings.TrimSpace(string(bodyBytes)), 240) - slog.Warn("grok_quota_probe_failed", "account_id", account.ID, "model", probeModel, "status", resp.StatusCode, "body", bodyText) - return nil, infraerrors.Newf(mapUpstreamStatus(resp.StatusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "upstream returned %d for probe model %q: %s", resp.StatusCode, probeModel, bodyText) + const reason = "GROK_QUOTA_PROBE_UPSTREAM_ERROR" + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10)) + slog.Warn( + "grok_quota_probe_failed", + "account_id", account.ID, + "model", probeModel, + "status", resp.StatusCode, + "reason", reason, + ) + return nil, infraerrors.Newf( + mapUpstreamStatus(resp.StatusCode), + reason, + "upstream returned %d for probe model %q", + resp.StatusCode, + probeModel, + ) } return result, nil } +// ProbeBilling only calls the xAI billing endpoints. Account usage refreshes +// use this method so opening the account list never consumes model quota. +func (s *GrokQuotaService) ProbeBilling(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { + return s.runProbeFlight(ctx, "billing:"+strconv.FormatInt(accountID, 10), func(sharedCtx context.Context) (*GrokQuotaProbeResult, error) { + return s.probeBilling(sharedCtx, accountID) + }) +} + +func (s *GrokQuotaService) probeBilling(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { + account, token, proxyURL, err := s.prepareProbe(ctx, accountID) + if err != nil { + return nil, err + } + + probeCtx, cancel := context.WithTimeout(ctx, grokQuotaUpstreamTimeout) + defer cancel() + type billingResult struct { + summary *xai.BillingSummary + status int + err error + } + var weekly, monthly billingResult + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + weekly.summary, weekly.status, weekly.err = s.fetchBilling(probeCtx, account, token, proxyURL, true) + }() + go func() { + defer wg.Done() + monthly.summary, monthly.status, monthly.err = s.fetchBilling(probeCtx, account, token, proxyURL, false) + }() + wg.Wait() + + weeklyOK := weekly.summary != nil + monthlyOK := monthly.summary != nil + if !weeklyOK && !monthlyOK { + return nil, mergeGrokBillingProbeErrors(weekly.status, monthly.status, weekly.err, monthly.err) + } + statusCode := preferSuccessfulBillingStatus(weekly.status, monthly.status, weeklyOK, monthlyOK) + previous, _ := grokBillingSnapshotFromExtra(account.Extra) + billing := xai.MergeBillingProbeResult(previous, weekly.summary, monthly.summary, weeklyOK, monthlyOK) + billing = xai.StampBillingSummary(billing, statusCode, "billing_probe") + persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ + grokBillingExtraKey: billing, + }) + if persistErr != nil { + slog.Warn("grok_billing_persist_failed", "account_id", account.ID, "error", persistErr) + } + now := time.Now().UTC() + localUsage24h, localUsage7d, localUsageMonthly := grokLocalUsageForQuota(ctx, s.usageLogRepo, account.ID, billing, now) + return &GrokQuotaProbeResult{ + Source: "billing_probe", + Billing: billing, + LocalUsage24h: localUsage24h, + LocalUsage7d: localUsage7d, + LocalUsageMonthly: localUsageMonthly, + StatusCode: statusCode, + FetchedAt: now.Unix(), + Persisted: persistErr == nil, + }, nil +} + +func (s *GrokQuotaService) runProbeFlight( + ctx context.Context, + key string, + probe func(context.Context) (*GrokQuotaProbeResult, error), +) (*GrokQuotaProbeResult, error) { + if s == nil { + return nil, infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_NOT_CONFIGURED", "grok quota service is not configured") + } + resultCh := s.probeFlight.DoChan(key, func() (any, error) { + sharedCtx, cancel := context.WithTimeout(context.Background(), grokQuotaUpstreamTimeout+5*time.Second) + defer cancel() + return probe(sharedCtx) + }) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case flightResult := <-resultCh: + if flightResult.Err != nil { + return nil, flightResult.Err + } + result, ok := flightResult.Val.(*GrokQuotaProbeResult) + if !ok || result == nil { + return nil, infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_PROBE_RESULT_INVALID", "invalid Grok quota probe result") + } + cloned := *result + return &cloned, nil + } +} + +func (s *GrokQuotaService) fetchBilling( + ctx context.Context, + account *Account, + token string, + proxyURL string, + weekly bool, +) (*xai.BillingSummary, int, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, xai.BuildBillingURL(weekly), nil) + if err != nil { + return nil, 0, infraerrors.Newf(http.StatusInternalServerError, "GROK_QUOTA_PROBE_REQUEST_BUILD_FAILED", "failed to build billing request: %v", err) + } + xai.ApplyCLIBillingHeaders(req, token) + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 2)) + if err != nil { + return nil, 0, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", err) + } + defer func() { _ = resp.Body.Close() }() + + bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if resp.StatusCode == http.StatusTooManyRequests { + return nil, resp.StatusCode, nil + } + if resp.StatusCode >= 400 { + bodyText := truncate(strings.TrimSpace(string(bodyBytes)), 240) + slog.Warn("grok_quota_billing_failed", "account_id", account.ID, "weekly", weekly, "status", resp.StatusCode, "body", bodyText) + return nil, resp.StatusCode, infraerrors.Newf(mapUpstreamStatus(resp.StatusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "billing returned %d: %s", resp.StatusCode, bodyText) + } + payload, err := xai.ParseBillingPayload(bodyBytes) + if err != nil { + return nil, resp.StatusCode, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_BILLING_PARSE_ERROR", "failed to parse billing body: %v", err) + } + return xai.BuildBillingSummary(payload.Config), resp.StatusCode, nil +} + +func mergeGrokBillingProbeErrors(weeklyStatus, monthlyStatus int, weeklyErr, monthlyErr error) error { + weeklyKey := grokBillingProbeErrorKey(weeklyStatus, weeklyErr) + monthlyKey := grokBillingProbeErrorKey(monthlyStatus, monthlyErr) + if weeklyKey == monthlyKey { + switch { + case weeklyErr != nil: + return weeklyErr + case monthlyErr != nil: + return monthlyErr + case weeklyStatus == http.StatusTooManyRequests: + return infraerrors.New(http.StatusTooManyRequests, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "billing rate limited") + case weeklyStatus != 0 && weeklyStatus != http.StatusOK: + return infraerrors.New(mapUpstreamStatus(weeklyStatus), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "xAI billing endpoints returned the same upstream error") + default: + return infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_BILLING_EMPTY", "xAI billing endpoints returned no quota data") + } + } + slog.Warn("grok_quota_probe_parts_failed", "weekly_status", weeklyStatus, "weekly_error", weeklyErr, "monthly_status", monthlyStatus, "monthly_error", monthlyErr) + return infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_PROBE_PARTS_FAILED", "weekly and monthly billing probes failed differently").WithMetadata(map[string]string{ + "weekly_status": strconv.Itoa(weeklyStatus), "monthly_status": strconv.Itoa(monthlyStatus), + }) +} + +func grokBillingProbeErrorKey(status int, err error) string { + if err != nil { + return strconv.Itoa(status) + ":" + strconv.Itoa(infraerrors.Code(err)) + ":" + infraerrors.Reason(err) + } + return strconv.Itoa(status) + ":empty" +} + +func preferSuccessfulBillingStatus(weeklyStatus, monthlyStatus int, weeklyOK, monthlyOK bool) int { + if weeklyOK && weeklyStatus >= 200 && weeklyStatus < 300 { + return weeklyStatus + } + if monthlyOK && monthlyStatus >= 200 && monthlyStatus < 300 { + return monthlyStatus + } + if weeklyStatus != 0 { + return weeklyStatus + } + return monthlyStatus +} + func (s *GrokQuotaService) ResetQuota(ctx context.Context, accountID int64) (*GrokQuotaResetResult, error) { if _, err := s.loadGrokOAuthAccount(ctx, accountID); err != nil { return nil, err diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 2248674899..ba9b6cbcb8 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -3,14 +3,19 @@ package service import ( + "bytes" "context" "io" + "log/slog" "net/http" + "strconv" "strings" + "sync" "testing" "time" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" @@ -63,6 +68,113 @@ type grokQuotaProxyRepo struct { calls int } +type grokQuotaUsageLogRepo struct { + UsageLogRepository + stats *usagestats.AccountStats + err error + calls int + startTimes []time.Time +} + +func (r *grokQuotaUsageLogRepo) GetAccountWindowStats(_ context.Context, _ int64, start time.Time) (*usagestats.AccountStats, error) { + r.calls++ + r.startTimes = append(r.startTimes, start) + return r.stats, r.err +} + +func (r *grokQuotaUsageLogRepo) GetAccountTodayStats(context.Context, int64) (*usagestats.AccountStats, error) { + return nil, nil +} + +type grokHybridUpstream struct { + httpUpstreamRecorder + mu sync.Mutex + requests []*http.Request + bodies [][]byte + weeklyUsagePercent *float64 + monthlyLimitCents *float64 + activeStatus int + activeHeaders http.Header + billingStarted chan struct{} + billingRelease <-chan struct{} + billingStartOnce sync.Once + billingStatus int + billingHeaders http.Header +} + +func (u *grokHybridUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + var body []byte + if req != nil && req.Body != nil { + body, _ = io.ReadAll(req.Body) + } + u.mu.Lock() + u.requests = append(u.requests, req) + u.bodies = append(u.bodies, body) + u.mu.Unlock() + + if req.URL.Path == "/v1/responses" { + status := u.activeStatus + if status == 0 { + status = http.StatusOK + } + headers := u.activeHeaders + if headers == nil { + headers = http.Header{ + "X-Ratelimit-Limit-Tokens": []string{"2000000"}, + "X-Ratelimit-Remaining-Tokens": []string{"1500000"}, + } + } + return &http.Response{StatusCode: status, Header: headers, Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`))}, nil + } + if u.billingStarted != nil { + u.billingStartOnce.Do(func() { close(u.billingStarted) }) + } + if u.billingRelease != nil { + select { + case <-u.billingRelease: + case <-req.Context().Done(): + return nil, req.Context().Err() + } + } + if u.billingStatus != 0 && u.billingStatus != http.StatusOK { + return &http.Response{ + StatusCode: u.billingStatus, + Header: u.billingHeaders, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"billing limited"}}`)), + }, nil + } + + if req.URL.RawQuery == "format=credits" { + usage := "" + if u.weeklyUsagePercent != nil { + usage = `,"creditUsagePercent":` + strconv.FormatFloat(*u.weeklyUsagePercent, 'f', -1, 64) + } + payload := `{"config":{"currentPeriod":{"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"}` + usage + `}}` + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(payload))}, nil + } + monthlyLimit := "" + if u.monthlyLimitCents != nil { + monthlyLimit = `,"monthlyLimit":{"val":` + strconv.FormatFloat(*u.monthlyLimitCents, 'f', -1, 64) + `}` + } + monthlyPayload := `{"config":{"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"` + monthlyLimit + `}}` + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(monthlyPayload)), + }, nil +} + +func (u *grokHybridUpstream) snapshot() ([]*http.Request, [][]byte) { + u.mu.Lock() + defer u.mu.Unlock() + requests := append([]*http.Request(nil), u.requests...) + bodies := make([][]byte, len(u.bodies)) + for i := range u.bodies { + bodies[i] = append([]byte(nil), u.bodies[i]...) + } + return requests, bodies +} + func (r *grokQuotaProxyRepo) GetByID(_ context.Context, id int64) (*Proxy, error) { r.calls++ return r.proxies[id], nil @@ -102,7 +214,7 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) { result, err := svc.ProbeUsage(context.Background(), 42) require.NoError(t, err) require.Equal(t, http.StatusOK, result.StatusCode) - require.Equal(t, "grok-4.3", result.Model) + require.Equal(t, "grok-4.5", result.Model) require.True(t, result.HeadersObserved) require.NotNil(t, result.Snapshot) require.True(t, result.Snapshot.HeadersObserved) @@ -115,7 +227,7 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) { require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version")) - require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) require.Contains(t, string(upstream.lastBody), `"max_output_tokens":1`) require.Contains(t, string(upstream.lastBody), `"store":false`) require.NotNil(t, repo.updates[42][grokQuotaSnapshotExtraKey]) @@ -152,8 +264,8 @@ func TestGrokQuotaServiceProbeUsageIgnoresAccountGrokMapping(t *testing.T) { result, err := svc.ProbeUsage(context.Background(), 47) require.NoError(t, err) - require.Equal(t, "grok-4.3", result.Model) - require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "grok-4.5", result.Model) + require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) require.NotContains(t, string(upstream.lastBody), "grok-composer") } @@ -185,7 +297,55 @@ func TestGrokQuotaServiceProbeUsageReportsProbeModelOnUpstreamError(t *testing.T _, err := svc.ProbeUsage(context.Background(), 48) require.Error(t, err) require.Equal(t, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", infraerrors.Reason(err)) - require.Contains(t, infraerrors.Message(err), `probe model "grok-4.3"`) + require.Contains(t, infraerrors.Message(err), `probe model "grok-4.5"`) +} + +func TestGrokQuotaServiceProbeUsageRedactsUpstreamErrorBodyFromErrorAndLogs(t *testing.T) { + const upstreamSecret = "upstream-secret-refresh-token" + account := &Account{ + ID: 49, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{ + mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{49: account}, + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader( + `{"error":"` + upstreamSecret + `","detail":"credential rejected"}`, + )), + }} + svc := NewGrokQuotaService( + repo, + nil, + NewGrokTokenProvider(repo, nil), + upstream, + ) + + var logs bytes.Buffer + previousLogger := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&logs, nil))) + defer slog.SetDefault(previousLogger) + + _, err := svc.ProbeUsage(context.Background(), account.ID) + require.Error(t, err) + require.Equal(t, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", infraerrors.Reason(err)) + require.Contains(t, infraerrors.Message(err), `probe model "grok-4.5"`) + require.NotContains(t, err.Error(), upstreamSecret) + require.NotContains(t, infraerrors.Message(err), upstreamSecret) + require.Contains(t, logs.String(), "GROK_QUOTA_PROBE_UPSTREAM_ERROR") + require.NotContains(t, logs.String(), upstreamSecret) + require.NotContains(t, logs.String(), "credential rejected") + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) } func TestGrokQuotaServiceProbeUsageLoadsProxyWhenAccountEdgeMissing(t *testing.T) { @@ -308,6 +468,385 @@ func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) { require.Zero(t, repo.tempUnschedCalls) } +func TestGrokQuotaServiceQueryQuotaFreeFallsBackToGrok45(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 51, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{} + usageRepo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 1_000_000}} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, usageRepo) + + result, err := svc.QueryQuota(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, "hybrid_probe", result.Source) + require.Equal(t, "grok-4.5", result.Model) + require.NotNil(t, result.Billing) + require.Nil(t, result.Billing.UsagePercent) + require.NotNil(t, result.LocalUsage24h) + require.EqualValues(t, 1_000_000, result.LocalUsage24h.Tokens) + require.Equal(t, 1, usageRepo.calls) + require.WithinDuration(t, time.Now().UTC().Add(-24*time.Hour), usageRepo.startTimes[0], time.Second) + require.NotNil(t, result.Snapshot) + require.NotNil(t, result.Snapshot.Tokens) + require.EqualValues(t, 2_000_000, *result.Snapshot.Tokens.Limit) + require.True(t, result.HeadersObserved) + + requests, bodies := upstream.snapshot() + require.Len(t, requests, 3) + responseCalls := 0 + for i, req := range requests { + if req.URL.Path != "/v1/responses" { + continue + } + responseCalls++ + require.Equal(t, http.MethodPost, req.Method) + require.Equal(t, "grok-4.5", gjson.GetBytes(bodies[i], "model").String()) + require.EqualValues(t, 1, gjson.GetBytes(bodies[i], "max_output_tokens").Int()) + } + require.Equal(t, 1, responseCalls) +} + +func TestGrokQuotaServiceQueryQuotaPaidBillingSkipsActiveProbe(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 52, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + usagePercent := 25.0 + upstream := &grokHybridUpstream{weeklyUsagePercent: &usagePercent} + usageRepo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 1_000_000}} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, usageRepo) + + result, err := svc.QueryQuota(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, "billing_probe", result.Source) + require.NotNil(t, result.Billing) + require.InDelta(t, usagePercent, *result.Billing.UsagePercent, 1e-9) + require.Nil(t, result.Snapshot) + require.Empty(t, result.Model) + require.Nil(t, result.LocalUsage24h) + + requests, _ := upstream.snapshot() + require.Len(t, requests, 2) + for _, req := range requests { + require.Equal(t, "/v1/billing", req.URL.Path) + } +} + +func TestGrokQuotaServiceQueryQuotaCustomPaidMonthlyLimitSkipsActiveProbe(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 57, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + monthlyLimit := 25_000.0 + upstream := &grokHybridUpstream{monthlyLimitCents: &monthlyLimit} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + result, err := svc.QueryQuota(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, "billing_probe", result.Source) + require.NotNil(t, result.Billing) + require.InDelta(t, monthlyLimit, *result.Billing.MonthlyLimitCents, 1e-9) + require.Nil(t, result.Snapshot) + + requests, _ := upstream.snapshot() + require.Len(t, requests, 2) + for _, req := range requests { + require.Equal(t, "/v1/billing", req.URL.Path) + } +} + +func TestGrokLocalUsage24hUsesRollingUTCWindow(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 7, 14, 20, 30, 0, 0, time.FixedZone("UTC+8", 8*60*60)) + + t.Run("returns usage from exact rolling window", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 1_250_000}} + stats := grokLocalUsage24h(context.Background(), repo, 57, now) + + require.NotNil(t, stats) + require.EqualValues(t, 1_250_000, stats.Tokens) + require.Equal(t, []time.Time{now.UTC().Add(-24 * time.Hour)}, repo.startTimes) + }) + + t.Run("query failure returns no stats", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{err: context.DeadlineExceeded} + stats := grokLocalUsage24h(context.Background(), repo, 57, now) + + require.Nil(t, stats) + require.Equal(t, []time.Time{now.UTC().Add(-24 * time.Hour)}, repo.startTimes) + }) + + t.Run("missing repository returns no stats", func(t *testing.T) { + require.Nil(t, grokLocalUsage24h(context.Background(), nil, 57, now)) + }) + + t.Run("invalid account returns no stats without query", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{} + require.Nil(t, grokLocalUsage24h(context.Background(), repo, 0, now)) + require.Zero(t, repo.calls) + }) +} + +func TestGrokLocalUsageForQuotaSelectsFreeOrPaidWindows(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 7, 14, 12, 0, 0, 0, time.UTC) + billing := &xai.BillingSummary{ + PeriodType: "weekly", + PeriodStart: now.Add(-4 * 24 * time.Hour).Format(time.RFC3339), + PeriodEnd: now.Add(3 * 24 * time.Hour).Format(time.RFC3339), + BillingPeriodStart: now.Add(-13 * 24 * time.Hour).Format(time.RFC3339), + BillingPeriodEnd: now.Add(17 * 24 * time.Hour).Format(time.RFC3339), + } + + t.Run("free queries only rolling 24h", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 500_000}} + rolling, weekly, monthly := grokLocalUsageForQuota(context.Background(), repo, 57, billing, now) + + require.NotNil(t, rolling) + require.Nil(t, weekly) + require.Nil(t, monthly) + require.Equal(t, []time.Time{now.Add(-24 * time.Hour)}, repo.startTimes) + }) + + t.Run("paid queries only billing windows", func(t *testing.T) { + usagePercent := 25.0 + paidBilling := *billing + paidBilling.UsagePercent = &usagePercent + repo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 500_000}} + rolling, weekly, monthly := grokLocalUsageForQuota(context.Background(), repo, 57, &paidBilling, now) + + require.Nil(t, rolling) + require.NotNil(t, weekly) + require.NotNil(t, monthly) + require.Equal(t, []time.Time{ + now.Add(-4 * 24 * time.Hour), + now.Add(-13 * 24 * time.Hour), + }, repo.startTimes) + }) +} + +func TestGrokLocalUsageForBillingOnlyReturnsAvailableWindows(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + billing := &xai.BillingSummary{ + PeriodType: "weekly", + PeriodStart: now.Add(-4 * 24 * time.Hour).Format(time.RFC3339), + PeriodEnd: now.Add(3 * 24 * time.Hour).Format(time.RFC3339), + } + + t.Run("valid weekly window", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 1_500_000}} + weekly, monthly := grokLocalUsageForBilling(context.Background(), repo, 57, billing, now) + require.NotNil(t, weekly) + require.EqualValues(t, 1_500_000, weekly.Tokens) + require.Nil(t, monthly) + require.Equal(t, 1, repo.calls) + }) + + t.Run("query failure", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{err: context.DeadlineExceeded} + weekly, monthly := grokLocalUsageForBilling(context.Background(), repo, 57, billing, now) + require.Nil(t, weekly) + require.Nil(t, monthly) + require.Equal(t, 1, repo.calls) + }) + + t.Run("missing billing window", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{} + weekly, monthly := grokLocalUsageForBilling(context.Background(), repo, 57, nil, now) + require.Nil(t, weekly) + require.Nil(t, monthly) + require.Zero(t, repo.calls) + }) +} + +func TestAccountUsageServiceGrokRefreshUsesBillingOnly(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 54, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{} + usageRepo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 750_000}} + quotaService := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, usageRepo) + usageService := &AccountUsageService{ + grokQuotaFetcher: NewGrokQuotaFetcher(), + grokQuotaService: quotaService, + usageLogRepo: usageRepo, + cache: NewUsageCache(), + } + + usage, err := usageService.getGrokUsage(context.Background(), account, false) + require.NoError(t, err) + require.NotNil(t, usage.GrokBilling) + require.Nil(t, usage.GrokBilling.UsagePercent) + require.NotNil(t, usage.GrokLocalUsage24h) + require.EqualValues(t, 750_000, usage.GrokLocalUsage24h.Tokens) + require.Equal(t, 1, usageRepo.calls) + require.Len(t, usageRepo.startTimes, 1) + require.WithinDuration(t, time.Now().UTC().Add(-24*time.Hour), usageRepo.startTimes[0], time.Second) + + requests, _ := upstream.snapshot() + require.Len(t, requests, 2) + for _, req := range requests { + require.Equal(t, http.MethodGet, req.Method) + require.Equal(t, "/v1/billing", req.URL.Path) + } +} + +func TestGrokQuotaServiceProbeFlightsDeduplicateBillingAndSeparateActive(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 55, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + billingStarted := make(chan struct{}) + billingRelease := make(chan struct{}) + upstream := &grokHybridUpstream{billingStarted: billingStarted, billingRelease: billingRelease} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + type probeOutcome struct { + result *GrokQuotaProbeResult + err error + } + billingOutcomes := make(chan probeOutcome, 2) + go func() { + result, err := svc.ProbeBilling(context.Background(), account.ID) + billingOutcomes <- probeOutcome{result: result, err: err} + }() + <-billingStarted + secondStarted := make(chan struct{}) + go func() { + close(secondStarted) + result, err := svc.ProbeBilling(context.Background(), account.ID) + billingOutcomes <- probeOutcome{result: result, err: err} + }() + <-secondStarted + time.Sleep(25 * time.Millisecond) + + activeResult, err := svc.ProbeUsage(context.Background(), account.ID) + require.NoError(t, err) + require.NotNil(t, activeResult.Snapshot) + close(billingRelease) + for range 2 { + outcome := <-billingOutcomes + require.NoError(t, outcome.err) + require.NotNil(t, outcome.result.Billing) + } + + requests, _ := upstream.snapshot() + billingCalls := 0 + activeCalls := 0 + for _, req := range requests { + switch req.URL.Path { + case "/v1/billing": + billingCalls++ + case "/v1/responses": + activeCalls++ + } + } + require.Equal(t, 2, billingCalls) + require.Equal(t, 1, activeCalls) +} + +func TestGrokQuotaServiceBilling429DoesNotPauseModelScheduling(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 56, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{ + billingStatus: http.StatusTooManyRequests, + billingHeaders: http.Header{"Retry-After": []string{"45"}}, + } + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + result, err := svc.ProbeBilling(context.Background(), account.ID) + + require.Error(t, err) + require.Nil(t, result) + require.Zero(t, repo.rateLimitedCalls) +} + +func TestGrokQuotaServiceQueryQuotaFree429PersistsLimitAndKeepsBilling(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 53, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{ + activeStatus: http.StatusTooManyRequests, + activeHeaders: http.Header{"Retry-After": []string{"45"}}, + } + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + result, err := svc.QueryQuota(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, http.StatusTooManyRequests, result.StatusCode) + require.NotNil(t, result.Billing) + require.NotNil(t, result.Snapshot) + require.Equal(t, 45, *result.Snapshot.RetryAfterSeconds) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Equal(t, account.ID, repo.lastRateLimitedID) + require.WithinDuration(t, time.Now().Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) +} + func TestGrokQuotaServiceResetQuotaUnsupported(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/image_generation_intent.go b/backend/internal/service/image_generation_intent.go index 5a063a7d71..613f11d048 100644 --- a/backend/internal/service/image_generation_intent.go +++ b/backend/internal/service/image_generation_intent.go @@ -9,9 +9,23 @@ import ( const ( openAIResponsesEndpoint = "/v1/responses" openAIResponsesCompactEndpoint = "/v1/responses/compact" + responsesLiteHeader = "X-OpenAI-Internal-Codex-Responses-Lite" + responsesLiteHeaderKey = "x-openai-internal-codex-responses-lite" + responsesLiteWSMetadataKey = "ws_request_header_x_openai_internal_codex_responses_lite" imageGenerationPermissionMessage = "Image generation is not enabled for this group" ) +func isOpenAIResponsesLiteHeader(value string) bool { + return strings.EqualFold(strings.TrimSpace(value), "true") +} + +func isOpenAIResponsesLiteWebSocketPayload(body []byte) bool { + if len(body) == 0 || !gjson.ValidBytes(body) { + return false + } + return isOpenAIResponsesLiteHeader(gjson.GetBytes(body, "client_metadata."+responsesLiteWSMetadataKey).String()) +} + // ImageGenerationPermissionMessage returns the stable end-user error text for disabled groups. func ImageGenerationPermissionMessage() string { return imageGenerationPermissionMessage diff --git a/backend/internal/service/model_not_found_error.go b/backend/internal/service/model_not_found_error.go index 910a97d844..de4a004d1e 100644 --- a/backend/internal/service/model_not_found_error.go +++ b/backend/internal/service/model_not_found_error.go @@ -22,6 +22,30 @@ func isModelNotFoundError(statusCode int, body []byte) bool { return isUpstreamModelNotFoundError(statusCode, body) || statusCode == http.StatusNotFound } +// openAICodexPlanGatedModelPhrase matches the deterministic Codex 400 returned +// when a ChatGPT OAuth account's plan cannot serve the requested model, e.g. +// {"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."} +// The phrase is compared against the normalized body (lowercased, "_"/"-" +// folded to spaces), so it also matches the same message embedded in +// error.message-style payloads. +const openAICodexPlanGatedModelPhrase = "model is not supported when using codex" + +// isOpenAICodexPlanGatedModelError reports whether the upstream response is the +// deterministic Codex rejection of a plan-gated model on a ChatGPT account. +// Unlike transient failures, retrying the same account cannot succeed until the +// account's plan changes, so callers should treat it like model-not-found and +// cool the (account, model) pair down instead of re-selecting the account. +func isOpenAICodexPlanGatedModelError(statusCode int, body []byte) bool { + if statusCode != http.StatusBadRequest { + return false + } + normalized := normalizeModelNotFoundBody(body) + if normalized == "" { + return false + } + return strings.Contains(normalized, openAICodexPlanGatedModelPhrase) +} + func containsModelNotFoundKeyword(normalizedBody string) bool { if normalizedBody == "" { return false diff --git a/backend/internal/service/model_not_found_error_test.go b/backend/internal/service/model_not_found_error_test.go index a87340eb55..2f8c83466f 100644 --- a/backend/internal/service/model_not_found_error_test.go +++ b/backend/internal/service/model_not_found_error_test.go @@ -64,3 +64,51 @@ func TestAntigravityModelNotFoundKeepsBare404Fallback(t *testing.T) { t.Fatal("antigravity model-not-found helper should keep bare 404 fallback") } } + +func TestIsOpenAICodexPlanGatedModelError(t *testing.T) { + tests := []struct { + name string + statusCode int + body []byte + want bool + }{ + { + name: "400 codex plan gated detail payload", + statusCode: http.StatusBadRequest, + body: []byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`), + want: true, + }, + { + name: "400 codex plan gated error message payload", + statusCode: http.StatusBadRequest, + body: []byte(`{"error":{"message":"The 'gpt-5.4' model is not supported when using Codex with a ChatGPT account."}}`), + want: true, + }, + { + name: "400 unrelated invalid request does not match", + statusCode: http.StatusBadRequest, + body: []byte(`{"error":{"message":"Invalid schema for response_format 'agentic_plan'"}}`), + want: false, + }, + { + name: "404 with plan gated message does not match", + statusCode: http.StatusNotFound, + body: []byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`), + want: false, + }, + { + name: "400 empty body does not match", + statusCode: http.StatusBadRequest, + body: nil, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isOpenAICodexPlanGatedModelError(tt.statusCode, tt.body); got != tt.want { + t.Fatalf("isOpenAICodexPlanGatedModelError() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/backend/internal/service/oauth_service.go b/backend/internal/service/oauth_service.go index 1369dd9e89..0b3888a73f 100644 --- a/backend/internal/service/oauth_service.go +++ b/backend/internal/service/oauth_service.go @@ -22,6 +22,7 @@ type OpenAIOAuthClient interface { type GrokOAuthClient interface { ExchangeCode(ctx context.Context, code, codeVerifier, redirectURI, proxyURL, clientID string) (*xai.TokenResponse, error) RefreshToken(ctx context.Context, refreshToken, proxyURL, clientID string) (*xai.TokenResponse, error) + ConvertSSOToBuild(ctx context.Context, ssoToken, proxyURL string) (*xai.TokenResponse, error) } // GrokOAuthTokenService is the narrow refresh port used by Grok token providers. diff --git a/backend/internal/service/openai_alpha_search_billing_test.go b/backend/internal/service/openai_alpha_search_billing_test.go index 1251ee43f9..7151725763 100644 --- a/backend/internal/service/openai_alpha_search_billing_test.go +++ b/backend/internal/service/openai_alpha_search_billing_test.go @@ -50,7 +50,7 @@ func TestCalculateOpenAIRecordUsageCostWebSearchPerCall(t *testing.T) { // 即使 token 倍率（含高峰，3.0）更高也不采用。 apiKey := &APIKey{ID: 1, GroupID: &groupID, Group: &Group{ID: groupID, Platform: PlatformOpenAI}} result := &OpenAIForwardResult{Model: "gpt-5.6-sol", UpstreamModel: "gpt-5.6-sol", WebSearchCalls: 1} - cost, err := svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey, []string{"gpt-5.6-sol"}, 3.0, 1.0, 1.0, 2.0, UsageTokens{}, "") + cost, err := svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey, []string{"gpt-5.6-sol"}, 3.0, 1.0, 1.0, 2.0, UsageTokens{}, "", false) require.NoError(t, err) require.Equal(t, string(BillingModePerRequest), cost.BillingMode) require.InDelta(t, 0.01, cost.TotalCost, 1e-12) @@ -58,7 +58,7 @@ func TestCalculateOpenAIRecordUsageCostWebSearchPerCall(t *testing.T) { // 分组配置单价 0.005 apiKey.Group.WebSearchPricePerCall = float64Ptr(0.005) - cost, err = svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey, []string{"gpt-5.6-sol"}, 1.0, 1.0, 1.0, 1.0, UsageTokens{}, "") + cost, err = svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey, []string{"gpt-5.6-sol"}, 1.0, 1.0, 1.0, 1.0, UsageTokens{}, "", false) require.NoError(t, err) require.InDelta(t, 0.005, cost.TotalCost, 1e-12) require.InDelta(t, 0.005, cost.ActualCost, 1e-12) @@ -66,7 +66,7 @@ func TestCalculateOpenAIRecordUsageCostWebSearchPerCall(t *testing.T) { // WebSearchCalls = 0 时不得走按次分支（无定价数据会返回 pricing 错误， // 证明回落到了 token 路径而不是被按次分支吞掉）。 result.WebSearchCalls = 0 - _, err = svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey, []string{"gpt-5.6-sol"}, 1.0, 1.0, 1.0, 1.0, UsageTokens{InputTokens: 10}, "") + _, err = svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey, []string{"gpt-5.6-sol"}, 1.0, 1.0, 1.0, 1.0, UsageTokens{InputTokens: 10}, "", false) require.Error(t, err) } diff --git a/backend/internal/service/openai_codex_models_service.go b/backend/internal/service/openai_codex_models_service.go index 8a919fa2b0..a7331f64c9 100644 --- a/backend/internal/service/openai_codex_models_service.go +++ b/backend/internal/service/openai_codex_models_service.go @@ -2,21 +2,36 @@ package service import ( "context" + "crypto/sha256" + "errors" + "fmt" "io" + "net" "net/http" "net/url" + "sort" "strings" + "sync" "time" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/httpclient" + "golang.org/x/net/http2" + "golang.org/x/sync/singleflight" ) // chatgptCodexModelsURL is the ChatGPT Codex models manifest endpoint. // Package-level variable so tests can point it at a stub server. var chatgptCodexModelsURL = "https://chatgpt.com/backend-api/codex/models" -const codexModelsManifestBodyLimit int64 = 8 << 20 +const ( + codexModelsManifestBodyLimit int64 = 8 << 20 + codexModelsManifestCacheBodyLimit = 1 << 20 + codexModelsManifestCacheMaxEntries = 64 + codexModelsManifestCacheTTL = 30 * time.Second + codexModelsManifestCacheStaleTTL = 5 * time.Minute + codexModelsManifestRequestTimeout = 15 * time.Second +) // CodexModelsManifest carries the raw upstream manifest payload plus caching // metadata so handlers can pass both through to the client untouched. @@ -26,8 +41,180 @@ type CodexModelsManifest struct { NotModified bool } -// FetchCodexModelsManifest fetches the live Codex models manifest from the -// ChatGPT backend using the account's OAuth credentials. +type codexModelsManifestUpstreamError struct { + err error + retryable bool +} + +func (e *codexModelsManifestUpstreamError) Error() string { return e.err.Error() } + +func (e *codexModelsManifestUpstreamError) Unwrap() error { return e.err } + +// IsRetryableCodexModelsManifestError reports whether another selected account +// may succeed without changing the request. Configuration and upstream 4xx +// responses, except 429, are intentionally not retried. +func IsRetryableCodexModelsManifestError(err error) bool { + var upstreamErr *codexModelsManifestUpstreamError + return errors.As(err, &upstreamErr) && upstreamErr.retryable +} + +func isRetryableCodexModelsManifestTransportError(err error) bool { + if err == nil || errors.Is(err, context.Canceled) { + return false + } + if errors.Is(err, context.DeadlineExceeded) || + errors.Is(err, io.EOF) || + errors.Is(err, io.ErrUnexpectedEOF) || + errors.Is(err, net.ErrClosed) { + return true + } + + var opErr *net.OpError + if errors.As(err, &opErr) { + return true + } + var dnsErr *net.DNSError + if errors.As(err, &dnsErr) { + return true + } + var goAwayErr http2.GoAwayError + if errors.As(err, &goAwayErr) { + return true + } + var streamErr http2.StreamError + if errors.As(err, &streamErr) { + return true + } + var connectionErr http2.ConnectionError + if errors.As(err, &connectionErr) { + return true + } + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return true + } + + // net/http uses unexported HTTP/2 error types, so typed matching is not + // possible for errors produced by the standard library transport. + message := strings.ToLower(err.Error()) + if strings.Contains(message, "http2:") && + (strings.Contains(message, "goaway") || + strings.Contains(message, "refused_stream") || + strings.Contains(message, "frame too large")) { + return true + } + if strings.Contains(message, "stream error: stream id ") { + return true + } + for _, code := range []http2.ErrCode{ + http2.ErrCodeNo, + http2.ErrCodeProtocol, + http2.ErrCodeInternal, + http2.ErrCodeFlowControl, + http2.ErrCodeSettingsTimeout, + http2.ErrCodeStreamClosed, + http2.ErrCodeFrameSize, + http2.ErrCodeRefusedStream, + http2.ErrCodeCancel, + http2.ErrCodeCompression, + http2.ErrCodeConnect, + http2.ErrCodeEnhanceYourCalm, + http2.ErrCodeInadequateSecurity, + http2.ErrCodeHTTP11Required, + } { + if strings.Contains(message, "connection error: "+strings.ToLower(code.String())) { + return true + } + } + return false +} + +type codexModelsManifestRequest struct { + url string + headers http.Header + proxyURL string + accountID int64 + credentialAccountID int64 + accountConcurrency int + useAPIKeyUpstream bool +} + +type codexModelsManifestCacheEntry struct { + manifest *CodexModelsManifest + order uint64 + expiresAt time.Time + staleUntil time.Time +} + +type codexModelsManifestCacheState uint8 + +const ( + codexModelsManifestCacheMiss codexModelsManifestCacheState = iota + codexModelsManifestCacheFresh + codexModelsManifestCacheStale +) + +type codexModelsManifestCache struct { + mu sync.Mutex + entries map[string]codexModelsManifestCacheEntry + nextOrder uint64 + refresh singleflight.Group +} + +func (c *codexModelsManifestCache) get(key string, now time.Time) (*CodexModelsManifest, codexModelsManifestCacheState) { + c.mu.Lock() + defer c.mu.Unlock() + entry, ok := c.entries[key] + if !ok { + return nil, codexModelsManifestCacheMiss + } + if !now.Before(entry.staleUntil) { + delete(c.entries, key) + return nil, codexModelsManifestCacheMiss + } + if now.Before(entry.expiresAt) { + return entry.manifest, codexModelsManifestCacheFresh + } + return entry.manifest, codexModelsManifestCacheStale +} + +func (c *codexModelsManifestCache) set(key string, manifest *CodexModelsManifest, now time.Time) { + if manifest == nil || len(manifest.Body) > codexModelsManifestCacheBodyLimit { + return + } + c.mu.Lock() + defer c.mu.Unlock() + if c.entries == nil { + c.entries = make(map[string]codexModelsManifestCacheEntry) + } + if _, exists := c.entries[key]; !exists && len(c.entries) >= codexModelsManifestCacheMaxEntries { + oldestKey := "" + var oldestOrder uint64 + for candidateKey, entry := range c.entries { + if !now.Before(entry.staleUntil) { + delete(c.entries, candidateKey) + continue + } + if oldestKey == "" || entry.order < oldestOrder { + oldestKey = candidateKey + oldestOrder = entry.order + } + } + if len(c.entries) >= codexModelsManifestCacheMaxEntries && oldestKey != "" { + delete(c.entries, oldestKey) + } + } + c.nextOrder++ + c.entries[key] = codexModelsManifestCacheEntry{ + manifest: manifest, + order: c.nextOrder, + expiresAt: now.Add(codexModelsManifestCacheTTL), + staleUntil: now.Add(codexModelsManifestCacheStaleTTL), + } +} + +// FetchCodexModelsManifest fetches the live Codex models manifest from either +// the ChatGPT backend for OAuth accounts or a custom upstream for API key accounts. // // The response body is passed through verbatim: the manifest schema evolves // with Codex client releases, and interpreting it here would force the gateway @@ -41,49 +228,171 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc if err != nil { return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_CREDENTIALS_FAILED", "resolve credential account: %v", err) } - accessToken := credAccount.GetOpenAIAccessToken() - if accessToken == "" { - return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_CODEX_MODELS_TOKEN_MISSING", "account has no Codex backend access token") - } clientVersion = strings.TrimSpace(clientVersion) if clientVersion == "" { clientVersion = openAICodexProbeVersion } - requestURL := chatgptCodexModelsURL + "?client_version=" + url.QueryEscape(clientVersion) - reqCtx, cancel := context.WithTimeout(ctx, 15*time.Second) - defer cancel() - req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, requestURL, nil) + requestEndpoint := chatgptCodexModelsURL + authToken := "" + useAPIKeyUpstream := false + appendModelsPath := false + switch { + case credAccount.IsOpenAIOAuth(): + authToken = strings.TrimSpace(credAccount.GetOpenAIAccessToken()) + if authToken == "" { + return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_CODEX_MODELS_TOKEN_MISSING", "account has no Codex backend access token") + } + case credAccount.IsOpenAIApiKey(): + baseURL := strings.TrimSpace(credAccount.GetCredential("base_url")) + if baseURL == "" || isOfficialOpenAIModelsBaseURL(baseURL) { + return nil, infraerrors.New( + http.StatusBadGateway, + "OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_UNSUPPORTED", + "Codex models manifest requires a custom API key upstream base URL", + ) + } + authToken = strings.TrimSpace(credAccount.GetOpenAIApiKey()) + if authToken == "" { + return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_CODEX_MODELS_API_KEY_MISSING", "account has no API key for the Codex models upstream") + } + normalizedBaseURL, validateErr := s.validateUpstreamBaseURL(baseURL) + if validateErr != nil { + return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_INVALID", "invalid Codex models upstream base URL: %v", validateErr) + } + requestEndpoint = normalizedBaseURL + useAPIKeyUpstream = true + appendModelsPath = true + default: + return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_ACCOUNT_TYPE_UNSUPPORTED", "account type %q cannot fetch the Codex models manifest", credAccount.Type) + } + + requestURL, err := buildCodexModelsManifestURL(requestEndpoint, appendModelsPath, clientVersion) if err != nil { - return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_REQUEST_FAILED", "create codex models request: %v", err) + if useAPIKeyUpstream { + return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_INVALID", "invalid Codex models upstream base URL: %v", err) + } + return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_REQUEST_FAILED", "parse codex models request URL: %v", err) } - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - req.Header.Set("Originator", "codex_cli_rs") - req.Header.Set("Version", clientVersion) - req.Header.Set("User-Agent", codexCLIUserAgent) - if ifNoneMatch = strings.TrimSpace(ifNoneMatch); ifNoneMatch != "" { - req.Header.Set("If-None-Match", ifNoneMatch) + + headers := make(http.Header) + headers.Set("Authorization", "Bearer "+authToken) + headers.Set("Accept", "application/json") + headers.Set("Originator", "codex_cli_rs") + headers.Set("Version", clientVersion) + headers.Set("User-Agent", codexCLIUserAgent) + if useAPIKeyUpstream { + credAccount.ApplyHeaderOverrides(headers) + } else { + setOpenAIChatGPTAccountHeaders(headers, credAccount) } - setOpenAIChatGPTAccountHeaders(req.Header, credAccount) proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { proxyURL = account.Proxy.URL() } - client, err := httpclient.GetClient(httpclient.Options{ - ProxyURL: proxyURL, - Timeout: 15 * time.Second, - ResponseHeaderTimeout: 10 * time.Second, + + request := codexModelsManifestRequest{ + url: requestURL.String(), + headers: headers, + proxyURL: proxyURL, + accountID: account.ID, + credentialAccountID: credAccount.ID, + accountConcurrency: account.Concurrency, + useAPIKeyUpstream: useAPIKeyUpstream, + } + if useAPIKeyUpstream { + return s.fetchCachedAPIKeyCodexModelsManifest(ctx, request, ifNoneMatch) + } + return s.fetchCodexModelsManifestUpstream(ctx, request, ifNoneMatch) +} + +func (s *OpenAIGatewayService) fetchCachedAPIKeyCodexModelsManifest(ctx context.Context, request codexModelsManifestRequest, ifNoneMatch string) (*CodexModelsManifest, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + cacheKey := buildCodexModelsManifestCacheKey(request) + manifest, state := s.codexModelsManifestCache.get(cacheKey, time.Now()) + if state == codexModelsManifestCacheFresh { + return codexModelsManifestForClient(manifest, ifNoneMatch), nil + } + resultCh := s.refreshCachedAPIKeyCodexModelsManifest(cacheKey, request) + if state == codexModelsManifestCacheStale { + return codexModelsManifestForClient(manifest, ifNoneMatch), nil + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case result := <-resultCh: + if result.Err != nil { + return nil, result.Err + } + manifest, ok := result.Val.(*CodexModelsManifest) + if !ok || manifest == nil { + return nil, infraerrors.New(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_REQUEST_FAILED", "invalid shared Codex models manifest result") + } + return codexModelsManifestForClient(manifest, ifNoneMatch), nil + } +} + +func (s *OpenAIGatewayService) refreshCachedAPIKeyCodexModelsManifest(cacheKey string, request codexModelsManifestRequest) <-chan singleflight.Result { + return s.codexModelsManifestCache.refresh.DoChan(cacheKey, func() (any, error) { + cached, _ := s.codexModelsManifestCache.get(cacheKey, time.Now()) + ifNoneMatch := "" + if cached != nil { + ifNoneMatch = cached.ETag + } + manifest, err := s.fetchCodexModelsManifestUpstream(context.Background(), request, ifNoneMatch) + if err != nil { + return nil, err + } + if manifest.NotModified && cached != nil { + s.codexModelsManifestCache.set(cacheKey, cached, time.Now()) + return cached, nil + } + if !manifest.NotModified { + s.codexModelsManifestCache.set(cacheKey, manifest, time.Now()) + } + return manifest, nil }) +} + +func (s *OpenAIGatewayService) fetchCodexModelsManifestUpstream(ctx context.Context, request codexModelsManifestRequest, ifNoneMatch string) (*CodexModelsManifest, error) { + reqCtx, cancel := context.WithTimeout(ctx, codexModelsManifestRequestTimeout) + defer cancel() + req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, request.url, nil) if err != nil { - return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_PROXY_INVALID", "invalid proxy configuration: %v", err) + return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_REQUEST_FAILED", "create codex models request: %v", err) + } + req.Header = request.headers.Clone() + if ifNoneMatch = strings.TrimSpace(ifNoneMatch); ifNoneMatch != "" { + req.Header.Set("If-None-Match", ifNoneMatch) } - resp, err := client.Do(req) + var resp *http.Response + if request.useAPIKeyUpstream { + if s.httpUpstream == nil { + return nil, infraerrors.New(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_UPSTREAM_NOT_CONFIGURED", "Codex models upstream HTTP client is not configured") + } + req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI)) + resp, err = s.httpUpstream.Do(req, request.proxyURL, request.accountID, request.accountConcurrency) + } else { + client, clientErr := httpclient.GetClient(httpclient.Options{ + ProxyURL: request.proxyURL, + Timeout: codexModelsManifestRequestTimeout, + ResponseHeaderTimeout: 10 * time.Second, + }) + if clientErr != nil { + return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_PROXY_INVALID", "invalid proxy configuration: %v", clientErr) + } + resp, err = client.Do(req) + } if err != nil { - return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest request failed: %v", err) + return nil, &codexModelsManifestUpstreamError{ + err: infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest request failed: %v", err), + retryable: isRetryableCodexModelsManifestTransportError(err), + } } defer func() { _ = resp.Body.Close() }() @@ -96,12 +405,100 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc if message == "" { message = resp.Status } - return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest upstream error %d: %s", resp.StatusCode, message) + return nil, &codexModelsManifestUpstreamError{ + err: infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest upstream error %d: %s", resp.StatusCode, message), + retryable: resp.StatusCode == http.StatusTooManyRequests || + (resp.StatusCode >= http.StatusInternalServerError && resp.StatusCode < 600), + } } body, err := io.ReadAll(io.LimitReader(resp.Body, codexModelsManifestBodyLimit)) if err != nil { - return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "read codex models manifest response: %v", err) + return nil, &codexModelsManifestUpstreamError{ + err: infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "read codex models manifest response: %v", err), + retryable: isRetryableCodexModelsManifestTransportError(err), + } } return &CodexModelsManifest{Body: body, ETag: resp.Header.Get("ETag")}, nil } + +func buildCodexModelsManifestCacheKey(request codexModelsManifestRequest) string { + hasher := sha256.New() + _, _ = fmt.Fprintf(hasher, "%d\n%d\n%s\n%s\n", request.accountID, request.credentialAccountID, request.proxyURL, request.url) + headerNames := make([]string, 0, len(request.headers)) + for name := range request.headers { + headerNames = append(headerNames, name) + } + sort.Strings(headerNames) + for _, name := range headerNames { + _, _ = fmt.Fprintf(hasher, "%s\n", strings.ToLower(name)) + for _, value := range request.headers[name] { + _, _ = fmt.Fprintf(hasher, "%s\n", value) + } + } + return fmt.Sprintf("%x", hasher.Sum(nil)) +} + +func codexModelsManifestForClient(manifest *CodexModelsManifest, ifNoneMatch string) *CodexModelsManifest { + if manifest == nil { + return nil + } + if codexModelsManifestETagMatches(ifNoneMatch, manifest.ETag) { + return &CodexModelsManifest{ETag: manifest.ETag, NotModified: true} + } + return manifest +} + +func codexModelsManifestETagMatches(ifNoneMatch, etag string) bool { + etag = strings.TrimSpace(etag) + if etag == "" { + return false + } + normalize := func(value string) string { + value = strings.TrimSpace(value) + if len(value) >= 2 && strings.EqualFold(value[:2], "W/") { + value = strings.TrimSpace(value[2:]) + } + return value + } + want := normalize(etag) + for _, candidate := range strings.Split(ifNoneMatch, ",") { + candidate = strings.TrimSpace(candidate) + if candidate == "*" || normalize(candidate) == want { + return true + } + } + return false +} + +func isOfficialOpenAIModelsBaseURL(raw string) bool { + parsed, err := url.Parse(strings.TrimSpace(raw)) + if err != nil { + return false + } + hostname := strings.TrimSuffix(parsed.Hostname(), ".") + return strings.EqualFold(hostname, "api.openai.com") +} + +func buildCodexModelsManifestURL(endpoint string, appendModelsPath bool, clientVersion string) (*url.URL, error) { + requestURL, err := url.Parse(endpoint) + if err != nil { + return nil, err + } + if requestURL.Fragment != "" { + return nil, fmt.Errorf("URL fragments are not supported") + } + + query := requestURL.Query() + requestURL.RawQuery = "" + requestURL.ForceQuery = false + if appendModelsPath { + requestURL, err = url.Parse(buildOpenAIModelsURL(requestURL.String())) + if err != nil { + return nil, err + } + } + query.Set("client_version", clientVersion) + requestURL.RawQuery = query.Encode() + return requestURL, nil +} diff --git a/backend/internal/service/openai_codex_models_service_test.go b/backend/internal/service/openai_codex_models_service_test.go index c9eae35629..2c87861f60 100644 --- a/backend/internal/service/openai_codex_models_service_test.go +++ b/backend/internal/service/openai_codex_models_service_test.go @@ -2,11 +2,146 @@ package service import ( "context" + "errors" + "io" + "net" "net/http" "net/http/httptest" + "net/url" + "strings" + "sync" + "sync/atomic" "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "golang.org/x/net/http2" ) +type codexModelsHTTPUpstreamStub struct { + do func(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) +} + +type codexModelsBlockingBody struct { + ctx context.Context + readStarted chan struct{} + startedOnce *sync.Once + release <-chan struct{} + body *strings.Reader +} + +func (b *codexModelsBlockingBody) Read(p []byte) (int, error) { + b.startedOnce.Do(func() { close(b.readStarted) }) + select { + case <-b.release: + return b.body.Read(p) + case <-b.ctx.Done(): + return 0, b.ctx.Err() + } +} + +func (b *codexModelsBlockingBody) Close() error { return nil } + +func (s *codexModelsHTTPUpstreamStub) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) { + return s.do(req, proxyURL, accountID, accountConcurrency) +} + +func (s *codexModelsHTTPUpstreamStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) { + return s.Do(req, proxyURL, accountID, accountConcurrency) +} + +func TestIsRetryableCodexModelsManifestTransportError(t *testing.T) { + tests := []struct { + name string + err error + retryable bool + }{ + {name: "nil", err: nil}, + {name: "configuration error", err: errors.New("invalid proxy URL")}, + {name: "upstream configuration error", err: errors.New("upstream error: invalid proxy")}, + {name: "proxy connection configuration error", err: errors.New("proxy connection error: invalid configuration")}, + {name: "canceled request", err: context.Canceled}, + { + name: "redirect policy error", + err: &url.Error{ + Op: "Get", + URL: "https://upstream.example/v1/models", + Err: errors.New("stopped after 10 redirects"), + }, + }, + {name: "deadline exceeded", err: context.DeadlineExceeded, retryable: true}, + {name: "unexpected EOF", err: io.ErrUnexpectedEOF, retryable: true}, + {name: "closed connection", err: net.ErrClosed, retryable: true}, + { + name: "network operation", + err: &net.OpError{ + Op: "read", + Net: "tcp", + Err: errors.New("connection reset"), + }, + retryable: true, + }, + { + name: "DNS error", + err: &net.DNSError{Err: "temporary failure", Name: "upstream.example"}, + retryable: true, + }, + { + name: "typed HTTP2 GOAWAY", + err: http2.GoAwayError{ErrCode: http2.ErrCodeNo}, + retryable: true, + }, + { + name: "stdlib HTTP2 GOAWAY", + err: errors.New("http2: server sent GOAWAY and closed the connection; LastStreamID=1, ErrCode=NO_ERROR"), + retryable: true, + }, + { + name: "stdlib HTTP2 refused stream", + err: errors.New("stream error: stream ID 3; REFUSED_STREAM"), + retryable: true, + }, + { + name: "stdlib HTTP2 connection error", + err: errors.New(`Get "https://upstream.example/v1/models": connection error: PROTOCOL_ERROR`), + retryable: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isRetryableCodexModelsManifestTransportError(tt.err); got != tt.retryable { + t.Fatalf("retryable = %v, want %v", got, tt.retryable) + } + }) + } +} + +func newCodexModelsAPIKeyTestService(upstream HTTPUpstream) *OpenAIGatewayService { + return &OpenAIGatewayService{ + cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{ + Enabled: false, + }}}, + httpUpstream: upstream, + } +} + +func newCodexModelsAPIKeyTestAccount(baseURL string) *Account { + credentials := map[string]any{"api_key": "sk-upstream"} + if baseURL != "" { + credentials["base_url"] = baseURL + } + return &Account{ + ID: 2, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: credentials, + Concurrency: 3, + } +} + func newCodexModelsTestAccount() *Account { return &Account{ ID: 1, @@ -136,3 +271,679 @@ func TestFetchCodexModelsManifestMissingToken(t *testing.T) { t.Fatal("expected error for missing access token, got nil") } } + +func TestFetchCodexModelsManifestAPIKeyCustomUpstream(t *testing.T) { + manifestBody := `{"models":[{"slug":"gpt-5.6"}]}` + var gotRequest *http.Request + var gotProxyURL string + var gotAccountID int64 + var gotConcurrency int + upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) { + gotRequest = req + gotProxyURL = proxyURL + gotAccountID = accountID + gotConcurrency = accountConcurrency + header := make(http.Header) + header.Set("ETag", `W/"api-key-manifest"`) + return &http.Response{ + StatusCode: http.StatusOK, + Header: header, + Body: io.NopCloser(strings.NewReader(manifestBody)), + }, nil + }} + + s := newCodexModelsAPIKeyTestService(upstream) + manifest, err := s.FetchCodexModelsManifest( + context.Background(), + newCodexModelsAPIKeyTestAccount("https://upstream.example/v1"), + "0.144.0", + "", + ) + if err != nil { + t.Fatalf("FetchCodexModelsManifest returned error: %v", err) + } + + if gotRequest == nil { + t.Fatal("expected request to custom API key upstream") + } + if gotRequest.Method != http.MethodGet { + t.Errorf("method: got %q", gotRequest.Method) + } + if gotRequest.URL.String() != "https://upstream.example/v1/models?client_version=0.144.0" { + t.Errorf("request URL: got %q", gotRequest.URL.String()) + } + if gotRequest.Header.Get("Authorization") != "Bearer sk-upstream" { + t.Errorf("authorization header: got %q", gotRequest.Header.Get("Authorization")) + } + if gotRequest.Header.Get("Originator") != "codex_cli_rs" { + t.Errorf("originator header: got %q", gotRequest.Header.Get("Originator")) + } + if gotRequest.Header.Get("Version") != "0.144.0" { + t.Errorf("version header: got %q", gotRequest.Header.Get("Version")) + } + if gotRequest.Header.Get("User-Agent") != codexCLIUserAgent { + t.Errorf("user-agent header: got %q", gotRequest.Header.Get("User-Agent")) + } + if gotRequest.Header.Get("chatgpt-account-id") != "" { + t.Errorf("chatgpt-account-id must not be sent to API key upstream: got %q", gotRequest.Header.Get("chatgpt-account-id")) + } + if gotProxyURL != "" || gotAccountID != 2 || gotConcurrency != 3 { + t.Errorf("upstream routing metadata: proxy=%q account_id=%d concurrency=%d", gotProxyURL, gotAccountID, gotConcurrency) + } + if string(manifest.Body) != manifestBody { + t.Errorf("body not passed through verbatim: got %q", manifest.Body) + } + if manifest.ETag != `W/"api-key-manifest"` { + t.Errorf("etag not passed through: got %q", manifest.ETag) + } +} + +func TestFetchCodexModelsManifestAPIKeySharedRefreshSurvivesCallerCancellation(t *testing.T) { + const manifestBody = `{"models":[{"slug":"gpt-5.6"}]}` + var calls atomic.Int32 + var readStartedOnce sync.Once + readStarted := make(chan struct{}) + deadlineRemaining := make(chan time.Duration, 1) + release := make(chan struct{}) + upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + calls.Add(1) + deadline, ok := req.Context().Deadline() + if !ok { + deadlineRemaining <- 0 + } else { + deadlineRemaining <- time.Until(deadline) + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Etag": []string{`W/"shared"`}}, + Body: &codexModelsBlockingBody{ + ctx: req.Context(), + readStarted: readStarted, + startedOnce: &readStartedOnce, + release: release, + body: strings.NewReader(manifestBody), + }, + }, nil + }} + + s := newCodexModelsAPIKeyTestService(upstream) + account := newCodexModelsAPIKeyTestAccount("https://upstream.example") + firstCtx, cancelFirst := context.WithCancel(context.Background()) + firstErr := make(chan error, 1) + go func() { + _, err := s.FetchCodexModelsManifest(firstCtx, account, "0.144.0", "") + firstErr <- err + }() + + select { + case <-readStarted: + case <-time.After(time.Second): + t.Fatal("upstream body read did not start") + } + remaining := <-deadlineRemaining + if remaining < 14*time.Second || remaining > codexModelsManifestRequestTimeout { + t.Errorf("detached refresh deadline: got %s, want approximately %s", remaining, codexModelsManifestRequestTimeout) + } + cancelFirst() + select { + case err := <-firstErr: + if !errors.Is(err, context.Canceled) { + t.Fatalf("first caller error: got %v, want context.Canceled", err) + } + case <-time.After(time.Second): + t.Fatal("canceled caller did not return promptly") + } + + secondResult := make(chan struct { + manifest *CodexModelsManifest + err error + }, 1) + go func() { + manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") + secondResult <- struct { + manifest *CodexModelsManifest + err error + }{manifest: manifest, err: err} + }() + + time.Sleep(50 * time.Millisecond) + if got := calls.Load(); got != 1 { + t.Errorf("upstream calls before shared refresh completed: got %d, want 1", got) + } + close(release) + select { + case result := <-secondResult: + if result.err != nil { + t.Fatalf("second caller returned error: %v", result.err) + } + if string(result.manifest.Body) != manifestBody { + t.Errorf("second caller body: got %q", result.manifest.Body) + } + case <-time.After(time.Second): + t.Fatal("second caller did not receive shared refresh result") + } + if got := calls.Load(); got != 1 { + t.Errorf("total upstream calls: got %d, want 1", got) + } +} + +func TestFetchCodexModelsManifestAPIKeyConcurrentRequestsShareRefresh(t *testing.T) { + const callers = 8 + var calls atomic.Int32 + started := make(chan struct{}) + var startedOnce sync.Once + release := make(chan struct{}) + upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + calls.Add(1) + startedOnce.Do(func() { close(started) }) + <-release + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"models":[]}`)), + }, nil + }} + + s := newCodexModelsAPIKeyTestService(upstream) + account := newCodexModelsAPIKeyTestAccount("https://upstream.example") + begin := make(chan struct{}) + errs := make(chan error, callers) + for i := 0; i < callers; i++ { + go func() { + <-begin + _, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") + errs <- err + }() + } + close(begin) + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("upstream request did not start") + } + time.Sleep(50 * time.Millisecond) + if got := calls.Load(); got != 1 { + t.Errorf("concurrent upstream calls: got %d, want 1", got) + } + close(release) + for i := 0; i < callers; i++ { + if err := <-errs; err != nil { + t.Errorf("caller %d returned error: %v", i, err) + } + } +} + +func TestFetchCodexModelsManifestAPIKeyFreshCacheHandlesETagLocally(t *testing.T) { + var calls atomic.Int32 + upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + calls.Add(1) + if got := req.Header.Get("If-None-Match"); got != "" { + t.Errorf("cache refresh must not inherit a caller's If-None-Match: got %q", got) + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Etag": []string{`W/"cached"`}}, + Body: io.NopCloser(strings.NewReader(`{"models":[]}`)), + }, nil + }} + + s := newCodexModelsAPIKeyTestService(upstream) + account := newCodexModelsAPIKeyTestAccount("https://upstream.example") + if _, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", ""); err != nil { + t.Fatalf("initial fetch returned error: %v", err) + } + manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", `W/"cached"`) + if err != nil { + t.Fatalf("cached fetch returned error: %v", err) + } + if !manifest.NotModified { + t.Fatal("matching cached ETag must return NotModified") + } + if got := calls.Load(); got != 1 { + t.Errorf("upstream calls: got %d, want 1", got) + } +} + +func TestFetchCodexModelsManifestAPIKeyCacheKeyIsolatesRequestIdentity(t *testing.T) { + var calls atomic.Int32 + upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + calls.Add(1) + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"models":[]}`)), + }, nil + }} + s := newCodexModelsAPIKeyTestService(upstream) + + base := newCodexModelsAPIKeyTestAccount("https://upstream.example") + fetch := func(account *Account, version string) { + t.Helper() + if _, err := s.FetchCodexModelsManifest(context.Background(), account, version, ""); err != nil { + t.Fatalf("fetch returned error: %v", err) + } + } + fetch(base, "0.144.0") + fetch(base, "0.144.0") + + differentAccount := newCodexModelsAPIKeyTestAccount("https://upstream.example") + differentAccount.ID = 3 + fetch(differentAccount, "0.144.0") + + differentToken := newCodexModelsAPIKeyTestAccount("https://upstream.example") + differentToken.Credentials["api_key"] = "sk-other" + fetch(differentToken, "0.144.0") + + differentUpstream := newCodexModelsAPIKeyTestAccount("https://other-upstream.example") + fetch(differentUpstream, "0.144.0") + fetch(base, "0.145.0") + + differentHeaders := newCodexModelsAPIKeyTestAccount("https://upstream.example") + differentHeaders.Credentials[credKeyHeaderOverrideEnabled] = true + differentHeaders.Credentials[credKeyHeaderOverrides] = map[string]any{"x-tenant": "other"} + fetch(differentHeaders, "0.144.0") + + proxyID := int64(9) + differentProxy := newCodexModelsAPIKeyTestAccount("https://upstream.example") + differentProxy.ProxyID = &proxyID + differentProxy.Proxy = &Proxy{Protocol: "http", Host: "127.0.0.1", Port: 8080} + fetch(differentProxy, "0.144.0") + fetch(differentProxy, "0.144.0") + + if got := calls.Load(); got != 7 { + t.Errorf("isolated upstream calls: got %d, want 7", got) + } +} + +func TestFetchCodexModelsManifestAPIKeyCacheBoundsEntriesAndBodySize(t *testing.T) { + var calls atomic.Int32 + upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + calls.Add(1) + body := `{"models":[]}` + if strings.Contains(req.URL.Host, "large") { + body = strings.Repeat("x", (1<<20)+1) + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + }, nil + }} + s := newCodexModelsAPIKeyTestService(upstream) + fetch := func(account *Account) { + t.Helper() + if _, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", ""); err != nil { + t.Fatalf("fetch returned error: %v", err) + } + } + + small := newCodexModelsAPIKeyTestAccount("https://small.example") + fetch(small) + fetch(small) + large := newCodexModelsAPIKeyTestAccount("https://large.example") + large.ID = 3 + fetch(large) + fetch(large) + if got := calls.Load(); got != 3 { + t.Fatalf("body-size bounded cache calls: got %d, want 3", got) + } + + for i := int64(10); i < 75; i++ { + account := newCodexModelsAPIKeyTestAccount("https://bounded.example") + account.ID = i + fetch(account) + } + last := newCodexModelsAPIKeyTestAccount("https://bounded.example") + last.ID = 74 + fetch(last) + if got := calls.Load(); got != 68 { + t.Fatalf("most recent cache entry was not retained: calls=%d, want 68", got) + } + first := newCodexModelsAPIKeyTestAccount("https://bounded.example") + first.ID = 10 + fetch(first) + if got := calls.Load(); got != 69 { + t.Errorf("oldest cache entry was not evicted: calls=%d, want 69", got) + } +} + +func TestFetchCodexModelsManifestAPIKeyServesStaleWhileRefreshing(t *testing.T) { + var calls atomic.Int32 + refreshStarted := make(chan struct{}) + releaseRefresh := make(chan struct{}) + upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + call := calls.Add(1) + body := `{"models":[{"slug":"old"}]}` + if call > 1 { + if call == 2 { + close(refreshStarted) + } + <-releaseRefresh + body = `{"models":[{"slug":"new"}]}` + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + }, nil + }} + s := newCodexModelsAPIKeyTestService(upstream) + account := newCodexModelsAPIKeyTestAccount("https://upstream.example") + if _, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", ""); err != nil { + t.Fatalf("initial fetch returned error: %v", err) + } + + s.codexModelsManifestCache.mu.Lock() + for key, entry := range s.codexModelsManifestCache.entries { + entry.expiresAt = time.Now().Add(-time.Second) + s.codexModelsManifestCache.entries[key] = entry + } + s.codexModelsManifestCache.mu.Unlock() + + resultCh := make(chan struct { + manifest *CodexModelsManifest + err error + }, 1) + go func() { + manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") + resultCh <- struct { + manifest *CodexModelsManifest + err error + }{manifest: manifest, err: err} + }() + select { + case <-refreshStarted: + case <-time.After(time.Second): + t.Fatal("background refresh did not start") + } + + var staleResult struct { + manifest *CodexModelsManifest + err error + } + select { + case staleResult = <-resultCh: + case <-time.After(100 * time.Millisecond): + t.Error("stale manifest was not returned while refresh was blocked") + close(releaseRefresh) + staleResult = <-resultCh + } + if staleResult.err != nil { + t.Fatalf("stale fetch returned error: %v", staleResult.err) + } + if got := string(staleResult.manifest.Body); got != `{"models":[{"slug":"old"}]}` { + t.Errorf("stale body: got %q", got) + } + if got := calls.Load(); got != 2 { + t.Errorf("upstream calls during stale refresh: got %d, want 2", got) + } + + select { + case <-releaseRefresh: + default: + close(releaseRefresh) + } + deadline := time.Now().Add(time.Second) + for { + manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") + if err == nil && string(manifest.Body) == `{"models":[{"slug":"new"}]}` { + break + } + if time.Now().After(deadline) { + t.Fatalf("refreshed manifest was not cached: manifest=%v err=%v", manifest, err) + } + time.Sleep(10 * time.Millisecond) + } + if got := calls.Load(); got != 2 { + t.Errorf("stale refresh was not deduplicated: calls=%d, want 2", got) + } +} + +func TestFetchCodexModelsManifestAPIKeyRevalidatesStaleETag(t *testing.T) { + var calls atomic.Int32 + refreshDone := make(chan struct{}) + upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + call := calls.Add(1) + if call == 1 { + header := make(http.Header) + header.Set("ETag", `W/"cached"`) + return &http.Response{ + StatusCode: http.StatusOK, + Header: header, + Body: io.NopCloser(strings.NewReader(`{"models":[{"slug":"cached"}]}`)), + }, nil + } + if got := req.Header.Get("If-None-Match"); got != `W/"cached"` { + t.Errorf("background revalidation If-None-Match: got %q", got) + } + close(refreshDone) + header := make(http.Header) + header.Set("ETag", `W/"cached"`) + return &http.Response{StatusCode: http.StatusNotModified, Header: header, Body: http.NoBody}, nil + }} + s := newCodexModelsAPIKeyTestService(upstream) + account := newCodexModelsAPIKeyTestAccount("https://upstream.example") + if _, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", ""); err != nil { + t.Fatalf("initial fetch returned error: %v", err) + } + s.codexModelsManifestCache.mu.Lock() + for key, entry := range s.codexModelsManifestCache.entries { + entry.expiresAt = time.Now().Add(-time.Second) + s.codexModelsManifestCache.entries[key] = entry + } + s.codexModelsManifestCache.mu.Unlock() + + manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") + if err != nil { + t.Fatalf("stale fetch returned error: %v", err) + } + if got := string(manifest.Body); got != `{"models":[{"slug":"cached"}]}` { + t.Fatalf("stale body: got %q", got) + } + select { + case <-refreshDone: + case <-time.After(time.Second): + t.Fatal("ETag revalidation did not complete") + } + + deadline := time.Now().Add(time.Second) + for { + s.codexModelsManifestCache.mu.Lock() + fresh := false + for _, entry := range s.codexModelsManifestCache.entries { + fresh = time.Now().Before(entry.expiresAt) + } + s.codexModelsManifestCache.mu.Unlock() + if fresh { + break + } + if time.Now().After(deadline) { + t.Fatal("304 revalidation did not renew the cached manifest") + } + time.Sleep(10 * time.Millisecond) + } + manifest, err = s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") + if err != nil || string(manifest.Body) != `{"models":[{"slug":"cached"}]}` { + t.Fatalf("renewed cached manifest: body=%q err=%v", manifest.Body, err) + } + if got := calls.Load(); got != 2 { + t.Errorf("upstream calls: got %d, want 2", got) + } +} + +func TestFetchCodexModelsManifestAPIKeyColdCacheHandlesNotModifiedLocally(t *testing.T) { + var gotIfNoneMatch string + upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + gotIfNoneMatch = req.Header.Get("If-None-Match") + header := make(http.Header) + header.Set("ETag", `W/"api-key-manifest"`) + return &http.Response{ + StatusCode: http.StatusOK, + Header: header, + Body: io.NopCloser(strings.NewReader(`{"models":[]}`)), + }, nil + }} + + s := newCodexModelsAPIKeyTestService(upstream) + manifest, err := s.FetchCodexModelsManifest( + context.Background(), + newCodexModelsAPIKeyTestAccount("https://upstream.example"), + "0.144.0", + `W/"api-key-manifest"`, + ) + if err != nil { + t.Fatalf("FetchCodexModelsManifest returned error: %v", err) + } + if !manifest.NotModified { + t.Error("expected NotModified to be true") + } + if manifest.ETag != `W/"api-key-manifest"` { + t.Errorf("etag not passed through: got %q", manifest.ETag) + } + if gotIfNoneMatch != "" { + t.Errorf("cold shared refresh must not inherit caller if-none-match: got %q", gotIfNoneMatch) + } +} + +func TestFetchCodexModelsManifestAPIKeyDoesNotCacheUnexpectedColdNotModified(t *testing.T) { + var calls atomic.Int32 + upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + calls.Add(1) + if got := req.Header.Get("If-None-Match"); got != "" { + t.Errorf("cold shared refresh If-None-Match: got %q", got) + } + header := make(http.Header) + header.Set("ETag", `W/"unexpected"`) + return &http.Response{StatusCode: http.StatusNotModified, Header: header, Body: http.NoBody}, nil + }} + s := newCodexModelsAPIKeyTestService(upstream) + account := newCodexModelsAPIKeyTestAccount("https://upstream.example") + for i := 0; i < 2; i++ { + manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") + if err != nil { + t.Fatalf("fetch %d returned error: %v", i, err) + } + if !manifest.NotModified { + t.Fatalf("fetch %d: expected upstream NotModified response", i) + } + } + if got := calls.Load(); got != 2 { + t.Errorf("unexpected cold 304 was cached: upstream calls=%d, want 2", got) + } +} + +func TestFetchCodexModelsManifestAPIKeyPreservesBaseURLQuery(t *testing.T) { + var gotURL string + upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + gotURL = req.URL.String() + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"models":[]}`)), + }, nil + }} + + s := newCodexModelsAPIKeyTestService(upstream) + _, err := s.FetchCodexModelsManifest( + context.Background(), + newCodexModelsAPIKeyTestAccount("https://upstream.example/v1?tenant=acme"), + "0.144.0", + "", + ) + if err != nil { + t.Fatalf("FetchCodexModelsManifest returned error: %v", err) + } + if gotURL != "https://upstream.example/v1/models?client_version=0.144.0&tenant=acme" { + t.Errorf("request URL: got %q", gotURL) + } +} + +func TestFetchCodexModelsManifestAPIKeyRejectsBaseURLFragment(t *testing.T) { + called := false + upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + called = true + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"models":[]}`)), + }, nil + }} + + s := newCodexModelsAPIKeyTestService(upstream) + _, err := s.FetchCodexModelsManifest( + context.Background(), + newCodexModelsAPIKeyTestAccount("https://upstream.example/v1#models"), + "0.144.0", + "", + ) + if err == nil { + t.Fatal("expected invalid upstream base URL error, got nil") + } + if infraerrors.Reason(err) != "OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_INVALID" { + t.Errorf("error reason: got %q", infraerrors.Reason(err)) + } + if called { + t.Fatal("fragment-bearing base URL must be rejected before the upstream request") + } +} + +func TestFetchCodexModelsManifestAPIKeyUpstreamError(t *testing.T) { + upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Status: "429 Too Many Requests", + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"error":"rate limited"}`)), + }, nil + }} + + s := newCodexModelsAPIKeyTestService(upstream) + _, err := s.FetchCodexModelsManifest( + context.Background(), + newCodexModelsAPIKeyTestAccount("https://upstream.example"), + "0.144.0", + "", + ) + if err == nil { + t.Fatal("expected error for upstream 429, got nil") + } + if infraerrors.Code(err) != http.StatusBadGateway { + t.Errorf("error status: got %d, want %d", infraerrors.Code(err), http.StatusBadGateway) + } + if infraerrors.Reason(err) != "OPENAI_CODEX_MODELS_UPSTREAM_FAILED" { + t.Errorf("error reason: got %q", infraerrors.Reason(err)) + } +} + +func TestFetchCodexModelsManifestAPIKeyRejectsOfficialOpenAIBaseURL(t *testing.T) { + tests := []struct { + name string + baseURL string + }{ + {name: "missing base URL"}, + {name: "official host", baseURL: "https://api.openai.com"}, + {name: "official versioned URL", baseURL: "https://API.OPENAI.COM:443/v1/"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := newCodexModelsAPIKeyTestService(&codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + t.Fatal("official OpenAI API key must not be used as a Codex manifest upstream") + return nil, nil + }}) + + _, err := s.FetchCodexModelsManifest( + context.Background(), + newCodexModelsAPIKeyTestAccount(tt.baseURL), + "0.144.0", + "", + ) + if err == nil { + t.Fatal("expected unsupported API key upstream error, got nil") + } + if infraerrors.Reason(err) != "OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_UNSUPPORTED" { + t.Errorf("error reason: got %q", infraerrors.Reason(err)) + } + }) + } +} diff --git a/backend/internal/service/openai_codex_transform.go b/backend/internal/service/openai_codex_transform.go index 3869e97c99..3bde4abcfb 100644 --- a/backend/internal/service/openai_codex_transform.go +++ b/backend/internal/service/openai_codex_transform.go @@ -838,6 +838,9 @@ func ensureOpenAIResponsesImageGenerationTool(reqBody map[string]any) bool { if isCodexSparkModel(firstNonEmptyString(reqBody["model"])) { return false } + if hasOpenAIImageGenerationTool(reqBody) { + return false + } tool := map[string]any{ "type": "image_generation", @@ -855,16 +858,6 @@ func ensureOpenAIResponsesImageGenerationTool(reqBody map[string]any) bool { reqBody["tools"] = []any{tool} return true } - for _, rawTool := range tools { - toolMap, ok := rawTool.(map[string]any) - if !ok { - continue - } - if strings.TrimSpace(firstNonEmptyString(toolMap["type"])) == "image_generation" { - return false - } - } - reqBody["tools"] = append(tools, tool) return true } diff --git a/backend/internal/service/openai_codex_transform_test.go b/backend/internal/service/openai_codex_transform_test.go index b226655eeb..456740136b 100644 --- a/backend/internal/service/openai_codex_transform_test.go +++ b/backend/internal/service/openai_codex_transform_test.go @@ -617,6 +617,65 @@ func TestEnsureOpenAIResponsesImageGenerationTool_PreservesExistingImageTool(t * require.Equal(t, "webp", tool["output_format"]) } +func TestEnsureOpenAIResponsesImageGenerationTool_PreservesImageGenNamespace(t *testing.T) { + tests := []struct { + name string + reqBody map[string]any + }{ + { + name: "top-level tools", + reqBody: map[string]any{ + "model": "gpt-5.5", + "tools": []any{ + map[string]any{ + "type": "namespace", + "name": "image_gen", + "tools": []any{ + map[string]any{"type": "function", "name": "imagegen"}, + }, + }, + }, + }, + }, + { + name: "responses lite additional_tools", + reqBody: map[string]any{ + "model": "gpt-5.5", + "input": []any{ + map[string]any{ + "type": "additional_tools", + "tools": []any{ + map[string]any{ + "type": "namespace", + "name": "image_gen", + "tools": []any{ + map[string]any{"type": "function", "name": "imagegen"}, + }, + }, + }, + }, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.True(t, hasOpenAIImageGenerationTool(tt.reqBody)) + + modified := ensureOpenAIResponsesImageGenerationTool(tt.reqBody) + + require.False(t, modified) + tools, _ := tt.reqBody["tools"].([]any) + for _, rawTool := range tools { + tool, ok := rawTool.(map[string]any) + require.True(t, ok) + require.NotEqual(t, "image_generation", firstNonEmptyString(tool["type"])) + } + }) + } +} + func TestApplyCodexImageGenerationBridgeInstructions_AppendsBridgeOnce(t *testing.T) { reqBody := map[string]any{ "model": "gpt-5.4", diff --git a/backend/internal/service/openai_compat_model_test.go b/backend/internal/service/openai_compat_model_test.go index 69b6ddbca2..e1007c507a 100644 --- a/backend/internal/service/openai_compat_model_test.go +++ b/backend/internal/service/openai_compat_model_test.go @@ -124,6 +124,55 @@ func TestApplyOpenAICompatModelNormalization(t *testing.T) { }) } +func TestForwardAsAnthropic_UsesExactFableMessagesDispatchModel(t *testing.T) { + t.Parallel() + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := []byte(`{"model":"claude-fable-5","max_tokens":16,"messages":[{"role":"user","content":"hello"}],"stream":false}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + upstreamBody := strings.Join([]string{ + `data: {"type":"response.completed","response":{"id":"resp_fable","object":"response","model":"gpt-5.6-sol","status":"completed","output":[{"type":"message","id":"msg_fable","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7}}}`, + "", + "data: [DONE]", + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_fable"}}, + Body: io.NopCloser(strings.NewReader(upstreamBody)), + }} + + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}}, + } + account := &Account{ + ID: 1, + Name: "openai-oauth", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "oauth-token", + "chatgpt_account_id": "chatgpt-acc", + }, + } + + result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "gpt-5.6-sol") + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "claude-fable-5", result.Model) + require.Equal(t, "gpt-5.6-sol", result.BillingModel) + require.Equal(t, "gpt-5.6-sol", result.UpstreamModel) + require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(upstream.lastBody, "model").String()) + require.NotContains(t, string(upstream.lastBody), "claude-fable-5") + require.Equal(t, "claude-fable-5", gjson.GetBytes(rec.Body.Bytes(), "model").String()) +} + func TestForwardAsAnthropic_NormalizesRoutingAndEffortForGpt54XHigh(t *testing.T) { t.Parallel() gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_content_session_seed.go b/backend/internal/service/openai_content_session_seed.go index 7c2ba25140..fce85f11bd 100644 --- a/backend/internal/service/openai_content_session_seed.go +++ b/backend/internal/service/openai_content_session_seed.go @@ -11,6 +11,10 @@ import ( // and explicit session IDs (e.g. "sess-xxx" or "compat_cc_xxx"). const contentSessionSeedPrefix = "compat_cs_" +// contentStablePrefixSessionSeedPrefix distinguishes cache identities derived +// only from request fields that remain stable across independent prompts. +const contentStablePrefixSessionSeedPrefix = "compat_csp_" + // deriveOpenAIContentSessionSeed builds a stable session seed from an // OpenAI-format request body. Only fields constant across conversation turns // are included: model, tools/functions definitions, system/developer prompts, @@ -105,3 +109,156 @@ func deriveOpenAIContentSessionSeed(body []byte) string { } return contentSessionSeedPrefix + b.String() } + +// deriveOpenAIAnchoredContentSessionSeed returns the legacy content-derived +// seed only when it contains a meaningful user/input anchor. This preserves +// the existing session derivation while preventing model-only requests from +// becoming a tenant-wide cache routing identity. +func deriveOpenAIAnchoredContentSessionSeed(body []byte) string { + if !hasOpenAIContentSessionUserAnchor(body) { + return "" + } + return deriveOpenAIContentSessionSeed(body) +} + +func hasOpenAIContentSessionUserAnchor(body []byte) bool { + if len(body) == 0 { + return false + } + + if messages := gjson.GetBytes(body, "messages"); messages.Exists() && messages.IsArray() { + anchored := false + messages.ForEach(func(_, message gjson.Result) bool { + if strings.TrimSpace(message.Get("role").String()) != "user" { + return true + } + anchored = hasMeaningfulOpenAIContent(message.Get("content")) + return false + }) + return anchored + } + + input := gjson.GetBytes(body, "input") + if !input.Exists() { + return false + } + if input.Type == gjson.String { + return strings.TrimSpace(input.String()) != "" + } + if !input.IsArray() { + return false + } + + anchored := false + input.ForEach(func(_, item gjson.Result) bool { + if strings.TrimSpace(item.Get("role").String()) == "user" { + anchored = hasMeaningfulOpenAIContent(item.Get("content")) + return false + } + if strings.TrimSpace(item.Get("type").String()) == "input_text" { + anchored = strings.TrimSpace(item.Get("text").String()) != "" + return false + } + return true + }) + return anchored +} + +func hasMeaningfulOpenAIContent(content gjson.Result) bool { + if !content.Exists() || content.Type == gjson.Null { + return false + } + if content.Type == gjson.String { + return strings.TrimSpace(content.String()) != "" + } + if !content.IsArray() { + normalized, ok := normalizeNonEmptyCompatSeedJSON(content) + return ok && strings.TrimSpace(normalized) != "" + } + + meaningful := false + content.ForEach(func(_, item gjson.Result) bool { + if item.Type == gjson.String { + meaningful = strings.TrimSpace(item.String()) != "" + } else if text := item.Get("text"); text.Exists() { + meaningful = strings.TrimSpace(text.String()) != "" + } else { + _, meaningful = normalizeNonEmptyCompatSeedJSON(item) + } + return !meaningful + }) + return meaningful +} + +// deriveOpenAIStablePrefixSessionSeed builds a seed from the reusable prefix +// of an OpenAI-format request. User and assistant content are deliberately +// excluded so independent prompts with the same system/tool prefix can share +// an upstream prompt-cache routing identity. +// +// An empty result means the request has no meaningful stable prefix. Callers +// must then use a narrower fallback instead of grouping all requests by tenant +// and model alone. +func deriveOpenAIStablePrefixSessionSeed(body []byte) string { + if len(body) == 0 { + return "" + } + + var b strings.Builder + hasStablePrefix := false + appendJSON := func(label string, value gjson.Result) { + normalized, ok := normalizeNonEmptyCompatSeedJSON(value) + if !ok { + return + } + _, _ = b.WriteString("|") + _, _ = b.WriteString(label) + _, _ = b.WriteString("=") + _, _ = b.WriteString(normalized) + hasStablePrefix = true + } + + if tools := gjson.GetBytes(body, "tools"); tools.Exists() && tools.IsArray() { + appendJSON("tools", tools) + } + if funcs := gjson.GetBytes(body, "functions"); funcs.Exists() && funcs.IsArray() { + appendJSON("functions", funcs) + } + if instructions := gjson.GetBytes(body, "instructions"); strings.TrimSpace(instructions.String()) != "" { + appendJSON("instructions", instructions) + } + + appendSystemMessages := func(items gjson.Result) { + items.ForEach(func(_, item gjson.Result) bool { + role := strings.TrimSpace(item.Get("role").String()) + switch role { + case "system", "developer": + appendJSON(role, item.Get("content")) + } + return true + }) + } + + if messages := gjson.GetBytes(body, "messages"); messages.Exists() && messages.IsArray() { + appendSystemMessages(messages) + } else if input := gjson.GetBytes(body, "input"); input.Exists() && input.IsArray() { + appendSystemMessages(input) + } + + if !hasStablePrefix { + return "" + } + return contentStablePrefixSessionSeedPrefix + b.String() +} + +func normalizeNonEmptyCompatSeedJSON(value gjson.Result) (string, bool) { + if !value.Exists() || value.Type == gjson.Null { + return "", false + } + normalized := normalizeCompatSeedJSON(json.RawMessage(value.Raw)) + switch normalized { + case "", `""`, "[]", "{}", "null": + return "", false + default: + return normalized, true + } +} diff --git a/backend/internal/service/openai_content_session_seed_test.go b/backend/internal/service/openai_content_session_seed_test.go index 65a0bf1808..6dadc5cf53 100644 --- a/backend/internal/service/openai_content_session_seed_test.go +++ b/backend/internal/service/openai_content_session_seed_test.go @@ -216,3 +216,154 @@ func TestDeriveOpenAIContentSessionSeed_ResponsesAPI_TypedMessageItem(t *testing require.Contains(t, seed, "|first_user=") require.Contains(t, seed, "Hello from typed message") } + +func TestDeriveOpenAIStablePrefixSessionSeed_IgnoresUserContent(t *testing.T) { + first := []byte(`{ + "model": "grok", + "instructions": "Be concise.", + "tools": [{"type":"function","name":"lookup","parameters":{"type":"object"}}], + "input": [{"role":"user","content":"Question A"}] + }`) + second := []byte(`{ + "model": "grok", + "instructions": "Be concise.", + "tools": [{"parameters":{"type":"object"},"name":"lookup","type":"function"}], + "input": [{"role":"user","content":"Question B"}] + }`) + + firstSeed := deriveOpenAIStablePrefixSessionSeed(first) + secondSeed := deriveOpenAIStablePrefixSessionSeed(second) + + require.NotEmpty(t, firstSeed) + require.Equal(t, firstSeed, secondSeed) + require.NotContains(t, firstSeed, "Question A") + require.NotContains(t, firstSeed, "first_user") +} + +func TestDeriveOpenAIStablePrefixSessionSeed_IsolatesStablePrefixFields(t *testing.T) { + base := []byte(`{ + "instructions":"Be concise.", + "tools":[{"type":"function","name":"lookup"}], + "input":[{"role":"system","content":"System A"},{"role":"user","content":"Question"}] + }`) + differentInstructions := []byte(`{ + "instructions":"Be detailed.", + "tools":[{"type":"function","name":"lookup"}], + "input":[{"role":"system","content":"System A"},{"role":"user","content":"Question"}] + }`) + differentTools := []byte(`{ + "instructions":"Be concise.", + "tools":[{"type":"function","name":"search"}], + "input":[{"role":"system","content":"System A"},{"role":"user","content":"Question"}] + }`) + differentSystem := []byte(`{ + "instructions":"Be concise.", + "tools":[{"type":"function","name":"lookup"}], + "input":[{"role":"system","content":"System B"},{"role":"user","content":"Question"}] + }`) + + baseSeed := deriveOpenAIStablePrefixSessionSeed(base) + require.NotEqual(t, baseSeed, deriveOpenAIStablePrefixSessionSeed(differentInstructions)) + require.NotEqual(t, baseSeed, deriveOpenAIStablePrefixSessionSeed(differentTools)) + require.NotEqual(t, baseSeed, deriveOpenAIStablePrefixSessionSeed(differentSystem)) +} + +func TestDeriveOpenAIStablePrefixSessionSeed_ChatSystemAndDeveloper(t *testing.T) { + first := []byte(`{ + "messages":[ + {"role":"system","content":"System prompt"}, + {"role":"developer","content":[{"type":"text","text":"Developer prompt"}]}, + {"role":"user","content":"Question A"} + ] + }`) + second := []byte(`{ + "messages":[ + {"role":"system","content":"System prompt"}, + {"role":"developer","content":[{"text":"Developer prompt","type":"text"}]}, + {"role":"user","content":"Question B"} + ] + }`) + + firstSeed := deriveOpenAIStablePrefixSessionSeed(first) + require.Equal(t, firstSeed, deriveOpenAIStablePrefixSessionSeed(second)) + require.Contains(t, firstSeed, "System prompt") + require.Contains(t, firstSeed, "Developer prompt") +} + +func TestDeriveOpenAIStablePrefixSessionSeed_EncodesSystemAndDeveloperRoles(t *testing.T) { + systemThenDeveloper := []byte(`{ + "messages":[ + {"role":"system","content":"Prompt A"}, + {"role":"developer","content":"Prompt B"} + ] + }`) + developerThenSystem := []byte(`{ + "messages":[ + {"role":"developer","content":"Prompt A"}, + {"role":"system","content":"Prompt B"} + ] + }`) + + firstSeed := deriveOpenAIStablePrefixSessionSeed(systemThenDeveloper) + secondSeed := deriveOpenAIStablePrefixSessionSeed(developerThenSystem) + + require.NotEqual(t, firstSeed, secondSeed) + require.Contains(t, firstSeed, "|system=") + require.Contains(t, firstSeed, "|developer=") +} + +func TestDeriveOpenAIStablePrefixSessionSeed_EncodesInstructionDelimiters(t *testing.T) { + instructionOnly := []byte(`{ + "instructions":"foo|system=\"bar\"" + }`) + instructionAndSystem := []byte(`{ + "instructions":"foo", + "input":[{"role":"system","content":"bar"}] + }`) + + firstSeed := deriveOpenAIStablePrefixSessionSeed(instructionOnly) + secondSeed := deriveOpenAIStablePrefixSessionSeed(instructionAndSystem) + + require.NotEmpty(t, firstSeed) + require.NotEmpty(t, secondSeed) + require.NotEqual(t, firstSeed, secondSeed) +} + +func TestDeriveOpenAIAnchoredContentSessionSeed_RequiresMeaningfulAnchor(t *testing.T) { + emptyAnchors := [][]byte{ + nil, + []byte(`{"model":"grok"}`), + []byte(`{"model":"grok","messages":[{"role":"assistant","content":"answer"}]}`), + []byte(`{"model":"grok","messages":[{"role":"user","content":" "}]}`), + []byte(`{"model":"grok","messages":[{"role":"user","content":[{"type":"text","text":""}]}]}`), + []byte(`{"model":"grok","input":" "}`), + []byte(`{"model":"grok","input":[{"type":"input_text","text":""}]}`), + } + for _, body := range emptyAnchors { + require.Empty(t, deriveOpenAIAnchoredContentSessionSeed(body)) + } + + meaningfulAnchors := [][]byte{ + []byte(`{"model":"grok","messages":[{"role":"user","content":"question"}]}`), + []byte(`{"model":"grok","messages":[{"role":"user","content":[{"type":"text","text":"question"}]}]}`), + []byte(`{"model":"grok","input":"question"}`), + []byte(`{"model":"grok","input":[{"type":"input_text","text":"question"}]}`), + } + for _, body := range meaningfulAnchors { + require.NotEmpty(t, deriveOpenAIAnchoredContentSessionSeed(body)) + } +} + +func TestDeriveOpenAIStablePrefixSessionSeed_RequiresMeaningfulPrefix(t *testing.T) { + tests := [][]byte{ + nil, + []byte(`{}`), + []byte(`{"model":"grok","input":"Question A"}`), + []byte(`{"model":"grok","tools":[],"input":"Question A"}`), + []byte(`{"model":"grok","functions":[],"instructions":" ","messages":[{"role":"system","content":""},{"role":"user","content":"Question A"}]}`), + } + + for _, body := range tests { + require.Empty(t, deriveOpenAIStablePrefixSessionSeed(body)) + } +} diff --git a/backend/internal/service/openai_gateway_chat_completions_test.go b/backend/internal/service/openai_gateway_chat_completions_test.go index b85ee33947..5186598a70 100644 --- a/backend/internal/service/openai_gateway_chat_completions_test.go +++ b/backend/internal/service/openai_gateway_chat_completions_test.go @@ -98,7 +98,7 @@ func TestNormalizeResponsesBodyServiceTier(t *testing.T) { require.False(t, gjson.GetBytes(body, "service_tier").Exists()) } -func TestForwardAsChatCompletions_UnknownModelDoesNotUseDefaultMappedModel(t *testing.T) { +func TestForwardAsChatCompletions_UnknownModelWithoutMessagesDispatchKeepsRequestedModel(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() @@ -129,7 +129,7 @@ func TestForwardAsChatCompletions_UnknownModelDoesNotUseDefaultMappedModel(t *te }, } - result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.4") + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") require.Error(t, err) require.Nil(t, result) require.Equal(t, "gpt6", gjson.GetBytes(upstream.lastBody, "model").String()) diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index f8dfdd2247..980d2ffe33 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -183,7 +183,11 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco if apiKey != nil { imageGenerationAllowed = GroupAllowsImageGeneration(apiKey.Group) } - codexImageGenerationBridgeEnabled := isCodexCLI && imageGenerationAllowed && codexImageGenerationExplicitToolPolicy != codexImageGenerationExplicitToolPolicyStrip && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) + codexImageGenerationBridgeEnabled := isCodexCLI && + !isOpenAIResponsesLiteHeader(c.GetHeader(responsesLiteHeader)) && + imageGenerationAllowed && + codexImageGenerationExplicitToolPolicy != codexImageGenerationExplicitToolPolicyStrip && + s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) var imageIntent bool if isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip { decoded, decodeErr := ensureReqBody() diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 379b586136..8bd1d456dc 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -23,6 +23,7 @@ const ( grokComposerImageBridgeMaxOutputTokens = 512 grokUpstreamUserAgent = "sub2api-grok/1.0" grokCLIVersion = "0.2.93" + grokDefaultResponsesModel = "grok-4.5" grokRateLimitFallbackCooldown = 2 * time.Minute ) @@ -41,7 +42,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( upstreamModel := account.GetMappedModel(originalModel) if strings.TrimSpace(upstreamModel) == "" { - upstreamModel = "grok-4.3" + upstreamModel = grokDefaultResponsesModel } cacheIdentity := resolveGrokCacheIdentity(c, body, "", upstreamModel) patchedBody, err := patchGrokResponsesBody(body, upstreamModel) @@ -184,6 +185,10 @@ func patchGrokResponsesBody(body []byte, upstreamModel string) ([]byte, error) { if err != nil { return nil, err } + out, err = sanitizeGrokReasoningNullContent(out) + if err != nil { + return nil, err + } out, err = sanitizeGrokResponsesTools(out) if err != nil { return nil, err @@ -303,6 +308,35 @@ func sanitizeGrokResponsesInput(body []byte) ([]byte, error) { return sjson.SetRawBytes(body, "input", encoded) } +// sanitizeGrokReasoningNullContent 删除 reasoning 项中的 "content": null。 +// xAI 的 untagged enum 反序列化器拒收该字段，返回 422。 +func sanitizeGrokReasoningNullContent(body []byte) ([]byte, error) { + input := gjson.GetBytes(body, "input") + if !input.Exists() || !input.IsArray() { + return body, nil + } + + items := input.Array() + changed := false + for i := len(items) - 1; i >= 0; i-- { + item := items[i] + if strings.TrimSpace(item.Get("type").String()) != "reasoning" { + continue + } + contentResult := item.Get("content") + if contentResult.Exists() && contentResult.Type == gjson.Null { + var err error + body, err = sjson.DeleteBytes(body, fmt.Sprintf("input.%d.content", i)) + if err != nil { + return nil, err + } + changed = true + } + } + _ = changed + return body, nil +} + var grokResponsesSupportedToolTypes = map[string]struct{}{ "code_execution": {}, "code_interpreter": {}, diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index 20934b94c3..1d689bce8a 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -42,7 +42,13 @@ func resolveGrokCacheIdentity(c *gin.Context, body []byte, explicitKey, upstream seed := explicitGrokCacheSeed(c, body, explicitKey) if seed == "" { - seed = deriveOpenAIContentSessionSeed(body) + seed = deriveOpenAIStablePrefixSessionSeed(body) + if seed == "" { + // A model alone is too broad for cache routing. Preserve the + // existing first-user-derived identity when no reusable prefix is + // available so unrelated prompts do not share one tenant-wide key. + seed = deriveOpenAIAnchoredContentSessionSeed(body) + } } if seed == "" { return "" diff --git a/backend/internal/service/openai_gateway_grok_cache_test.go b/backend/internal/service/openai_gateway_grok_cache_test.go index 556f19304f..42abfc5800 100644 --- a/backend/internal/service/openai_gateway_grok_cache_test.go +++ b/backend/internal/service/openai_gateway_grok_cache_test.go @@ -37,6 +37,63 @@ func TestResolveGrokCacheIdentityStableAcrossAppendOnlyTurns(t *testing.T) { require.Equal(t, first, second) } +func TestResolveGrokCacheIdentityStableAcrossIndependentPromptsWithSamePrefix(t *testing.T) { + gin.SetMode(gin.TestMode) + c := newGrokCacheTestContext(102) + firstBody := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"user","content":"Question A"}]}`) + secondBody := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"user","content":"Question B"}]}`) + + first := resolveGrokCacheIdentity(c, firstBody, "", "grok-4.5") + second := resolveGrokCacheIdentity(c, secondBody, "", "grok-4.5") + + require.NotEmpty(t, first) + require.Equal(t, first, second) +} + +func TestResolveGrokCacheIdentityStablePrefixIsolation(t *testing.T) { + gin.SetMode(gin.TestMode) + baseBody := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"system","content":"System A"},{"role":"user","content":"Question A"}]}`) + differentInstructions := []byte(`{"model":"grok","instructions":"be detailed","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"system","content":"System A"},{"role":"user","content":"Question B"}]}`) + differentSystem := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"system","content":"System B"},{"role":"user","content":"Question B"}]}`) + differentTools := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"search"}],"input":[{"role":"system","content":"System A"},{"role":"user","content":"Question B"}]}`) + + base := resolveGrokCacheIdentity(newGrokCacheTestContext(103), baseBody, "", "grok-4.5") + require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(104), baseBody, "", "grok-4.5")) + require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(103), baseBody, "", "grok-4.3")) + require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(103), differentInstructions, "", "grok-4.5")) + require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(103), differentSystem, "", "grok-4.5")) + require.NotEqual(t, base, resolveGrokCacheIdentity(newGrokCacheTestContext(103), differentTools, "", "grok-4.5")) +} + +func TestResolveGrokCacheIdentityFallsBackWhenStablePrefixIsEmpty(t *testing.T) { + gin.SetMode(gin.TestMode) + c := newGrokCacheTestContext(105) + firstBody := []byte(`{"model":"grok","tools":[],"input":"Question A"}`) + secondBody := []byte(`{"model":"grok","tools":[],"input":"Question B"}`) + + first := resolveGrokCacheIdentity(c, firstBody, "", "grok-4.5") + second := resolveGrokCacheIdentity(c, secondBody, "", "grok-4.5") + + require.NotEmpty(t, first) + require.NotEmpty(t, second) + require.NotEqual(t, first, second) +} + +func TestResolveGrokCacheIdentitySkipsUnanchoredFallback(t *testing.T) { + gin.SetMode(gin.TestMode) + c := newGrokCacheTestContext(106) + tests := [][]byte{ + []byte(`{"model":"grok"}`), + []byte(`{"model":"grok","messages":[{"role":"assistant","content":"answer"}]}`), + []byte(`{"model":"grok","messages":[{"role":"user","content":""}]}`), + []byte(`{"model":"grok","input":" "}`), + } + + for _, body := range tests { + require.Empty(t, resolveGrokCacheIdentity(c, body, "", "grok-4.5")) + } +} + func TestResolveGrokCacheIdentityIsolatesAPIKeyAndMappedModel(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"grok","input":"same prompt"}`) diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 71e427c790..3b64cbc16b 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -271,7 +271,22 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) require.Equal(t, `{"model":"grok-4.3"}`, strings.TrimSpace(string(data))) } -func TestBuildGrokResponsesRequestRejectsUnsafeAccountBaseURL(t *testing.T) { +func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "base_url": "https://grok.example.test/v1/", + }, + } + + req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "") + require.NoError(t, err) + require.Equal(t, "https://grok.example.test/v1/responses", req.URL.String()) + require.Equal(t, "Bearer api-key", req.Header.Get("Authorization")) +} + +func TestBuildGrokResponsesRequestPinsOAuthCustomBaseURLByDefault(t *testing.T) { t.Parallel() account := &Account{ @@ -282,9 +297,9 @@ func TestBuildGrokResponsesRequestRejectsUnsafeAccountBaseURL(t *testing.T) { }, } - _, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "") - require.Error(t, err) - require.Contains(t, err.Error(), "invalid base url") + req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "") + require.NoError(t, err) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String()) } func TestGrokMediaGenerationGateCoversImagesAndVideo(t *testing.T) { @@ -596,6 +611,42 @@ func TestForwardGrokMediaVideoGenerationPreservesImageToVideoModel(t *testing.T) require.Equal(t, VideoBillingDefaultDurationSeconds, result.VideoDurationSeconds) } +func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + imageData := strings.Repeat("A", 2*1024*1024) + body := []byte(`{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"image_url":"data:image/png;base64,` + imageData + `"}}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 66, + Name: "grok-oauth", + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "oauth-access-token", + "base_url": xai.DefaultCLIBaseURL, + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + }, + Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-oauth"}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + _, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json") + require.NoError(t, err) + require.Equal(t, xai.DefaultBaseURL+"/videos/generations", upstream.lastReq.URL.String()) + require.Equal(t, "data:image/png;base64,"+imageData, gjson.GetBytes(upstream.lastBody, "image.image_url").String()) +} + func TestForwardGrokMediaVideoStatusUsesGETWithoutBody(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") gin.SetMode(gin.TestMode) @@ -637,6 +688,49 @@ func TestForwardGrokMediaVideoStatusUsesGETWithoutBody(t *testing.T) { require.Equal(t, "xai-video-req", result.RequestID) } +func TestForwardGrokMediaVideoMutationEndpoints(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + endpoint GrokMediaEndpoint + path string + }{ + {name: "edit", endpoint: GrokMediaEndpointVideosEdits, path: "/videos/edits"}, + {name: "extension", endpoint: GrokMediaEndpointVideosExtensions, path: "/videos/extensions"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok-imagine-video","prompt":"continue","video":{"url":"https://example.com/in.mp4"},"duration":6}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1"+tt.path, bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 71, Name: "grok", Platform: PlatformGrok, Type: AccountTypeAPIKey, Concurrency: 1, + Credentials: map[string]any{"api_key": "api-key", "base_url": "https://xai.test/v1"}, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"request_id":"video-mutation-123"}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.ForwardGrokMedia(context.Background(), c, account, tt.endpoint, "", body, "application/json") + require.NoError(t, err) + require.Equal(t, "https://xai.test/v1"+tt.path, upstream.lastReq.URL.String()) + require.Equal(t, http.MethodPost, upstream.lastReq.Method) + require.JSONEq(t, string(body), string(upstream.lastBody)) + require.Equal(t, "video-mutation-123", result.ResponseID) + require.Equal(t, 1, result.VideoCount) + require.Equal(t, 6, result.VideoDurationSeconds) + }) + } +} + func TestBindGrokMediaVideoRequestAccountUsesRequestIDStickyHash(t *testing.T) { ctx := context.Background() groupID := int64(7) @@ -759,12 +853,12 @@ func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *tes require.Equal(t, http.StatusOK, recorder.Code) } -func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T) { +func TestForwardGrokResponsesStreamingDefaultsEmptyModelTo45AndSnapshots(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) - body := []byte(`{"model":"grok","input":"hi","stream":true,"reasoning_effort":"high"}`) + body := []byte(`{"input":"hi","stream":true,"reasoning_effort":"high"}`) c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) c.Request.Header.Set("Content-Type", "application/json") c.Request.Header.Set("OpenAI-Beta", "responses=experimental") @@ -811,7 +905,7 @@ func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T) accountRepo: repo, } - result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", true, time.Now()) + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "", true, time.Now()) require.NoError(t, err) require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) @@ -1677,3 +1771,70 @@ func TestFailoverOpenAIUpstreamHTTPErrorUsesOnlyGrokRateLimitPolicy(t *testing.T require.Equal(t, 1, repo.rateLimitedCalls) require.Zero(t, repo.tempUnschedCalls) } + +func TestPatchGrokResponsesBody_StripsReasoningContentNull(t *testing.T) { + t.Parallel() + + body := []byte(`{ + "model": "grok-latest", + "input": [ + {"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"thinking..."}],"content":null,"encrypted_content":null}, + {"type":"message","role":"assistant","content":[{"type":"output_text","text":"Hello!"}]} + ] + }`) + + patched, err := patchGrokResponsesBody(body, "grok-4.5") + require.NoError(t, err) + require.True(t, json.Valid(patched)) + + input := gjson.GetBytes(patched, "input") + require.True(t, input.IsArray()) + + items := input.Array() + require.Len(t, items, 3) + + reasoning := items[1] + require.Equal(t, "reasoning", reasoning.Get("type").String()) + require.True(t, reasoning.Get("summary").Exists(), "summary should be preserved") + require.False(t, reasoning.Get("content").Exists(), "content: null should be stripped") +} + +func TestPatchGrokResponsesBody_KeepsReasoningContentNonNull(t *testing.T) { + t.Parallel() + + body := []byte(`{ + "model": "grok-latest", + "input": [ + {"type":"reasoning","summary":[{"type":"summary_text","text":"ok"}],"content":"real content"} + ] + }`) + + patched, err := patchGrokResponsesBody(body, "grok-4.5") + require.NoError(t, err) + + reasoning := gjson.GetBytes(patched, "input.0") + require.Equal(t, "real content", reasoning.Get("content").String(), "non-null content must not be stripped") +} + +func TestPatchGrokResponsesBody_MultipleReasoningContentNull(t *testing.T) { + t.Parallel() + + body := []byte(`{ + "model": "grok-latest", + "input": [ + {"type":"reasoning","summary":[{"type":"summary_text","text":"r1"}],"content":null}, + {"type":"message","role":"user","content":"hi"}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"r2"}],"content":null} + ] + }`) + + patched, err := patchGrokResponsesBody(body, "grok-4.5") + require.NoError(t, err) + + items := gjson.GetBytes(patched, "input").Array() + require.Len(t, items, 3) + + require.False(t, items[0].Get("content").Exists()) + require.False(t, items[2].Get("content").Exists()) +} diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index 9f62362ee8..0a19a9d283 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -914,6 +914,11 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( trimmedData = strings.TrimSpace(string(normalizedData)) line = "data: " + string(normalizedData) } + if normalizedData, normalized := normalizeCompletedImageGenerationStatus(dataBytes); normalized { + dataBytes = normalizedData + trimmedData = strings.TrimSpace(string(normalizedData)) + line = "data: " + string(normalizedData) + } if trimmedData != "[DONE]" { restoredData, restoreErr := restoreOpenAIResponsesNamespacePayload(c, dataBytes) if restoreErr != nil { diff --git a/backend/internal/service/openai_gateway_record_usage_test.go b/backend/internal/service/openai_gateway_record_usage_test.go index 81578f630f..d2eca05d74 100644 --- a/backend/internal/service/openai_gateway_record_usage_test.go +++ b/backend/internal/service/openai_gateway_record_usage_test.go @@ -39,6 +39,17 @@ type openAIRecordUsageBillingRepoStub struct { lastCtxErr error } +type openAIRecordUsageAccountRepoStub struct { + AccountRepository + account *Account + calls int +} + +func (s *openAIRecordUsageAccountRepoStub) GetByID(_ context.Context, _ int64) (*Account, error) { + s.calls++ + return s.account, nil +} + func (s *openAIRecordUsageBillingRepoStub) Apply(ctx context.Context, cmd *UsageBillingCommand) (*UsageBillingApplyResult, error) { s.calls++ s.lastCmd = cmd @@ -1045,7 +1056,7 @@ func TestOpenAIGatewayServiceRecordUsage_GPT56SeparatesCacheWriteForBillingAndSt require.InDelta(t, usageRepo.lastLog.TotalCost*1.1, usageRepo.lastLog.ActualCost, 1e-12) } -func TestOpenAIGatewayServiceRecordUsage_Gpt54LongContextBillsWholeSession(t *testing.T) { +func TestOpenAIGatewayServiceRecordUsage_Gpt54LongContextBillingDisabledByDefault(t *testing.T) { usageRepo := &openAIRecordUsageLogRepoStub{inserted: true} userRepo := &openAIRecordUsageUserRepoStub{} subRepo := &openAIRecordUsageSubRepoStub{} @@ -1063,7 +1074,45 @@ func TestOpenAIGatewayServiceRecordUsage_Gpt54LongContextBillsWholeSession(t *te }, APIKey: &APIKey{ID: 1014}, User: &User{ID: 2014}, - Account: &Account{ID: 3014}, + Account: &Account{ID: 3014, Platform: PlatformOpenAI}, + }) + + require.NoError(t, err) + require.NotNil(t, usageRepo.lastLog) + + expectedInput := 300000 * 2.5e-6 + expectedOutput := 2000 * 15e-6 + require.InDelta(t, expectedInput, usageRepo.lastLog.InputCost, 1e-10) + require.InDelta(t, expectedOutput, usageRepo.lastLog.OutputCost, 1e-10) + require.InDelta(t, expectedInput+expectedOutput, usageRepo.lastLog.TotalCost, 1e-10) + require.InDelta(t, (expectedInput+expectedOutput)*1.1, usageRepo.lastLog.ActualCost, 1e-10) + require.False(t, usageRepo.lastLog.LongContextBillingApplied) + require.Equal(t, 1, userRepo.deductCalls) +} + +func TestOpenAIGatewayServiceRecordUsage_Gpt54LongContextBillingEnabledPerAccount(t *testing.T) { + usageRepo := &openAIRecordUsageLogRepoStub{inserted: true} + userRepo := &openAIRecordUsageUserRepoStub{} + subRepo := &openAIRecordUsageSubRepoStub{} + svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, subRepo, nil) + + err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{ + Result: &OpenAIForwardResult{ + RequestID: "resp_gpt54_long_context_disabled", + Usage: OpenAIUsage{ + InputTokens: 300000, + OutputTokens: 2000, + }, + Model: "gpt-5.4-2026-03-05", + Duration: time.Second, + }, + APIKey: &APIKey{ID: 1015}, + User: &User{ID: 2015}, + Account: &Account{ + ID: 3015, + Platform: PlatformOpenAI, + Extra: map[string]any{"openai_long_context_billing_enabled": true}, + }, }) require.NoError(t, err) @@ -1075,7 +1124,62 @@ func TestOpenAIGatewayServiceRecordUsage_Gpt54LongContextBillsWholeSession(t *te require.InDelta(t, expectedOutput, usageRepo.lastLog.OutputCost, 1e-10) require.InDelta(t, expectedInput+expectedOutput, usageRepo.lastLog.TotalCost, 1e-10) require.InDelta(t, (expectedInput+expectedOutput)*1.1, usageRepo.lastLog.ActualCost, 1e-10) - require.Equal(t, 1, userRepo.deductCalls) + require.True(t, usageRepo.lastLog.LongContextBillingApplied) +} + +func TestOpenAIGatewayServiceRecordUsage_SparkShadowUsesCurrentParentBillingSetting(t *testing.T) { + tests := []struct { + name string + parentEnabled bool + }{ + {name: "parent opt out overrides stale enabled shadow", parentEnabled: false}, + {name: "parent opt in overrides stale disabled shadow", parentEnabled: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + usageRepo := &openAIRecordUsageLogRepoStub{inserted: true} + accountRepo := &openAIRecordUsageAccountRepoStub{account: &Account{ + ID: 4016, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Extra: map[string]any{openAILongContextBillingEnabledKey: tt.parentEnabled}, + }} + svc := newOpenAIRecordUsageServiceForTest( + usageRepo, + &openAIRecordUsageUserRepoStub{}, + &openAIRecordUsageSubRepoStub{}, + nil, + ) + svc.accountRepo = accountRepo + parentID := int64(4016) + + err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{ + Result: &OpenAIForwardResult{ + RequestID: "resp_gpt54_shadow_parent_setting", + Usage: OpenAIUsage{InputTokens: 300000, OutputTokens: 2000}, + Model: "gpt-5.4-2026-03-05", + Duration: time.Second, + }, + APIKey: &APIKey{ID: 1016}, + User: &User{ID: 2016}, + Account: &Account{ + ID: 3016, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + ParentAccountID: &parentID, + QuotaDimension: QuotaDimensionSpark, + Extra: map[string]any{ + openAILongContextBillingEnabledKey: !tt.parentEnabled, + }, + }, + }) + + require.NoError(t, err) + require.Equal(t, 1, accountRepo.calls) + require.Equal(t, tt.parentEnabled, usageRepo.lastLog.LongContextBillingApplied) + }) + } } func TestOpenAIGatewayServiceRecordUsage_ServiceTierPriorityUsesFastPricing(t *testing.T) { diff --git a/backend/internal/service/openai_gateway_request_body.go b/backend/internal/service/openai_gateway_request_body.go index 935a32f58b..b47c2c8638 100644 --- a/backend/internal/service/openai_gateway_request_body.go +++ b/backend/internal/service/openai_gateway_request_body.go @@ -136,12 +136,20 @@ func sanitizeEncryptedReasoningInputItem(item any) (next any, changed bool, keep return item, false, true } - _, hasEncryptedContent := inputItem["encrypted_content"] - if !hasEncryptedContent { - return item, false, true + if _, has := inputItem["encrypted_content"]; has { + delete(inputItem, "encrypted_content") + changed = true } - delete(inputItem, "encrypted_content") + // xAI 422: "content": null 导致 untagged enum 反序列化失败 + if v, has := inputItem["content"]; has && v == nil { + delete(inputItem, "content") + changed = true + } + + if !changed { + return item, false, true + } if len(inputItem) == 1 { return nil, true, false } @@ -365,15 +373,57 @@ func newOpenAIRequestView(body []byte) openAIRequestView { if len(body) == 0 { return openAIRequestView{} } - return openAIRequestView{ - body: body, - Model: strings.TrimSpace(gjson.GetBytes(body, "model").String()), - Stream: gjson.GetBytes(body, "stream").Bool(), - PromptCacheKey: strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()), - PreviousResponseID: strings.TrimSpace(gjson.GetBytes(body, "previous_response_id").String()), - ServiceTier: strings.TrimSpace(gjson.GetBytes(body, "service_tier").String()), - ReasoningEffort: strings.TrimSpace(gjson.GetBytes(body, "reasoning.effort").String()), - } + + const ( + modelField uint8 = 1 << iota + streamField + promptCacheKeyField + previousResponseIDField + serviceTierField + reasoningField + allRequestViewFields = modelField | streamField | promptCacheKeyField | + previousResponseIDField | serviceTierField | reasoningField + ) + + view := openAIRequestView{body: body} + var seen uint8 + // parseRawJSONView reads body without copying; view keeps body alive for extracted strings. + parseRawJSONView(body).ForEach(func(key, value gjson.Result) bool { + switch key.Str { + case "model": + if seen&modelField == 0 { + view.Model = strings.TrimSpace(value.String()) + seen |= modelField + } + case "stream": + if seen&streamField == 0 { + view.Stream = value.Bool() + seen |= streamField + } + case "prompt_cache_key": + if seen&promptCacheKeyField == 0 { + view.PromptCacheKey = strings.TrimSpace(value.String()) + seen |= promptCacheKeyField + } + case "previous_response_id": + if seen&previousResponseIDField == 0 { + view.PreviousResponseID = strings.TrimSpace(value.String()) + seen |= previousResponseIDField + } + case "service_tier": + if seen&serviceTierField == 0 { + view.ServiceTier = strings.TrimSpace(value.String()) + seen |= serviceTierField + } + case "reasoning": + if seen&reasoningField == 0 { + view.ReasoningEffort = strings.TrimSpace(value.Get("effort").String()) + seen |= reasoningField + } + } + return seen != allRequestViewFields + }) + return view } // Decode 保留阶段一既有 full-map 行为；后续阶段会把调用点下沉到复杂分支。 diff --git a/backend/internal/service/openai_gateway_request_body_reasoning_test.go b/backend/internal/service/openai_gateway_request_body_reasoning_test.go new file mode 100644 index 0000000000..62a8855114 --- /dev/null +++ b/backend/internal/service/openai_gateway_request_body_reasoning_test.go @@ -0,0 +1,112 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTrimOpenAIEncryptedReasoningItems_ContentNull(t *testing.T) { + reqBody := map[string]any{ + "model": "grok-4.5", + "input": []any{ + map[string]any{"type": "message", "role": "user", "content": "hi"}, + map[string]any{ + "type": "reasoning", + "summary": []any{map[string]any{"type": "summary_text", "text": "thinking..."}}, + "content": nil, + "encrypted_content": nil, + }, + map[string]any{"type": "message", "role": "assistant", "content": "Hello!"}, + }, + } + + changed := trimOpenAIEncryptedReasoningItems(reqBody) + require.True(t, changed) + + input, ok := reqBody["input"].([]any) + require.True(t, ok) + require.Len(t, input, 3) + + reasoning, ok := input[1].(map[string]any) + require.True(t, ok) + assert.Equal(t, "reasoning", reasoning["type"]) + assert.NotNil(t, reasoning["summary"]) + _, hasContent := reasoning["content"] + assert.False(t, hasContent, "content: null should be stripped") + _, hasEncrypted := reasoning["encrypted_content"] + assert.False(t, hasEncrypted, "encrypted_content should be stripped") +} + +func TestTrimOpenAIEncryptedReasoningItems_ContentNullOnly(t *testing.T) { + reqBody := map[string]any{ + "model": "grok-4.5", + "input": []any{ + map[string]any{ + "type": "reasoning", + "summary": []any{map[string]any{"type": "summary_text", "text": "ok"}}, + "content": nil, + }, + }, + } + + changed := trimOpenAIEncryptedReasoningItems(reqBody) + require.True(t, changed) + + input, ok := reqBody["input"].([]any) + require.True(t, ok) + require.Len(t, input, 1) + + reasoning, ok := input[0].(map[string]any) + require.True(t, ok) + _, hasContent := reasoning["content"] + assert.False(t, hasContent, "content: null should be stripped even without encrypted_content") +} + +func TestTrimOpenAIEncryptedReasoningItems_ContentNonNull(t *testing.T) { + reqBody := map[string]any{ + "model": "grok-4.5", + "input": []any{ + map[string]any{ + "type": "reasoning", + "summary": []any{map[string]any{"type": "summary_text", "text": "ok"}}, + "content": "some actual content", + }, + }, + } + + changed := trimOpenAIEncryptedReasoningItems(reqBody) + assert.False(t, changed, "non-null content should not be stripped") + + input, ok := reqBody["input"].([]any) + require.True(t, ok) + reasoning, ok := input[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, "some actual content", reasoning["content"]) +} + +func TestTrimOpenAIEncryptedReasoningItems_NoReasoningItems(t *testing.T) { + reqBody := map[string]any{ + "model": "grok-4.5", + "input": []any{ + map[string]any{"type": "message", "role": "user", "content": "hi"}, + }, + } + + changed := trimOpenAIEncryptedReasoningItems(reqBody) + assert.False(t, changed) +} + +func TestTrimOpenAIEncryptedReasoningItems_ContentNullDropsBareSkeleton(t *testing.T) { + reqBody := map[string]any{ + "input": []any{ + map[string]any{"type": "reasoning", "content": nil}, + }, + } + + changed := trimOpenAIEncryptedReasoningItems(reqBody) + require.True(t, changed) + _, hasInput := reqBody["input"] + assert.False(t, hasInput, "bare reasoning skeleton should be dropped, emptying input") +} diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index 9efe649653..1d70245eb0 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -281,6 +281,11 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp forceFlushFailedEvent = true sawFailedEvent = true } + if normalizedData, normalized := normalizeCompletedImageGenerationStatus(dataBytes); normalized { + dataBytes = normalizedData + data = string(normalizedData) + line = "data: " + data + } imageCounter.AddSSEData(dataBytes) // Correct Codex tool calls if needed (apply_patch -> edit, etc.) @@ -1090,6 +1095,9 @@ func extractCodexFinalResponse(body string) ([]byte, bool) { if finalResponse != nil { return } + if normalized, changed := normalizeCompletedImageGenerationStatus(data); changed { + data = normalized + } eventType := gjson.GetBytes(data, "type").String() if eventType == "response.done" || eventType == "response.completed" { if response := gjson.GetBytes(data, "response"); response.Exists() && response.Type == gjson.JSON && response.Raw != "" { @@ -1103,6 +1111,59 @@ func extractCodexFinalResponse(body string) ([]byte, bool) { return nil, false } +func normalizeCompletedImageGenerationStatus(data []byte) ([]byte, bool) { + if len(data) == 0 || !gjson.ValidBytes(data) { + return data, false + } + + shouldNormalize := func(item gjson.Result) bool { + if !item.Exists() || !item.IsObject() || + strings.TrimSpace(item.Get("type").String()) != "image_generation_call" { + return false + } + switch strings.TrimSpace(item.Get("status").String()) { + case "generating", "in_progress": + return strings.TrimSpace(item.Get("result").String()) != "" + default: + return false + } + } + + eventType := strings.TrimSpace(gjson.GetBytes(data, "type").String()) + switch eventType { + case "response.output_item.done": + if !shouldNormalize(gjson.GetBytes(data, "item")) { + return data, false + } + updated, err := sjson.SetBytes(data, "item.status", "completed") + if err != nil { + return data, false + } + return updated, true + case "response.completed", "response.done": + output := gjson.GetBytes(data, "response.output") + if !output.Exists() || !output.IsArray() { + return data, false + } + updated := data + changed := false + for i, item := range output.Array() { + if !shouldNormalize(item) { + continue + } + next, err := sjson.SetBytes(updated, "response.output."+strconv.Itoa(i)+".status", "completed") + if err != nil { + return data, false + } + updated = next + changed = true + } + return updated, changed + default: + return data, false + } +} + func normalizeResponsesStreamingTerminalOutput(data []byte, acc *apicompat.BufferedResponseAccumulator, imageOutputs []json.RawMessage) ([]byte, bool) { eventType := strings.TrimSpace(gjson.GetBytes(data, "type").String()) switch eventType { @@ -1143,7 +1204,8 @@ func responsesStreamEventMayContributeToOutput(eventType string) bool { } // collectRawResponsesOutputItemsFromSSE 按到达顺序收集 SSE 流中 -// response.output_item.done 携带的原始 item。item 以 raw JSON 逐字节保留， +// response.output_item.done 携带的原始 item。除已产生结果但仍停留在进行中 +// 的图片状态外，item 以 raw JSON 逐字节保留， // 避免经窄结构体重建时丢弃 encrypted_content/summary/opaque 等 compact // 专属或未来新增字段（#3777 问题 2）。若整条流没有任何 done 事件，退回 // 收集 output_item.added 中的 compaction 类 item——compaction 结果没有 @@ -1170,6 +1232,9 @@ func collectRawResponsesOutputItemsFromSSE(bodyText string) ([]byte, bool) { items = append(items, json.RawMessage(item.Raw)) } forEachOpenAISSEDataPayload(bodyText, func(data []byte) { + if normalized, changed := normalizeCompletedImageGenerationStatus(data); changed { + data = normalized + } if strings.TrimSpace(gjson.GetBytes(data, "type").String()) != "response.output_item.done" { return } diff --git a/backend/internal/service/openai_gateway_responses_chat_fallback.go b/backend/internal/service/openai_gateway_responses_chat_fallback.go index f494429226..1082b020dd 100644 --- a/backend/internal/service/openai_gateway_responses_chat_fallback.go +++ b/backend/internal/service/openai_gateway_responses_chat_fallback.go @@ -43,9 +43,14 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions( // custom_tool_call 项，先记下名字集合；tool_search 工具同理，回程还原为 // tool_search_call 项；namespace 子工具（如 MCP 工具）摊平转发，回程按映射还原 // 为带 namespace 字段的 function_call 项。 - customTools := apicompat.CustomToolNames(responsesReq.Tools) - toolSearch := apicompat.HasToolSearchTool(responsesReq.Tools) - namespaceTools := apicompat.NamespaceToolNames(responsesReq.Tools) + effectiveTools, err := apicompat.EffectiveResponsesTools(&responsesReq) + if err != nil { + writeOpenAIResponsesFallbackError(c, http.StatusBadRequest, "invalid_request_error", err.Error()) + return nil, fmt.Errorf("resolve responses tools: %w", err) + } + customTools := apicompat.CustomToolNames(effectiveTools) + toolSearch := apicompat.HasToolSearchTool(effectiveTools) + namespaceTools := apicompat.NamespaceToolNames(effectiveTools) chatReq, err := apicompat.ResponsesToChatCompletionsRequest(&responsesReq) if err != nil { diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 29c7d968a2..11a8255bfa 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -70,6 +70,7 @@ var openaiAllowedHeaders = map[string]bool{ "x-codex-beta-features": true, "x-codex-turn-state": true, "x-codex-turn-metadata": true, + responsesLiteHeaderKey: true, } // OpenAI passthrough allowed headers whitelist. @@ -86,6 +87,7 @@ var openaiPassthroughAllowedHeaders = map[string]bool{ "x-codex-beta-features": true, "x-codex-turn-state": true, "x-codex-turn-metadata": true, + responsesLiteHeaderKey: true, } // codex_cli_only 拒绝时记录的请求头白名单（仅用于诊断日志，不参与上游透传） @@ -404,6 +406,7 @@ type OpenAIGatewayService struct { openaiWSRetryMetrics openAIWSRetryMetrics responseHeaderFilter *responseheaders.CompiledHeaderFilter codexSnapshotThrottle *accountWriteThrottle + codexModelsManifestCache codexModelsManifestCache openaiCompatSessionResponses sync.Map openaiCompatAnthropicDigestSessions sync.Map } diff --git a/backend/internal/service/openai_gateway_service_hotpath_test.go b/backend/internal/service/openai_gateway_service_hotpath_test.go index 1dde60c9f0..326fde534d 100644 --- a/backend/internal/service/openai_gateway_service_hotpath_test.go +++ b/backend/internal/service/openai_gateway_service_hotpath_test.go @@ -27,6 +27,33 @@ func TestOpenAIRequestView_ExtractsRawScalars(t *testing.T) { require.Equal(t, "medium", view.ReasoningEffort) } +func TestOpenAIRequestView_ExtractsFieldsAfterLargeInput(t *testing.T) { + body := []byte(`{"model":"gpt-5","input":[{"content":"` + strings.Repeat("payload", 1024) + `"}],"stream":true,"prompt_cache_key":"session-1","previous_response_id":"resp-1","service_tier":"flex","reasoning":{"effort":"high"}}`) + + view := newOpenAIRequestView(body) + + require.Equal(t, "gpt-5", view.Model) + require.True(t, view.Stream) + require.Equal(t, "session-1", view.PromptCacheKey) + require.Equal(t, "resp-1", view.PreviousResponseID) + require.Equal(t, "flex", view.ServiceTier) + require.Equal(t, "high", view.ReasoningEffort) +} + +func TestOpenAIRequestView_KeepsFirstDuplicateField(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5","model":"gpt-5.1","reasoning":{"effort":"low"},"reasoning":{"effort":"high"}}`)) + + require.Equal(t, "gpt-5", view.Model) + require.Equal(t, "low", view.ReasoningEffort) +} + +func TestOpenAIRequestView_KeepsLenientPrefixExtraction(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5","stream":true,"input":[`)) + + require.Equal(t, "gpt-5", view.Model) + require.True(t, view.Stream) +} + func TestOpenAIRequestView_DecodeKeepsFullMapBehavior(t *testing.T) { view := newOpenAIRequestView([]byte(`{"model":"gpt-5","stream":true,"input":[{"type":"message","content":"hi"}]}`)) diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index bc14350394..507fb9d5c3 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -2961,7 +2961,7 @@ func TestHandleSSEToJSON_ReconstructsImageGenerationOutputItemDone(t *testing.T) Header: http.Header{"Content-Type": []string{"text/event-stream"}}, } body := []byte(strings.Join([]string{ - `data: {"type":"response.output_item.done","item":{"id":"ig_123","type":"image_generation_call","result":"aGVsbG8=","revised_prompt":"draw a cat","output_format":"png"}}`, + `data: {"type":"response.output_item.done","item":{"id":"ig_123","type":"image_generation_call","status":"generating","result":"aGVsbG8=","revised_prompt":"draw a cat","output_format":"png"}}`, `data: {"type":"response.completed","response":{"id":"resp_img","model":"gpt-5.4","output":[],"usage":{"input_tokens":7,"output_tokens":9,"output_tokens_details":{"image_tokens":4}}}}`, `data: [DONE]`, }, "\n")) @@ -2972,6 +2972,7 @@ func TestHandleSSEToJSON_ReconstructsImageGenerationOutputItemDone(t *testing.T) require.Equal(t, 4, usage.ImageOutputTokens) require.NotContains(t, rec.Body.String(), "data:") require.Equal(t, "image_generation_call", gjson.Get(rec.Body.String(), "output.0.type").String()) + require.Equal(t, "completed", gjson.Get(rec.Body.String(), "output.0.status").String()) require.Equal(t, "aGVsbG8=", gjson.Get(rec.Body.String(), "output.0.result").String()) require.Equal(t, "draw a cat", gjson.Get(rec.Body.String(), "output.0.revised_prompt").String()) } diff --git a/backend/internal/service/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index 9431b73ecb..410b4d944d 100644 --- a/backend/internal/service/openai_gateway_usage.go +++ b/backend/internal/service/openai_gateway_usage.go @@ -178,7 +178,27 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec if result.ServiceTier != nil { serviceTier = strings.TrimSpace(*result.ServiceTier) } - cost, err = s.calculateOpenAIRecordUsageCost(ctx, result, apiKey, billingModels, multiplier, imageMultiplier, videoMultiplier, baseMultiplier, tokens, serviceTier) + billingAccount := account + if account.IsShadow() { + billingAccount, err = resolveCredentialAccount(ctx, s.accountRepo, account) + if err != nil { + return err + } + } + longContextBillingEnabled := billingAccount.IsOpenAILongContextBillingEnabled() + cost, err = s.calculateOpenAIRecordUsageCost( + ctx, + result, + apiKey, + billingModels, + multiplier, + imageMultiplier, + videoMultiplier, + baseMultiplier, + tokens, + serviceTier, + longContextBillingEnabled, + ) if err != nil { if !isUsagePricingUnavailableError(err) { return err @@ -257,6 +277,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec usageLog.CacheReadCost = cost.CacheReadCost usageLog.TotalCost = cost.TotalCost usageLog.ActualCost = cost.ActualCost + usageLog.LongContextBillingApplied = cost.LongContextBillingApplied } if isVideoUsage && (cost == nil || cost.BillingMode != string(BillingModeToken)) { usageLog.RateMultiplier = videoMultiplier @@ -366,6 +387,7 @@ func (s *OpenAIGatewayService) calculateOpenAIRecordUsageCost( webSearchMultiplier float64, tokens UsageTokens, serviceTier string, + longContextBillingEnabled bool, ) (*CostBreakdown, error) { billingModel := firstUsageBillingModel(billingModels) if result != nil && result.WebSearchCalls > 0 { @@ -395,7 +417,15 @@ func (s *OpenAIGatewayService) calculateOpenAIRecordUsageCost( if candidate == "" { continue } - cost, err := s.calculateOpenAIRecordUsageTokenCost(ctx, apiKey, candidate, multiplier, tokens, serviceTier) + cost, err := s.calculateOpenAIRecordUsageTokenCost( + ctx, + apiKey, + candidate, + multiplier, + tokens, + serviceTier, + longContextBillingEnabled, + ) if err == nil { return cost, nil } @@ -443,21 +473,29 @@ func (s *OpenAIGatewayService) calculateOpenAIRecordUsageTokenCost( multiplier float64, tokens UsageTokens, serviceTier string, + longContextBillingEnabled bool, ) (*CostBreakdown, error) { if s.resolver != nil && apiKey.Group != nil { gid := apiKey.Group.ID return s.billingService.CalculateCostUnified(CostInput{ - Ctx: ctx, - Model: billingModel, - GroupID: &gid, - Tokens: tokens, - RequestCount: 1, - RateMultiplier: multiplier, - ServiceTier: serviceTier, - Resolver: s.resolver, + Ctx: ctx, + Model: billingModel, + GroupID: &gid, + Tokens: tokens, + RequestCount: 1, + RateMultiplier: multiplier, + ServiceTier: serviceTier, + Resolver: s.resolver, + LongContextBillingEnabled: &longContextBillingEnabled, }) } - return s.billingService.CalculateCostWithServiceTier(billingModel, tokens, multiplier, serviceTier) + return s.billingService.calculateCostWithServiceTierPolicy( + billingModel, + tokens, + multiplier, + serviceTier, + longContextBillingEnabled, + ) } func (s *OpenAIGatewayService) calculateOpenAIImageCost( diff --git a/backend/internal/service/openai_image_generation_controls_test.go b/backend/internal/service/openai_image_generation_controls_test.go index af0cdf669c..090948afd4 100644 --- a/backend/internal/service/openai_image_generation_controls_test.go +++ b/backend/internal/service/openai_image_generation_controls_test.go @@ -87,11 +87,13 @@ func TestOpenAIGatewayServiceForward_CodexImageInjectionRespectsGroupCapability( name string allowImages bool bridgeEnabled bool + responsesLite bool wantInjected bool }{ {name: "disabled group skips injection", allowImages: false, bridgeEnabled: true, wantInjected: false}, {name: "enabled group skips injection by default", allowImages: true, bridgeEnabled: false, wantInjected: false}, {name: "enabled group injects image tool when bridge enabled", allowImages: true, bridgeEnabled: true, wantInjected: true}, + {name: "responses lite skips hosted image bridge", allowImages: true, bridgeEnabled: true, responsesLite: true, wantInjected: false}, } for _, tt := range tests { @@ -106,6 +108,9 @@ func TestOpenAIGatewayServiceForward_CodexImageInjectionRespectsGroupCapability( svc := newOpenAIImageGenerationControlTestService(upstream) svc.cfg.Gateway.CodexImageGenerationBridgeEnabled = tt.bridgeEnabled c, _ := newOpenAIImageGenerationControlTestContext(tt.allowImages, "codex_cli_rs/0.98.0") + if tt.responsesLite { + c.Request.Header.Set(responsesLiteHeader, "true") + } account := newOpenAIImageGenerationControlTestAccount() result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.4","input":"write code","stream":false}`)) @@ -115,6 +120,11 @@ func TestOpenAIGatewayServiceForward_CodexImageInjectionRespectsGroupCapability( require.NotNil(t, upstream.lastReq) hasImageTool := gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation")`).Exists() require.Equal(t, tt.wantInjected, hasImageTool) + expectedLiteHeader := "" + if tt.responsesLite { + expectedLiteHeader = "true" + } + require.Equal(t, expectedLiteHeader, upstream.lastReq.Header.Get(responsesLiteHeader)) instructions := gjson.GetBytes(upstream.lastBody, "instructions").String() require.Equal(t, tt.wantInjected, strings.Contains(instructions, "image_generation")) toolChoice := gjson.GetBytes(upstream.lastBody, "tool_choice") @@ -126,6 +136,24 @@ func TestOpenAIGatewayServiceForward_CodexImageInjectionRespectsGroupCapability( } } +func TestOpenAIBuildUpstreamRequestOpenAIPassthroughForwardsResponsesLiteHeader(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.98.0") + c.Request.Header.Set(responsesLiteHeader, "true") + + svc := newOpenAIImageGenerationControlTestService(&httpUpstreamRecorder{}) + req, err := svc.buildUpstreamRequestOpenAIPassthrough( + c.Request.Context(), + c, + newOpenAIImageGenerationControlTestAccount(), + []byte(`{"model":"gpt-5.4","input":"write code"}`), + "test-token", + ) + + require.NoError(t, err) + require.Equal(t, "true", req.Header.Get(responsesLiteHeader)) +} + func TestOpenAIGatewayServiceForward_ExplicitImageToolWorksWithBridgeDisabled(t *testing.T) { gin.SetMode(gin.TestMode) @@ -283,6 +311,44 @@ func TestOpenAIGatewayServiceForward_ChannelBridgeOverrideEnablesCodexInjection( require.Contains(t, instructions, "image_generation") } +func TestOpenAIGatewayServiceForward_CodexBridgeDoesNotInjectHostedToolAlongsideImageGenNamespace(t *testing.T) { + gin.SetMode(gin.TestMode) + + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_namespace_image","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}`)), + }, + } + svc := newOpenAIImageGenerationControlTestService(upstream) + svc.cfg.Gateway.CodexImageGenerationBridgeEnabled = true + c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.144.1") + account := newOpenAIImageGenerationControlTestAccount() + body := []byte(`{ + "model":"gpt-5.5", + "stream":false, + "tools":[ + {"type":"function","name":"shell","parameters":{"type":"object"}}, + {"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]} + ], + "input":[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"draw a cat"}]}, + {"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]} + ], + "tool_choice":"auto" + }`) + + result, err := svc.Forward(context.Background(), c, account, body) + + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, upstream.lastReq) + require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation")`).Exists()) + require.Equal(t, "namespace", gjson.GetBytes(upstream.lastBody, `tools.#(name=="image_gen").type`).String()) + require.Equal(t, "namespace", gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools").tools.#(name=="image_gen").type`).String()) +} + func TestOpenAIGatewayServiceForward_CodexBridgePreservesExistingToolChoice(t *testing.T) { gin.SetMode(gin.TestMode) @@ -455,13 +521,13 @@ func TestOpenAIGatewayServiceHandleResponsesImageOutputs_Streaming(t *testing.T) gin.SetMode(gin.TestMode) svc := newOpenAIImageGenerationControlTestService(&httpUpstreamRecorder{}) - c, _ := newOpenAIImageGenerationControlTestContext(true, "unit-test-agent/1.0") + c, recorder := newOpenAIImageGenerationControlTestContext(true, "unit-test-agent/1.0") resp := &http.Response{ StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader( - "data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"result\":\"final-image\"}}\n\n" + - "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_stream\",\"model\":\"gpt-5.5\",\"output\":[{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"result\":\"final-image\"}],\"usage\":{\"input_tokens\":11,\"output_tokens\":5,\"output_tokens_details\":{\"image_tokens\":4}}}}\n\n", + "data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"status\":\"generating\",\"result\":\"final-image\"}}\n\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_stream\",\"model\":\"gpt-5.5\",\"output\":[{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"status\":\"generating\",\"result\":\"final-image\"}],\"usage\":{\"input_tokens\":11,\"output_tokens\":5,\"output_tokens_details\":{\"image_tokens\":4}}}}\n\n", )), } @@ -474,6 +540,73 @@ func TestOpenAIGatewayServiceHandleResponsesImageOutputs_Streaming(t *testing.T) require.Equal(t, 11, result.usage.InputTokens) require.Equal(t, 5, result.usage.OutputTokens) require.Equal(t, 4, result.usage.ImageOutputTokens) + require.NotContains(t, recorder.Body.String(), `"status":"generating"`) + require.Equal(t, 2, strings.Count(recorder.Body.String(), `"status":"completed"`)) +} + +func TestOpenAIGatewayServiceHandleResponsesImageOutputs_StreamingPassthrough(t *testing.T) { + gin.SetMode(gin.TestMode) + + svc := newOpenAIImageGenerationControlTestService(&httpUpstreamRecorder{}) + c, recorder := newOpenAIImageGenerationControlTestContext(true, "unit-test-agent/1.0") + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"status\":\"in_progress\",\"result\":\"final-image\"}}\n\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_stream\",\"model\":\"gpt-5.5\",\"output\":[{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"status\":\"in_progress\",\"result\":\"final-image\"}],\"usage\":{\"input_tokens\":11,\"output_tokens\":5,\"output_tokens_details\":{\"image_tokens\":4}}}}\n\n", + )), + } + + result, err := svc.handleStreamingResponsePassthrough(context.Background(), resp, c, &Account{ID: 1}, time.Now(), "gpt-5.5", "gpt-5.5") + + require.NoError(t, err) + require.NotNil(t, result) + require.NotContains(t, recorder.Body.String(), `"status":"in_progress"`) + require.Equal(t, 2, strings.Count(recorder.Body.String(), `"status":"completed"`)) +} + +func TestNormalizeCompletedImageGenerationStatus(t *testing.T) { + tests := []struct { + name string + input string + want string + wantChanged bool + }{ + { + name: "output item done with result", + input: `{"type":"response.output_item.done","item":{"type":"image_generation_call","status":"generating","result":"image-data"}}`, + want: `{"type":"response.output_item.done","item":{"type":"image_generation_call","status":"completed","result":"image-data"}}`, + wantChanged: true, + }, + { + name: "terminal response only changes completed image result", + input: `{"type":"response.completed","response":{"output":[{"type":"image_generation_call","status":"in_progress","result":"image-data"},{"type":"image_generation_call","status":"failed","result":"partial-data"}]}}`, + want: `{"type":"response.completed","response":{"output":[{"type":"image_generation_call","status":"completed","result":"image-data"},{"type":"image_generation_call","status":"failed","result":"partial-data"}]}}`, + wantChanged: true, + }, + { + name: "done item without result", + input: `{"type":"response.output_item.done","item":{"type":"image_generation_call","status":"generating"}}`, + want: `{"type":"response.output_item.done","item":{"type":"image_generation_call","status":"generating"}}`, + wantChanged: false, + }, + { + name: "non-final image event", + input: `{"type":"response.output_item.added","item":{"type":"image_generation_call","status":"generating","result":"image-data"}}`, + want: `{"type":"response.output_item.added","item":{"type":"image_generation_call","status":"generating","result":"image-data"}}`, + wantChanged: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, changed := normalizeCompletedImageGenerationStatus([]byte(tt.input)) + + require.Equal(t, tt.wantChanged, changed) + require.JSONEq(t, tt.want, string(got)) + }) + } } // TestHandleStreamingResponse_CyberPolicyCapturesRealUpstreamTokens 锁定流式 diff --git a/backend/internal/service/openai_images_json_keepalive.go b/backend/internal/service/openai_images_json_keepalive.go new file mode 100644 index 0000000000..0d8b24fb0a --- /dev/null +++ b/backend/internal/service/openai_images_json_keepalive.go @@ -0,0 +1,268 @@ +package service + +import ( + "bufio" + "errors" + "net" + "net/http" + "sync" + "time" + + "github.com/gin-gonic/gin" +) + +const openAIImagesJSONKeepaliveKey = "openai_images_json_keepalive" + +// openAIImagesJSONKeepalive keeps non-streaming Images API requests alive while +// an OAuth upstream is producing SSE internally. JSON permits leading +// whitespace, so each heartbeat remains compatible with clients expecting one +// final JSON document. +// +// Once the first heartbeat is sent, the HTTP status is committed as 200. Late +// upstream errors are still returned as an OpenAI-compatible JSON error body, +// matching the status tradeoff used by the compact SSE keepalive path. +type openAIImagesJSONKeepalive struct { + mu sync.Mutex + writer gin.ResponseWriter + started bool + stopped bool + bytes int + stop chan struct{} +} + +// StartOpenAIImagesJSONKeepalive starts whitespace heartbeats for a +// non-streaming Images request. A non-positive interval disables the feature. +func StartOpenAIImagesJSONKeepalive(c *gin.Context, interval time.Duration) func() { + if c == nil || c.Writer == nil || interval <= 0 { + return func() {} + } + originalWriter := c.Writer + k := &openAIImagesJSONKeepalive{ + writer: originalWriter, + stop: make(chan struct{}), + } + c.Set(openAIImagesJSONKeepaliveKey, k) + wrappedWriter := &openAIImagesJSONKeepaliveWriter{ResponseWriter: originalWriter, k: k} + c.Writer = wrappedWriter + + var reqDone <-chan struct{} + if c.Request != nil { + reqDone = c.Request.Context().Done() + } + go func() { + timer := time.NewTimer(interval) + defer timer.Stop() + for { + select { + case <-k.stop: + return + case <-reqDone: + return + case <-timer.C: + } + if !k.beat() { + return + } + timer.Reset(interval) + } + }() + + return func() { + k.Stop() + if current, ok := c.Writer.(*openAIImagesJSONKeepaliveWriter); ok && current == wrappedWriter { + c.Writer = originalWriter + } + } +} + +func (k *openAIImagesJSONKeepalive) beat() bool { + k.mu.Lock() + defer k.mu.Unlock() + if k.stopped { + return false + } + if !k.started { + header := k.writer.Header() + header.Set("Content-Type", "application/json; charset=utf-8") + header.Set("Cache-Control", "no-cache") + header.Set("X-Accel-Buffering", "no") + k.writer.WriteHeader(http.StatusOK) + k.started = true + } + n, err := k.writer.Write([]byte(" \n")) + k.bytes += n + if err != nil { + k.stopped = true + return false + } + k.writer.Flush() + return true +} + +func (k *openAIImagesJSONKeepalive) Stop() { + k.mu.Lock() + k.markStoppedLocked() + k.mu.Unlock() +} + +func (k *openAIImagesJSONKeepalive) markStoppedLocked() { + if k.stopped { + return + } + k.stopped = true + close(k.stop) +} + +// StopOpenAIImagesJSONKeepaliveCommitted stops heartbeats and reports whether +// they already committed a 200 response. +func StopOpenAIImagesJSONKeepaliveCommitted(c *gin.Context) bool { + k := openAIImagesJSONKeepaliveFromContext(c) + if k == nil { + return false + } + k.mu.Lock() + k.markStoppedLocked() + committed := k.started + k.mu.Unlock() + return committed +} + +// OpenAIImagesJSONKeepaliveAdjustedWrittenSize excludes heartbeat whitespace +// from response-size checks so account retry and failover remain available. +func OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c *gin.Context) int { + if c == nil || c.Writer == nil { + return -1 + } + k := openAIImagesJSONKeepaliveFromContext(c) + if k == nil { + return c.Writer.Size() + } + k.mu.Lock() + defer k.mu.Unlock() + size := k.writer.Size() + if size < 0 { + return size + } + if real := size - k.bytes; real > 0 { + return real + } + return -1 +} + +func openAIImagesJSONKeepaliveFromContext(c *gin.Context) *openAIImagesJSONKeepalive { + if c == nil { + return nil + } + value, ok := c.Get(openAIImagesJSONKeepaliveKey) + if !ok { + return nil + } + k, _ := value.(*openAIImagesJSONKeepalive) + return k +} + +type openAIImagesJSONKeepaliveWriter struct { + gin.ResponseWriter + k *openAIImagesJSONKeepalive +} + +func (w *openAIImagesJSONKeepaliveWriter) suspend() { + if w.k != nil { + w.k.Stop() + } +} + +func (w *openAIImagesJSONKeepaliveWriter) Header() http.Header { + w.suspend() + if w.ResponseWriter == nil { + return http.Header{} + } + return w.ResponseWriter.Header() +} + +func (w *openAIImagesJSONKeepaliveWriter) Write(data []byte) (int, error) { + w.suspend() + if w.ResponseWriter == nil { + return 0, nil + } + return w.ResponseWriter.Write(data) +} + +func (w *openAIImagesJSONKeepaliveWriter) WriteString(s string) (int, error) { + w.suspend() + if w.ResponseWriter == nil { + return 0, nil + } + return w.ResponseWriter.WriteString(s) +} + +func (w *openAIImagesJSONKeepaliveWriter) WriteHeader(code int) { + w.suspend() + if w.ResponseWriter != nil { + w.ResponseWriter.WriteHeader(code) + } +} + +func (w *openAIImagesJSONKeepaliveWriter) WriteHeaderNow() { + w.suspend() + if w.ResponseWriter != nil { + w.ResponseWriter.WriteHeaderNow() + } +} + +func (w *openAIImagesJSONKeepaliveWriter) Flush() { + w.suspend() + if w.ResponseWriter != nil { + w.ResponseWriter.Flush() + } +} + +func (w *openAIImagesJSONKeepaliveWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + if w.ResponseWriter == nil { + return nil, nil, errors.New("response writer released") + } + return w.ResponseWriter.Hijack() +} + +func (w *openAIImagesJSONKeepaliveWriter) CloseNotify() <-chan bool { + if w.ResponseWriter == nil { + ch := make(chan bool) + close(ch) + return ch + } + return w.ResponseWriter.CloseNotify() +} + +func (w *openAIImagesJSONKeepaliveWriter) Pusher() http.Pusher { + if w.ResponseWriter == nil { + return nil + } + return w.ResponseWriter.Pusher() +} + +func (w *openAIImagesJSONKeepaliveWriter) Status() int { + if w.k == nil || w.ResponseWriter == nil { + return 0 + } + w.k.mu.Lock() + defer w.k.mu.Unlock() + return w.ResponseWriter.Status() +} + +func (w *openAIImagesJSONKeepaliveWriter) Size() int { + if w.k == nil || w.ResponseWriter == nil { + return 0 + } + w.k.mu.Lock() + defer w.k.mu.Unlock() + return w.ResponseWriter.Size() +} + +func (w *openAIImagesJSONKeepaliveWriter) Written() bool { + if w.k == nil || w.ResponseWriter == nil { + return false + } + w.k.mu.Lock() + defer w.k.mu.Unlock() + return w.ResponseWriter.Written() +} diff --git a/backend/internal/service/openai_images_json_keepalive_test.go b/backend/internal/service/openai_images_json_keepalive_test.go new file mode 100644 index 0000000000..a7207b3adb --- /dev/null +++ b/backend/internal/service/openai_images_json_keepalive_test.go @@ -0,0 +1,252 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestOpenAIImagesJSONKeepalive_PreservesValidJSONResponse(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + originalWriter := c.Writer + + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + waitForOpenAIImagesJSONKeepalive(t, c) + require.Equal(t, -1, OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)) + + c.JSON(http.StatusOK, gin.H{"data": []gin.H{{"b64_json": "aW1hZ2U="}}}) + stop() + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "application/json; charset=utf-8", rec.Header().Get("Content-Type")) + require.Equal(t, "no", rec.Header().Get("X-Accel-Buffering")) + require.True(t, rec.Flushed) + require.True(t, json.Valid(rec.Body.Bytes()), rec.Body.String()) + require.Equal(t, "aW1hZ2U=", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) + require.Greater(t, OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c), 0) + require.Same(t, originalWriter, c.Writer) +} + +func TestOpenAIImagesJSONKeepalive_DisabledIsNoop(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + originalWriter := c.Writer + + stop := StartOpenAIImagesJSONKeepalive(c, 0) + stop() + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": "invalid request"}}) + + require.Same(t, originalWriter, c.Writer) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Equal(t, "invalid request", gjson.Get(rec.Body.String(), "error.message").String()) +} + +func TestOpenAIImagesJSONKeepalive_FastErrorPreservesStatus(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + stop := StartOpenAIImagesJSONKeepalive(c, time.Second) + wrote := writeOpenAIImagesUpstreamErrorResponse(c, &OpenAIImagesUpstreamError{ + StatusCode: http.StatusBadRequest, + ErrorType: "invalid_request_error", + Message: "invalid size", + }) + stop() + + require.True(t, wrote) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.False(t, strings.HasPrefix(rec.Body.String(), " \n")) + require.Equal(t, "invalid size", gjson.Get(rec.Body.String(), "error.message").String()) +} + +func TestOpenAIImagesJSONKeepalive_LateErrorRemainsJSON(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + defer stop() + waitForOpenAIImagesJSONKeepalive(t, c) + + wrote := writeOpenAIImagesUpstreamErrorResponse(c, &OpenAIImagesUpstreamError{ + StatusCode: http.StatusBadRequest, + ErrorType: "image_generation_user_error", + Code: "moderation_blocked", + Message: "request rejected", + }) + + require.True(t, wrote) + require.Equal(t, http.StatusOK, rec.Code, "heartbeat already committed the status") + require.True(t, json.Valid(rec.Body.Bytes()), rec.Body.String()) + require.Equal(t, "moderation_blocked", gjson.Get(rec.Body.String(), "error.code").String()) + require.Equal(t, "request rejected", gjson.Get(rec.Body.String(), "error.message").String()) +} + +func TestOpenAIImagesJSONKeepalive_DoesNotBlockFailoverDetection(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + waitForOpenAIImagesJSONKeepalive(t, c) + + before := OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) + require.Equal(t, -1, before) + require.True(t, c.Writer.Written()) + require.Equal(t, before, OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)) + stop() + require.True(t, strings.TrimSpace(rec.Body.String()) == "") +} + +func TestOpenAIImagesJSONKeepalive_KeepsOAuthNonStreamResponseValid(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + reader, writer := io.Pipe() + go func() { + time.Sleep(20 * time.Millisecond) + _, _ = io.WriteString(writer, + "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000000,\"output\":[{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2U=\",\"output_format\":\"png\"}]}}\n\n"+ + "data: [DONE]\n\n", + ) + _ = writer.Close() + }() + + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: reader, + } + svc := &OpenAIGatewayService{} + _, imageCount, _, err := svc.handleOpenAIImagesOAuthNonStreamingResponse(resp, c, "b64_json", "gpt-image-2") + stop() + + require.NoError(t, err) + require.Equal(t, 1, imageCount) + require.True(t, rec.Flushed) + require.True(t, strings.HasPrefix(rec.Body.String(), " \n"), rec.Body.String()) + require.True(t, json.Valid(rec.Body.Bytes()), rec.Body.String()) + require.Equal(t, "aW1hZ2U=", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) +} + +func TestOpenAIImagesJSONKeepaliveWriter_NilGuards(t *testing.T) { + w := &openAIImagesJSONKeepaliveWriter{} + require.NotPanics(t, func() { + require.NotNil(t, w.Header()) + _, _ = w.Write([]byte("test")) + _, _ = w.WriteString("test") + w.WriteHeader(http.StatusOK) + w.WriteHeaderNow() + w.Flush() + require.Equal(t, 0, w.Status()) + require.Equal(t, 0, w.Size()) + require.False(t, w.Written()) + require.Nil(t, w.Pusher()) + }) + + conn, _, err := w.Hijack() + require.Error(t, err) + require.Nil(t, conn) + select { + case <-w.CloseNotify(): + default: + t.Fatal("nil writer CloseNotify channel should be closed") + } +} + +// 回归：failover 第 2+ 轮时，上一轮心跳残留的空白字节不得被误判为"已写响应"， +// 可重试上游错误必须仍转换为 UpstreamFailoverError（而非裸错误吞掉换号）。 +func TestOpenAIImagesJSONKeepalive_HeartbeatBeforeForwardStillFailsOver(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","response_format":"b64_json"}`) + + req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = req + + svc := &OpenAIGatewayService{ + httpUpstream: &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "X-Request-Id": []string{"req_img_heartbeat_failover"}, + }, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000021}}\n\n" + + "data: {\"type\":\"error\",\"error\":{\"type\":\"server_error\",\"code\":\"server_error\",\"message\":\"The image service is temporarily unavailable.\"}}\n\n", + )), + }, + }, + } + parsed, err := svc.ParseOpenAIImagesRequest(c, body) + require.NoError(t, err) + + // 模拟上一轮 failover 已发生：心跳已提交 200 并写出空白字节。 + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + defer stop() + waitForOpenAIImagesJSONKeepalive(t, c) + + account := &Account{ + ID: 22, + Name: "openai-oauth-heartbeat-failover", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "access_token": "token-123", + }, + } + + result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode) + require.Contains(t, string(failoverErr.ResponseBody), "temporarily unavailable") + require.Empty(t, strings.TrimSpace(rec.Body.String()), "only heartbeat whitespace may reach the client") + + rawEvents, ok := c.Get(OpsUpstreamErrorsKey) + require.True(t, ok) + events, ok := rawEvents.([]*OpsUpstreamErrorEvent) + require.True(t, ok) + require.Len(t, events, 1) + require.Equal(t, "failover", events[0].Kind) + require.Equal(t, account.ID, events[0].AccountID) + require.Equal(t, http.StatusBadGateway, events[0].UpstreamStatusCode) +} + +func waitForOpenAIImagesJSONKeepalive(t *testing.T, c *gin.Context) { + t.Helper() + k := openAIImagesJSONKeepaliveFromContext(c) + require.NotNil(t, k) + require.Eventually(t, func() bool { + k.mu.Lock() + defer k.mu.Unlock() + return k.started + }, time.Second, time.Millisecond) +} diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 226cbf8e88..363ba9f591 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -1010,9 +1010,13 @@ func buildOpenAIImagesStreamErrorBodyFromUpstream(err *OpenAIImagesUpstreamError } func writeOpenAIImagesUpstreamErrorResponse(c *gin.Context, err *OpenAIImagesUpstreamError) bool { - if c == nil || c.Writer == nil || c.Writer.Written() || err == nil { + if c == nil || c.Writer == nil || err == nil { return false } + if c.Writer.Written() && OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) >= 0 { + return false + } + StopOpenAIImagesJSONKeepaliveCommitted(c) errorObj := gin.H{ "type": err.clientErrorType(), "message": err.clientMessage(), @@ -1176,7 +1180,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( var sseData openAISSEDataAccumulator var processDataErr error processDataDone := false - writerSizeBeforeResponse := c.Writer.Size() + writerSizeBeforeResponse := OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) processData := func(dataBytes []byte) { if processDataDone || processDataErr != nil { @@ -1591,7 +1595,9 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( imageOutputSizes []string firstTokenMs *int ) - writerSizeBeforeResponse := c.Writer.Size() + // 与 handleOpenAIImagesOAuthResponseError 的比较端同口径：排除非流式 JSON + // keepalive 心跳字节，避免 failover 第 2 轮起把上一轮心跳残留误判为已写响应。 + writerSizeBeforeResponse := OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) if parsed.Stream { usage, imageCount, imageOutputSizes, firstTokenMs, err = s.handleOpenAIImagesOAuthStreamingResponse(resp, c, startTime, parsed.ResponseFormat, openAIImagesStreamPrefix(parsed), requestModel) if err != nil { @@ -1672,7 +1678,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthResponseError( } retryable := IsOpenAIImagesRetryableUpstreamError(upstreamErr) - responseWritten := c != nil && c.Writer != nil && c.Writer.Size() != writerSizeBeforeResponse + responseWritten := c != nil && c.Writer != nil && OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeResponse kind := "http_error" if retryable { kind = "failover" diff --git a/backend/internal/service/openai_messages_dispatch_test.go b/backend/internal/service/openai_messages_dispatch_test.go index bafd36449b..db7804a4f3 100644 --- a/backend/internal/service/openai_messages_dispatch_test.go +++ b/backend/internal/service/openai_messages_dispatch_test.go @@ -37,3 +37,25 @@ func TestGroupResolveMessagesDispatchModel_GrokMapsClaudeFamilyToGrok(t *testing require.Empty(t, group.ResolveMessagesDispatchModel("grok")) require.Empty(t, group.ResolveMessagesDispatchModel("gpt-5.3-codex")) } + +func TestSanitizeGroupMessagesDispatchFields_ClearsNonOpenAIPlatform(t *testing.T) { + t.Parallel() + + group := &Group{ + Platform: PlatformAnthropic, + AllowMessagesDispatch: true, + DefaultMappedModel: "gpt-5.6-sol", + MessagesDispatchModelConfig: OpenAIMessagesDispatchModelConfig{ + SonnetMappedModel: "gpt-5.3-codex", + ExactModelMappings: map[string]string{ + "claude-fable-5": "gpt-5.6-sol", + }, + }, + } + + sanitizeGroupMessagesDispatchFields(group) + + require.False(t, group.AllowMessagesDispatch) + require.Empty(t, group.DefaultMappedModel) + require.Equal(t, OpenAIMessagesDispatchModelConfig{}, group.MessagesDispatchModelConfig) +} diff --git a/backend/internal/service/openai_model_mapping.go b/backend/internal/service/openai_model_mapping.go index cb7a8ca84b..8ba1d6fe1b 100644 --- a/backend/internal/service/openai_model_mapping.go +++ b/backend/internal/service/openai_model_mapping.go @@ -3,19 +3,20 @@ package service import "strings" // resolveOpenAIForwardModel 解析 OpenAI 兼容转发使用的模型。 -// defaultMappedModel 只服务于 /v1/messages 的 Claude 系列显式调度映射， -// 不作为普通 OpenAI 请求的未知模型兜底。 -func resolveOpenAIForwardModel(account *Account, requestedModel, defaultMappedModel string) string { +// messagesDispatchMappedModel 是调用方已为 /v1/messages 解析的显式调度结果； +// 普通 OpenAI 请求必须传空，避免将分组配置作为通用模型兜底。 +func resolveOpenAIForwardModel(account *Account, requestedModel, messagesDispatchMappedModel string) string { + messagesDispatchMappedModel = strings.TrimSpace(messagesDispatchMappedModel) if account == nil { - if defaultMappedModel != "" && claudeMessagesDispatchFamily(requestedModel) != "" { - return defaultMappedModel + if messagesDispatchMappedModel != "" { + return messagesDispatchMappedModel } return requestedModel } mappedModel, matched := account.ResolveMappedModel(requestedModel) - if !matched && defaultMappedModel != "" && claudeMessagesDispatchFamily(requestedModel) != "" { - return defaultMappedModel + if !matched && messagesDispatchMappedModel != "" { + return messagesDispatchMappedModel } return mappedModel } diff --git a/backend/internal/service/openai_model_mapping_test.go b/backend/internal/service/openai_model_mapping_test.go index f2ceb3551c..7107a706ad 100644 --- a/backend/internal/service/openai_model_mapping_test.go +++ b/backend/internal/service/openai_model_mapping_test.go @@ -4,159 +4,156 @@ import "testing" func TestResolveOpenAIForwardModel(t *testing.T) { tests := []struct { - name string - account *Account - requestedModel string - defaultMappedModel string - expectedModel string + name string + account *Account + requestedModel string + messagesDispatchMappedModel string + expectedModel string }{ { - name: "uses messages dispatch default for claude model", + name: "uses messages dispatch model for known claude family", account: &Account{ Credentials: map[string]any{}, }, - requestedModel: "claude-opus-4-6", - defaultMappedModel: "gpt-4o-mini", - expectedModel: "gpt-4o-mini", + requestedModel: "claude-opus-4-6", + messagesDispatchMappedModel: "gpt-4o-mini", + expectedModel: "gpt-4o-mini", }, { - name: "does not fall back to group default for invalid gpt model", + name: "uses exact messages dispatch model for unknown claude family", account: &Account{ Credentials: map[string]any{}, }, - requestedModel: "gpt6", - defaultMappedModel: "gpt-5.4", - expectedModel: "gpt6", + requestedModel: "claude-fable-5", + messagesDispatchMappedModel: " gpt-5.6-sol ", + expectedModel: "gpt-5.6-sol", }, { - name: "preserves explicit gpt-5.4 instead of group default", + name: "nil account uses messages dispatch model", + requestedModel: "claude-fable-5", + messagesDispatchMappedModel: "gpt-5.6-sol", + expectedModel: "gpt-5.6-sol", + }, + { + name: "nil account without messages dispatch keeps requested model", + requestedModel: "claude-fable-5", + expectedModel: "claude-fable-5", + }, + { + name: "ordinary unknown gpt model has no messages dispatch fallback", account: &Account{ Credentials: map[string]any{}, }, - requestedModel: "gpt-5.4", - defaultMappedModel: "gpt-4o-mini", - expectedModel: "gpt-5.4", + requestedModel: "gpt6", + expectedModel: "gpt6", }, { - name: "preserves exact passthrough mapping instead of group default", + name: "account exact mapping overrides messages dispatch model", account: &Account{ Credentials: map[string]any{ "model_mapping": map[string]any{ - "gpt-5.4": "gpt-5.4", + "claude-fable-5": "gpt-5.5", }, }, }, - requestedModel: "gpt-5.4", - defaultMappedModel: "gpt-4o-mini", - expectedModel: "gpt-5.4", + requestedModel: "claude-fable-5", + messagesDispatchMappedModel: "gpt-5.6-sol", + expectedModel: "gpt-5.5", }, { - name: "preserves wildcard passthrough mapping instead of group default", + name: "account wildcard mapping overrides messages dispatch model", account: &Account{ Credentials: map[string]any{ "model_mapping": map[string]any{ - "gpt-*": "gpt-5.4", + "claude-*": "gpt-5.4", }, }, }, - requestedModel: "gpt-5.4", - defaultMappedModel: "gpt-4o-mini", - expectedModel: "gpt-5.4", + requestedModel: "claude-fable-5", + messagesDispatchMappedModel: "gpt-5.6-sol", + expectedModel: "gpt-5.4", }, { - name: "uses account remap when explicit target differs", + name: "account passthrough mapping overrides messages dispatch model", account: &Account{ Credentials: map[string]any{ "model_mapping": map[string]any{ - "gpt-5": "gpt-5.4", + "claude-fable-5": "claude-fable-5", }, }, }, - requestedModel: "gpt-5", - defaultMappedModel: "gpt-4o-mini", - expectedModel: "gpt-5.4", + requestedModel: "claude-fable-5", + messagesDispatchMappedModel: "gpt-5.6-sol", + expectedModel: "claude-fable-5", }, { - name: "preserves codex spark instead of group default", + name: "ordinary codex spark request keeps requested model", account: &Account{ Credentials: map[string]any{}, }, - requestedModel: "gpt-5.3-codex-spark", - defaultMappedModel: "gpt-5.4", - expectedModel: "gpt-5.3-codex-spark", + requestedModel: "gpt-5.3-codex-spark", + expectedModel: "gpt-5.3-codex-spark", }, { - name: "preserves gpt-5.5 instead of group default", + name: "ordinary gpt-5.5 request keeps requested model", account: &Account{ Credentials: map[string]any{}, }, - requestedModel: "gpt-5.5", - defaultMappedModel: "gpt-5.4", - expectedModel: "gpt-5.5", + requestedModel: "gpt-5.5", + expectedModel: "gpt-5.5", }, { - name: "preserves gpt-5.5-pro instead of group default", + name: "ordinary gpt-5.5-pro request keeps requested model", account: &Account{ Credentials: map[string]any{}, }, - requestedModel: "gpt-5.5-pro", - defaultMappedModel: "gpt-5.5", - expectedModel: "gpt-5.5-pro", + requestedModel: "gpt-5.5-pro", + expectedModel: "gpt-5.5-pro", }, { - name: "preserves compact-spelled gpt5.5 instead of group default", + name: "ordinary compact-spelled gpt5.5 request keeps requested model", account: &Account{ Credentials: map[string]any{}, }, - requestedModel: "gpt5.5", - defaultMappedModel: "gpt-5.4", - expectedModel: "gpt5.5", + requestedModel: "gpt5.5", + expectedModel: "gpt5.5", }, { - name: "preserves openai namespaced gpt-5.5 instead of group default", + name: "ordinary namespaced gpt-5.5 request keeps requested model", account: &Account{ Credentials: map[string]any{}, }, - requestedModel: "openai/gpt-5.5", - defaultMappedModel: "gpt-5.4", - expectedModel: "openai/gpt-5.5", + requestedModel: "openai/gpt-5.5", + expectedModel: "openai/gpt-5.5", }, { - name: "preserves compact gpt-5.5 instead of group default", + name: "ordinary compact gpt-5.5 request keeps requested model", account: &Account{ Credentials: map[string]any{}, }, - requestedModel: "gpt-5.5-openai-compact", - defaultMappedModel: "gpt-5.4", - expectedModel: "gpt-5.5-openai-compact", + requestedModel: "gpt-5.5-openai-compact", + expectedModel: "gpt-5.5-openai-compact", + }, + { + name: "whitespace-only messages dispatch model is ignored", + account: &Account{ + Credentials: map[string]any{}, + }, + requestedModel: "gpt-5.5", + messagesDispatchMappedModel: " ", + expectedModel: "gpt-5.5", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - if got := resolveOpenAIForwardModel(tt.account, tt.requestedModel, tt.defaultMappedModel); got != tt.expectedModel { + if got := resolveOpenAIForwardModel(tt.account, tt.requestedModel, tt.messagesDispatchMappedModel); got != tt.expectedModel { t.Fatalf("resolveOpenAIForwardModel(...) = %q, want %q", got, tt.expectedModel) } }) } } -func TestResolveOpenAIForwardModel_PreventsClaudeModelFromFallingBackToGpt54(t *testing.T) { - account := &Account{ - Credentials: map[string]any{}, - } - - withoutDefault := resolveOpenAIForwardModel(account, "claude-opus-4-6", "") - if withoutDefault != "claude-opus-4-6" { - t.Fatalf("resolveOpenAIForwardModel(...) = %q, want %q", withoutDefault, "claude-opus-4-6") - } - - withDefault := resolveOpenAIForwardModel(account, "claude-opus-4-6", "gpt-5.4") - if withDefault != "gpt-5.4" { - t.Fatalf("resolveOpenAIForwardModel(...) = %q, want %q", withDefault, "gpt-5.4") - } -} - func TestResolveOpenAICompactForwardModel(t *testing.T) { tests := []struct { name string diff --git a/backend/internal/service/openai_quota_reset_credits.go b/backend/internal/service/openai_quota_reset_credits.go new file mode 100644 index 0000000000..75756976db --- /dev/null +++ b/backend/internal/service/openai_quota_reset_credits.go @@ -0,0 +1,141 @@ +package service + +import ( + "bytes" + "encoding/json" + "strconv" + "strings" +) + +type openAIRateLimitResetCreditDetailPayload struct { + ExpiresAt string `json:"expires_at,omitempty"` + ExpiresAtCamel string `json:"expiresAt,omitempty"` + ResetType string `json:"reset_type,omitempty"` + ResetTypeCamel string `json:"resetType,omitempty"` + Status string `json:"status,omitempty"` +} + +type openAIRateLimitResetCreditDetailsPayload struct { + AvailableCount json.RawMessage `json:"available_count,omitempty"` + AvailableCountCamel json.RawMessage `json:"availableCount,omitempty"` + Credits json.RawMessage `json:"credits,omitempty"` + RateLimitResetCredits json.RawMessage `json:"rate_limit_reset_credits,omitempty"` + Items json.RawMessage `json:"items,omitempty"` + Data json.RawMessage `json:"data,omitempty"` +} + +type openAIRateLimitResetCreditDetails struct { + AvailableCount *int + AvailableCreditCount int + CreditListPresent bool + Credits []OpenAIRateLimitResetCreditDetail +} + +func parseOpenAIRateLimitResetCreditDetails(body []byte) (openAIRateLimitResetCreditDetails, error) { + trimmed := bytes.TrimSpace(body) + if len(trimmed) == 0 { + return openAIRateLimitResetCreditDetails{}, nil + } + + var rawCredits []*openAIRateLimitResetCreditDetailPayload + var availableCount *int + var creditListPresent bool + if trimmed[0] == '[' { + if err := json.Unmarshal(trimmed, &rawCredits); err != nil { + return openAIRateLimitResetCreditDetails{}, err + } + creditListPresent = true + } else { + var payload openAIRateLimitResetCreditDetailsPayload + if err := json.Unmarshal(trimmed, &payload); err != nil { + return openAIRateLimitResetCreditDetails{}, err + } + availableCount = parseOpenAIResetCreditAvailableCount(payload.AvailableCount, payload.AvailableCountCamel) + var err error + rawCredits, creditListPresent, err = firstPresentResetCreditPayload( + payload.Credits, + payload.RateLimitResetCredits, + payload.Items, + payload.Data, + ) + if err != nil { + return openAIRateLimitResetCreditDetails{}, err + } + } + + credits := make([]OpenAIRateLimitResetCreditDetail, 0, len(rawCredits)) + availableCreditCount := 0 + for _, raw := range rawCredits { + if raw == nil { + continue + } + resetType := strings.TrimSpace(raw.ResetType) + if resetType == "" { + resetType = strings.TrimSpace(raw.ResetTypeCamel) + } + if resetType != "" && !strings.EqualFold(resetType, "codex_rate_limits") { + continue + } + if status := strings.TrimSpace(raw.Status); status != "" && !strings.EqualFold(status, "available") { + continue + } + availableCreditCount++ + expiresAt := strings.TrimSpace(raw.ExpiresAt) + if expiresAt == "" { + expiresAt = strings.TrimSpace(raw.ExpiresAtCamel) + } + if expiresAt == "" { + continue + } + credits = append(credits, OpenAIRateLimitResetCreditDetail{ExpiresAt: expiresAt}) + } + return openAIRateLimitResetCreditDetails{ + AvailableCount: availableCount, + AvailableCreditCount: availableCreditCount, + CreditListPresent: creditListPresent, + Credits: credits, + }, nil +} + +func parseOpenAIResetCreditAvailableCount(values ...json.RawMessage) *int { + for _, value := range values { + trimmed := bytes.TrimSpace(value) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + continue + } + + var count int + if trimmed[0] == '"' { + var text string + if err := json.Unmarshal(trimmed, &text); err != nil { + continue + } + parsed, err := strconv.Atoi(strings.TrimSpace(text)) + if err != nil { + continue + } + count = parsed + } else if err := json.Unmarshal(trimmed, &count); err != nil { + continue + } + if count >= 0 { + return &count + } + } + return nil +} + +func firstPresentResetCreditPayload(values ...json.RawMessage) ([]*openAIRateLimitResetCreditDetailPayload, bool, error) { + for _, value := range values { + trimmed := bytes.TrimSpace(value) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + continue + } + var credits []*openAIRateLimitResetCreditDetailPayload + if err := json.Unmarshal(trimmed, &credits); err != nil { + return nil, false, err + } + return credits, true, nil + } + return nil, false, nil +} diff --git a/backend/internal/service/openai_quota_reset_credits_test.go b/backend/internal/service/openai_quota_reset_credits_test.go new file mode 100644 index 0000000000..5d994602e1 --- /dev/null +++ b/backend/internal/service/openai_quota_reset_credits_test.go @@ -0,0 +1,175 @@ +package service + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestParseOpenAIRateLimitResetCreditDetails_PreservesAvailableCreditOrder(t *testing.T) { + body := []byte(`{ + "availableCount":"2", + "credits":[ + {"reset_type":"codex_rate_limits","status":"redeemed","expires_at":"2026-07-01T04:05:06Z"}, + {"reset_type":"codex_rate_limits","status":"available","expires_at":"2026-07-04T04:05:06Z"}, + {"resetType":"codex_rate_limits","status":"available","expiresAt":"2026-07-03T04:05:06Z"}, + {"reset_type":"other","status":"available","expires_at":"2026-07-02T04:05:06Z"} + ] + }`) + + details, err := parseOpenAIRateLimitResetCreditDetails(body) + require.NoError(t, err) + require.NotNil(t, details.AvailableCount) + require.Equal(t, 2, *details.AvailableCount) + require.Equal(t, []OpenAIRateLimitResetCreditDetail{ + {ExpiresAt: "2026-07-04T04:05:06Z"}, + {ExpiresAt: "2026-07-03T04:05:06Z"}, + }, details.Credits) +} + +func TestQueryUsageResetCreditCountPrecedence(t *testing.T) { + tests := []struct { + name string + usageBody string + detailBody string + wantCount int + wantCredits int + wantNil bool + }{ + { + name: "detail count creates missing usage credits", + usageBody: `{}`, + detailBody: `{"available_count":3,"credits":[{"expires_at":"2026-07-03T04:05:06Z"}]}`, + wantCount: 3, wantCredits: 1, + }, + { + name: "explicit detail zero overrides usage and records", + usageBody: `{"rate_limit_reset_credits":{"available_count":4}}`, + detailBody: `{"available_count":0,"credits":[{"expires_at":"2026-07-03T04:05:06Z"}]}`, + wantCount: 0, wantCredits: 1, + }, + { + name: "available records override usage when detail count is absent", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[{"expires_at":"2026-07-03T04:05:06Z"},{"expiresAt":"2026-07-04T04:05:06Z"}]}`, + wantCount: 2, wantCredits: 2, + }, + { + name: "empty detail list overrides usage with zero", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[]}`, + wantCount: 0, + }, + { + name: "fully filtered list overrides usage with zero", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[{"reset_type":"codex_rate_limits","status":"redeemed","expires_at":"2026-07-03T04:05:06Z"},{"reset_type":"other","status":"available","expires_at":"2026-07-04T04:05:06Z"}]}`, + wantCount: 0, + }, + { + name: "available records without expiry still count", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[{"status":"available"},{"status":"available","expires_at":"2026-07-04T04:05:06Z"}]}`, + wantCount: 2, wantCredits: 1, + }, + { + name: "shape without count or list preserves usage details", + usageBody: `{"rate_limit_reset_credits":{"available_count":5,"credits":[{"expires_at":"usage-expiry"}]}}`, + detailBody: `{}`, + wantCount: 5, + wantCredits: 1, + }, + { + name: "negative detail count without list preserves usage", + usageBody: `{"rate_limit_reset_credits":{"available_count":4}}`, + detailBody: `{"available_count":-1}`, + wantCount: 4, + }, + { + name: "negative detail count falls back to available records", + usageBody: `{"rate_limit_reset_credits":{"available_count":4}}`, + detailBody: `{"available_count":-1,"credits":[{"status":"available","expires_at":"2026-07-04T04:05:06Z"}]}`, + wantCount: 1, wantCredits: 1, + }, + { + name: "empty object preserves missing usage credits", + usageBody: `{}`, + detailBody: `{}`, + wantNil: true, + }, + { + name: "null body preserves missing usage credits", + usageBody: `{}`, + detailBody: `null`, + wantNil: true, + }, + { + name: "empty body preserves missing usage credits", + usageBody: `{}`, + detailBody: ``, + wantNil: true, + }, + { + name: "null object record is not counted", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[null]}`, + wantCount: 0, + }, + { + name: "null top level record is not counted", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `[null]`, + wantCount: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := &Account{ + ID: 100, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{ + "chatgpt_account_id": "org-parent123", + }, + } + repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: account}} + tokenCache := &stubQuotaTokenCache{tokens: map[string]string{ + OpenAITokenCacheKey(account): "fake-token", + }} + tokenProvider := NewOpenAITokenProvider(repo, tokenCache, nil) + + var detailCalls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("content-type", "application/json") + switch r.URL.Path { + case "/backend-api/wham/usage": + _, _ = w.Write([]byte(tt.usageBody)) + case "/backend-api/wham/rate-limit-reset-credits": + detailCalls++ + _, _ = w.Write([]byte(tt.detailBody)) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv)) + usage, err := svc.QueryUsage(context.Background(), 100) + require.NoError(t, err) + require.NotNil(t, usage) + require.Equal(t, 1, detailCalls) + if tt.wantNil { + require.Nil(t, usage.RateLimitResetCredits) + return + } + require.NotNil(t, usage.RateLimitResetCredits) + require.Equal(t, tt.wantCount, usage.RateLimitResetCredits.AvailableCount) + require.Len(t, usage.RateLimitResetCredits.Credits, tt.wantCredits) + }) + } +} diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index 337f8c1e8f..91b511122e 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -1,11 +1,9 @@ package service import ( - "bytes" "context" "crypto/rand" "encoding/hex" - "encoding/json" "fmt" "log/slog" "net/http" @@ -171,13 +169,26 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* } payload.FetchedAt = time.Now().Unix() - if payload.RateLimitResetCredits != nil && payload.RateLimitResetCredits.AvailableCount > 0 { - payload.RateLimitResetCredits.Credits = s.queryResetCreditDetails(callCtx, client, accessToken, chatGPTAccountID, fedRAMP, accountID) + details := s.queryResetCreditDetails(callCtx, client, accessToken, chatGPTAccountID, fedRAMP, accountID) + if details != nil { + hasDetailCount := details.AvailableCount != nil + if payload.RateLimitResetCredits == nil { + payload.RateLimitResetCredits = &OpenAIRateLimitResetCredits{} + } + if details.CreditListPresent { + payload.RateLimitResetCredits.Credits = details.Credits + } + switch { + case hasDetailCount: + payload.RateLimitResetCredits.AvailableCount = *details.AvailableCount + case details.CreditListPresent: + payload.RateLimitResetCredits.AvailableCount = details.AvailableCreditCount + } } return &payload, nil } -func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client *req.Client, accessToken, chatGPTAccountID string, fedRAMP bool, accountID int64) []OpenAIRateLimitResetCreditDetail { +func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client *req.Client, accessToken, chatGPTAccountID string, fedRAMP bool, accountID int64) *openAIRateLimitResetCreditDetails { resp, err := client.R(). SetContext(ctx). SetHeaders(buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)). @@ -191,12 +202,15 @@ func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client return nil } - credits, err := parseOpenAIRateLimitResetCreditDetails(resp.Bytes()) + details, err := parseOpenAIRateLimitResetCreditDetails(resp.Bytes()) if err != nil { slog.Warn("openai_quota_reset_credit_details_parse_failed", "account_id", accountID, "error", err) return nil } - return credits + if details.AvailableCount == nil && !details.CreditListPresent { + return nil + } + return &details } // ResetCredit consumes one rate_limit_reset_credit for the given OpenAI account. @@ -372,65 +386,6 @@ func generateRedeemRequestID() (string, error) { return fmt.Sprintf("%s-%s-%s-%s-%s", hexStr[0:8], hexStr[8:12], hexStr[12:16], hexStr[16:20], hexStr[20:]), nil } -type openAIRateLimitResetCreditDetailPayload struct { - ExpiresAt string `json:"expires_at,omitempty"` - ExpiresAtCamel string `json:"expiresAt,omitempty"` -} - -type openAIRateLimitResetCreditDetailsPayload struct { - Credits []openAIRateLimitResetCreditDetailPayload `json:"credits,omitempty"` - RateLimitResetCredits []openAIRateLimitResetCreditDetailPayload `json:"rate_limit_reset_credits,omitempty"` - Items []openAIRateLimitResetCreditDetailPayload `json:"items,omitempty"` - Data []openAIRateLimitResetCreditDetailPayload `json:"data,omitempty"` -} - -func parseOpenAIRateLimitResetCreditDetails(body []byte) ([]OpenAIRateLimitResetCreditDetail, error) { - trimmed := bytes.TrimSpace(body) - if len(trimmed) == 0 { - return nil, nil - } - - var rawCredits []openAIRateLimitResetCreditDetailPayload - if trimmed[0] == '[' { - if err := json.Unmarshal(trimmed, &rawCredits); err != nil { - return nil, err - } - } else { - var payload openAIRateLimitResetCreditDetailsPayload - if err := json.Unmarshal(trimmed, &payload); err != nil { - return nil, err - } - rawCredits = firstNonEmptyResetCreditPayload( - payload.Credits, - payload.RateLimitResetCredits, - payload.Items, - payload.Data, - ) - } - - credits := make([]OpenAIRateLimitResetCreditDetail, 0, len(rawCredits)) - for _, raw := range rawCredits { - expiresAt := strings.TrimSpace(raw.ExpiresAt) - if expiresAt == "" { - expiresAt = strings.TrimSpace(raw.ExpiresAtCamel) - } - if expiresAt == "" { - continue - } - credits = append(credits, OpenAIRateLimitResetCreditDetail{ExpiresAt: expiresAt}) - } - return credits, nil -} - -func firstNonEmptyResetCreditPayload(lists ...[]openAIRateLimitResetCreditDetailPayload) []openAIRateLimitResetCreditDetailPayload { - for _, list := range lists { - if len(list) > 0 { - return list - } - } - return nil -} - // buildCodexSparkWindowExtraUpdates extracts Codex Spark usage windows from the // /wham/usage response body's additional_rate_limits, matching the entry with // MeteredFeature == "codex_bengalfox". It produces plain codex_* keys (NOT the diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index c56600d3f5..213669f5e6 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -251,11 +251,11 @@ func TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing. t.Run(tt.name, func(t *testing.T) { got, err := parseOpenAIRateLimitResetCreditDetails([]byte(tt.body)) require.NoError(t, err) - require.Len(t, got, len(tt.want)) + require.Len(t, got.Credits, len(tt.want)) for i := range tt.want { - require.Equal(t, tt.want[i], got[i].ExpiresAt) + require.Equal(t, tt.want[i], got.Credits[i].ExpiresAt) } - encoded, err := json.Marshal(got) + encoded, err := json.Marshal(got.Credits) require.NoError(t, err) require.NotContains(t, string(encoded), "secret-id") }) diff --git a/backend/internal/service/openai_ws_client.go b/backend/internal/service/openai_ws_client.go index 80b7553083..e336abcbd6 100644 --- a/backend/internal/service/openai_ws_client.go +++ b/backend/internal/service/openai_ws_client.go @@ -39,6 +39,13 @@ type openAIWSClientConn interface { Close() error } +// openAIWSIdlePingCapable is intentionally separate from openAIWSClientConn. +// A pool probe happens while no goroutine is reading an idle connection, which +// is not safe for every WebSocket implementation. +type openAIWSIdlePingCapable interface { + SupportsIdlePingWithoutReader() bool +} + // openAIWSClientDialer 抽象 WS 建连器。 type openAIWSClientDialer interface { Dial(ctx context.Context, wsURL string, headers http.Header, proxyURL string) (openAIWSClientConn, int, http.Header, error) @@ -301,6 +308,14 @@ func (c *coderOpenAIWSClientConn) Ping(ctx context.Context) error { return c.conn.Ping(ctx) } +// SupportsIdlePingWithoutReader reports the actual coder/websocket contract. +// Conn.Ping waits for a pong, while control frames are only consumed by Read. +// The pool deliberately has no reader on an idle connection, so using Ping as +// a health probe would deterministically time out a healthy socket. +func (*coderOpenAIWSClientConn) SupportsIdlePingWithoutReader() bool { + return false +} + func (c *coderOpenAIWSClientConn) Close() error { if c == nil || c.conn == nil { return nil diff --git a/backend/internal/service/openai_ws_client_test.go b/backend/internal/service/openai_ws_client_test.go index a88d626651..95614cdbfe 100644 --- a/backend/internal/service/openai_ws_client_test.go +++ b/backend/internal/service/openai_ws_client_test.go @@ -110,3 +110,7 @@ func TestCoderOpenAIWSClientDialer_ProxyTransportTLSHandshakeTimeout(t *testing. require.NotNil(t, transport) require.Equal(t, 10*time.Second, transport.TLSHandshakeTimeout) } + +func TestCoderOpenAIWSClientConn_DoesNotSupportIdlePingWithoutReader(t *testing.T) { + require.False(t, (&coderOpenAIWSClientConn{}).SupportsIdlePingWithoutReader()) +} diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index a2af4b760f..b632c451df 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -18,6 +18,13 @@ import ( "github.com/tidwall/sjson" ) +func (s *OpenAIGatewayService) openAIWSIngressInterTurnIdleTimeout() time.Duration { + if s == nil || s.cfg == nil || s.cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds <= 0 { + return 0 + } + return time.Duration(s.cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds) * time.Second +} + func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( ctx context.Context, c *gin.Context, @@ -231,7 +238,11 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( if isCodexCLI { codexImageGenerationExplicitToolPolicy = account.CodexImageGenerationExplicitToolPolicy() } - codexBridgeEnabled := isCodexCLI && imageGenerationAllowed && codexImageGenerationExplicitToolPolicy != codexImageGenerationExplicitToolPolicyStrip && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) + codexBridgeEnabled := isCodexCLI && + !isOpenAIResponsesLiteWebSocketPayload(normalized) && + imageGenerationAllowed && + codexImageGenerationExplicitToolPolicy != codexImageGenerationExplicitToolPolicyStrip && + s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) if codexBridgeEnabled { payloadMap := make(map[string]any) if err := json.Unmarshal(normalized, &payloadMap); err != nil { @@ -360,8 +371,23 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } readClientMessage := func() ([]byte, error) { - msgType, payload, readErr := clientConn.Read(ctx) + readCtx := ctx + idleTimeout := s.openAIWSIngressInterTurnIdleTimeout() + cancelRead := func() {} + if idleTimeout > 0 { + readCtx, cancelRead = context.WithTimeout(ctx, idleTimeout) + } + msgType, payload, readErr := clientConn.Read(readCtx) + cancelRead() if readErr != nil { + if idleTimeout > 0 && errors.Is(readErr, context.DeadlineExceeded) && ctx.Err() == nil { + logOpenAIWSModeInfo("ingress_ws_inter_turn_idle_timeout account_id=%d timeout_seconds=%d", account.ID, int(idleTimeout.Seconds())) + return nil, NewOpenAIWSClientCloseError( + coderws.StatusNormalClosure, + "websocket idle timeout", + readErr, + ) + } return nil, readErr } if msgType != coderws.MessageText && msgType != coderws.MessageBinary { @@ -1318,7 +1344,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( unpinSessionConn(sessionConnID) } } - shouldPreflightPing := turn > 1 && sessionLease != nil && turnRetry == 0 + shouldPreflightPing := turn > 1 && sessionLease != nil && sessionLease.SupportsIdlePingWithoutReader() && turnRetry == 0 if shouldPreflightPing && openAIWSIngressPreflightPingIdle > 0 && !lastTurnFinishedAt.IsZero() { if time.Since(lastTurnFinishedAt) < openAIWSIngressPreflightPingIdle { shouldPreflightPing = false diff --git a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go index 9ae1b855ea..5ec4b5333b 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -164,6 +164,111 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT require.Len(t, captureConn.writes, 2, "应向同一上游连接发送两轮 response.create") } +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_IdleTimeoutReleasesStoreDisabledSession(t *testing.T) { + gin.SetMode(gin.TestMode) + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 1 + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + captureConn := &openAIWSCaptureConn{events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_idle_timeout","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }} + captureDialer := &openAIWSCaptureDialer{conn: captureConn} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(captureDialer) + defer pool.Close() + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + account := &Account{ + ID: 116, + Name: "openai-ingress-idle-timeout", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + + serverErrCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + + readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second) + _, firstMessage, err := conn.Read(readCtx) + cancelRead() + if err != nil { + serverErrCh <- err + return + } + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + ginCtx.Request = r.Clone(r.Context()) + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false}`)) + cancelWrite() + require.NoError(t, err) + + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + _, event, err := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, err) + require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) + + select { + case proxyErr := <-serverErrCh: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, proxyErr, &closeErr) + require.Equal(t, coderws.StatusNormalClosure, closeErr.StatusCode()) + require.Equal(t, "websocket idle timeout", closeErr.Reason()) + case <-time.After(4 * time.Second): + t.Fatal("timed out waiting for idle ingress session to close") + } + + ap, ok := pool.getAccountPool(account.ID) + require.True(t, ok) + ap.mu.Lock() + require.Empty(t, ap.pinnedConns, "idle close must unpin a store=false session") + for _, conn := range ap.conns { + require.False(t, conn.isLeased(), "idle close must release the upstream lease") + } + ap.mu.Unlock() +} + func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_FollowupCreateCanOmitModel(t *testing.T) { gin.SetMode(gin.TestMode) @@ -298,7 +403,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_FollowupCreateCa require.Equal(t, "resp_omit_model_1", gjson.Get(requestToJSONString(captureConn.writes[1]), "previous_response_id").String()) } -func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_InjectsCodexImageBridge(t *testing.T) { +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_CodexImageBridgeRespectsResponsesLite(t *testing.T) { gin.SetMode(gin.TestMode) cfg := &config.Config{} @@ -319,6 +424,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_InjectsCodexImag captureConn := &openAIWSCaptureConn{ events: [][]byte{ []byte(`{"type":"response.completed","response":{"id":"resp_codex_image_bridge","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}}`), + []byte(`{"type":"response.completed","response":{"id":"resp_codex_image_lite","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}}`), }, } captureDialer := &openAIWSCaptureDialer{conn: captureConn} @@ -418,6 +524,28 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_InjectsCodexImag require.Equal(t, coderws.MessageText, msgType) require.Equal(t, "resp_codex_image_bridge", gjson.GetBytes(message, "response.id").String()) + writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{ + "type":"response.create", + "model":"gpt-5.5", + "stream":false, + "previous_response_id":"resp_codex_image_bridge", + "client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"}, + "input":[ + {"type":"additional_tools","role":"developer","tools":[{"type":"custom","name":"exec","description":"Execute code-mode tools, including image_gen.imagegen."}]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"draw a cat"}]} + ] + }`)) + cancelWrite() + require.NoError(t, err) + + readCtx, cancelRead = context.WithTimeout(context.Background(), 3*time.Second) + msgType, message, err = clientConn.Read(readCtx) + cancelRead() + require.NoError(t, err) + require.Equal(t, coderws.MessageText, msgType) + require.Equal(t, "resp_codex_image_lite", gjson.GetBytes(message, "response.id").String()) + _ = clientConn.Close(coderws.StatusNormalClosure, "done") select { @@ -427,12 +555,19 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_InjectsCodexImag t.Fatal("等待 ingress websocket 结束超时") } - require.Len(t, captureConn.writes, 1) - upstreamPayload := requestToJSONString(captureConn.writes[0]) - require.True(t, gjson.Get(upstreamPayload, `tools.#(type=="image_generation")`).Exists()) - require.Equal(t, "png", gjson.Get(upstreamPayload, `tools.#(type=="image_generation").output_format`).String()) - require.Equal(t, "auto", gjson.Get(upstreamPayload, "tool_choice").String()) - require.Contains(t, gjson.Get(upstreamPayload, "instructions").String(), "image_generation") + require.Len(t, captureConn.writes, 2) + nonLitePayload := requestToJSONString(captureConn.writes[0]) + require.True(t, gjson.Get(nonLitePayload, `tools.#(type=="image_generation")`).Exists()) + require.Equal(t, "png", gjson.Get(nonLitePayload, `tools.#(type=="image_generation").output_format`).String()) + require.Equal(t, "auto", gjson.Get(nonLitePayload, "tool_choice").String()) + require.Contains(t, gjson.Get(nonLitePayload, "instructions").String(), "image_generation") + + litePayload := requestToJSONString(captureConn.writes[1]) + require.False(t, gjson.Get(litePayload, `tools.#(type=="image_generation")`).Exists()) + require.False(t, gjson.Get(litePayload, "tool_choice").Exists()) + require.NotContains(t, gjson.Get(litePayload, "instructions").String(), "image_generation") + require.Equal(t, "exec", gjson.Get(litePayload, `input.#(type=="additional_tools").tools.0.name`).String()) + require.Contains(t, gjson.Get(litePayload, `input.#(type=="additional_tools").tools.0.description`).String(), "image_gen.imagegen") } func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_DedicatedModeDoesNotReuseConnAcrossSessions(t *testing.T) { diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index 0afb9181a1..dbc24c850a 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -201,6 +201,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( if err != nil { return nil, err } + if account.Platform != PlatformGrok && isOpenAIResponsesLiteWebSocketPayload(payload) { + upstreamReq.Header.Set(responsesLiteHeader, "true") + } proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { @@ -431,7 +434,7 @@ func resolveGrokWSUpstreamModel(account *Account, body []byte, originalModel str } } if upstreamModel == "" { - upstreamModel = "grok-4.3" + upstreamModel = grokDefaultResponsesModel } return upstreamModel } diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index da2ae77917..b1bfb4d342 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -91,7 +91,7 @@ func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) { Concurrency: 1, Status: StatusActive, } - payload := []byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"input":"hi"}`) + payload := []byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"},"input":"hi"}`) type bridgeResult struct { result *OpenAIForwardResult @@ -173,11 +173,57 @@ func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) { require.NotNil(t, upstream.lastReq) require.Equal(t, http.MethodPost, upstream.lastReq.Method) + require.Equal(t, "true", upstream.lastReq.Header.Get(responsesLiteHeader)) require.False(t, gjson.GetBytes(upstream.lastBody, "type").Exists()) require.False(t, gjson.GetBytes(upstream.lastBody, "generate").Exists()) require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) } +func TestProxyOpenAIWSHTTPBridgeTurnForGrokDefaultsEmptyModelTo45(t *testing.T) { + gin.SetMode(gin.TestMode) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_grok_default","model":"grok-4.5"}}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_grok_default","model":"grok-4.5","usage":{"input_tokens":1,"output_tokens":1}}}`, + "", + }, "\n"))), + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}, + httpUpstream: upstream, + } + account := &Account{ + ID: 72, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{"base_url": xai.DefaultCLIBaseURL}, + } + payload := []byte(`{"type":"response.create","generate":true,"stream":true,"input":"hi"}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + var events [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "access-token", payload, len(payload), + "", "", "", "", "", 1, + func(message []byte) error { + events = append(events, append([]byte(nil), message...)) + return nil + }, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, grokDefaultResponsesModel, gjson.GetBytes(upstream.lastBody, "model").String()) + require.Len(t, events, 2) +} + func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T) { gin.SetMode(gin.TestMode) @@ -632,3 +678,100 @@ func TestOpenAIWSHTTPBridgeKeepsContinuationFramesOnHTTPWithoutPreviousResponseI require.Equal(t, 0, captureDialer.DialCount()) require.Empty(t, captureConn.writes) } + +func TestOpenAIWSHTTPBridge_IdleTimeoutClosesClientSession(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"type":"response.completed","response":{"id":"resp_bridge_idle","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(sseBody)), + }} + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true + cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = 1 + cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 1 + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: upstream, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 20, + Name: "api-key-bridge-idle-timeout", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "sk-upstream"}, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + Concurrency: 1, + Status: StatusActive, + Schedulable: true, + } + + errCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + errCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + + readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second) + _, firstMessage, err := conn.Read(readCtx) + cancelRead() + if err != nil { + errCh <- err + return + } + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + ginCtx.Request = r.Clone(r.Context()) + errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false,"input":"hello"}`)) + cancelWrite() + require.NoError(t, err) + + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + _, event, err := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, err) + require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) + + select { + case proxyErr := <-errCh: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, proxyErr, &closeErr) + require.Equal(t, coderws.StatusNormalClosure, closeErr.StatusCode()) + require.Equal(t, "websocket idle timeout", closeErr.Reason()) + case <-time.After(4 * time.Second): + t.Fatal("timed out waiting for idle HTTP bridge session to close") + } + require.Len(t, upstream.bodies, 1, "an idle client must not leave a continuation request running") +} diff --git a/backend/internal/service/openai_ws_pool.go b/backend/internal/service/openai_ws_pool.go index 329908e762..8affe2a930 100644 --- a/backend/internal/service/openai_ws_pool.go +++ b/backend/internal/service/openai_ws_pool.go @@ -203,6 +203,14 @@ func (l *openAIWSConnLease) PingWithTimeout(timeout time.Duration) error { return conn.pingWithTimeout(timeout) } +func (l *openAIWSConnLease) SupportsIdlePingWithoutReader() bool { + conn, err := l.activeConn() + if err != nil { + return false + } + return conn.supportsIdlePingWithoutReader() +} + func (l *openAIWSConnLease) MarkBroken() { if l == nil || l.pool == nil || l.conn == nil || l.released.Load() { return @@ -437,6 +445,16 @@ func (c *openAIWSConn) pingWithTimeout(timeout time.Duration) error { return nil } +func (c *openAIWSConn) supportsIdlePingWithoutReader() bool { + if c == nil || c.ws == nil { + return false + } + capable, ok := c.ws.(openAIWSIdlePingCapable) + // Test and alternate implementations keep the historical probe behavior + // unless they explicitly declare it unsafe. + return !ok || capable.SupportsIdlePingWithoutReader() +} + func (c *openAIWSConn) touch() { if c == nil { return @@ -707,7 +725,7 @@ func (p *openAIWSConnPool) runBackgroundPingSweep() { g.SetLimit(10) for _, item := range candidates { item := item - if item.conn == nil || item.conn.isLeased() || item.conn.waiters.Load() > 0 { + if item.conn == nil || item.conn.isLeased() || item.conn.waiters.Load() > 0 || !item.conn.supportsIdlePingWithoutReader() { continue } g.Go(func() error { @@ -1613,7 +1631,7 @@ func (p *openAIWSConnPool) nextConnID(accountID int64) string { } func (p *openAIWSConnPool) shouldHealthCheckConn(conn *openAIWSConn) bool { - if conn == nil { + if conn == nil || !conn.supportsIdlePingWithoutReader() { return false } return conn.idleDuration(time.Now()) >= openAIWSConnHealthCheckIdle @@ -1669,7 +1687,7 @@ func (p *openAIWSConnPool) effectiveMaxConnsByAccount(account *Account) int { if account.Concurrency <= 0 { return 0 } - return account.Concurrency + return min(account.Concurrency, hardCap) } if account == nil || !p.dynamicMaxConnsEnabled() { return hardCap diff --git a/backend/internal/service/openai_ws_pool_test.go b/backend/internal/service/openai_ws_pool_test.go index ae9b94ce4a..8d339359ee 100644 --- a/backend/internal/service/openai_ws_pool_test.go +++ b/backend/internal/service/openai_ws_pool_test.go @@ -739,7 +739,7 @@ func TestOpenAIWSConnPool_EffectiveMaxConnsDisabledFallbackHardCap(t *testing.T) require.Equal(t, 8, pool.effectiveMaxConnsByAccount(account), "关闭动态模式后应保持旧行为") } -func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2UsesAccountConcurrency(t *testing.T) { +func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2RespectsHardCap(t *testing.T) { cfg := &config.Config{} cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8 @@ -750,7 +750,7 @@ func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2UsesAccountConc pool := newOpenAIWSConnPool(cfg) high := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 20} - require.Equal(t, 20, pool.effectiveMaxConnsByAccount(high), "v2 路径应直接使用账号并发数作为池上限") + require.Equal(t, 8, pool.effectiveMaxConnsByAccount(high), "v2 路径也必须受连接池硬上限约束") nonPositive := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 0} require.Equal(t, 0, pool.effectiveMaxConnsByAccount(nonPositive), "并发数<=0 时应不可调度") @@ -1319,6 +1319,9 @@ func TestOpenAIWSConnPool_UtilityBranches(t *testing.T) { conn := newOpenAIWSConn("health", 1, &openAIWSFakeConn{}, nil) conn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano()) require.True(t, pool.shouldHealthCheckConn(conn)) + unsafeConn := newOpenAIWSConn("unsafe_health", 1, &openAIWSIdlePingUnsupportedConn{}, nil) + unsafeConn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano()) + require.False(t, pool.shouldHealthCheckConn(unsafeConn)) } func TestOpenAIWSConn_LeaseAndTimeHelpers_NilAndClosedBranches(t *testing.T) { @@ -1610,6 +1613,14 @@ type openAIWSPingBlockingConn struct { release <-chan struct{} } +type openAIWSIdlePingUnsupportedConn struct { + openAIWSFakeConn +} + +func (c *openAIWSIdlePingUnsupportedConn) SupportsIdlePingWithoutReader() bool { + return false +} + func (c *openAIWSPingBlockingConn) WriteJSON(context.Context, any) error { return nil } diff --git a/backend/internal/service/ops_models.go b/backend/internal/service/ops_models.go index e33dcf82a8..d95bbadbb8 100644 --- a/backend/internal/service/ops_models.go +++ b/backend/internal/service/ops_models.go @@ -8,6 +8,7 @@ import ( type OpsSystemLog struct { ID int64 `json:"id"` CreatedAt time.Time `json:"created_at"` + Host string `json:"host"` Level string `json:"level"` Component string `json:"component"` Message string `json:"message"` diff --git a/backend/internal/service/ops_port.go b/backend/internal/service/ops_port.go index 46d171c7c3..2b73d4a694 100644 --- a/backend/internal/service/ops_port.go +++ b/backend/internal/service/ops_port.go @@ -194,6 +194,7 @@ type OpsInsertSystemMetricsInput struct { type OpsInsertSystemLogInput struct { CreatedAt time.Time + Host string Level string Component string Message string @@ -210,6 +211,7 @@ type OpsInsertSystemLogInput struct { type OpsSystemLogFilter struct { StartTime *time.Time EndTime *time.Time + Host string Level string Component string @@ -230,6 +232,7 @@ type OpsSystemLogFilter struct { type OpsSystemLogCleanupFilter struct { StartTime *time.Time EndTime *time.Time + Host string Level string Component string diff --git a/backend/internal/service/ops_system_log_service.go b/backend/internal/service/ops_system_log_service.go index b3be37e8ae..b96ae89d92 100644 --- a/backend/internal/service/ops_system_log_service.go +++ b/backend/internal/service/ops_system_log_service.go @@ -89,6 +89,7 @@ func marshalSystemLogCleanupConditions(filter *OpsSystemLogCleanupFilter) string return "{}" } payload := map[string]any{ + "host": strings.TrimSpace(filter.Host), "level": strings.TrimSpace(filter.Level), "component": strings.TrimSpace(filter.Component), "request_id": strings.TrimSpace(filter.RequestID), diff --git a/backend/internal/service/ops_system_log_service_test.go b/backend/internal/service/ops_system_log_service_test.go index 8b5a84c1f0..e8c6199f17 100644 --- a/backend/internal/service/ops_system_log_service_test.go +++ b/backend/internal/service/ops_system_log_service_test.go @@ -101,6 +101,7 @@ func TestOpsServiceCleanupSystemLogs_SuccessAndAudit(t *testing.T) { now := time.Now().UTC() filter := &OpsSystemLogCleanupFilter{ StartTime: &now, + Host: "api-node-1", Level: "warn", RequestID: "req-1", ClientRequestID: "creq-1", @@ -119,6 +120,9 @@ func TestOpsServiceCleanupSystemLogs_SuccessAndAudit(t *testing.T) { if audit == nil { t.Fatalf("expected cleanup audit") } + if !strings.Contains(audit.Conditions, `"host":"api-node-1"`) { + t.Fatalf("audit conditions should include host: %s", audit.Conditions) + } if !strings.Contains(audit.Conditions, `"client_request_id":"creq-1"`) { t.Fatalf("audit conditions should include client_request_id: %s", audit.Conditions) } diff --git a/backend/internal/service/ops_system_log_sink.go b/backend/internal/service/ops_system_log_sink.go index 2ff273be53..2e6f5515c8 100644 --- a/backend/internal/service/ops_system_log_sink.go +++ b/backend/internal/service/ops_system_log_sink.go @@ -27,6 +27,7 @@ type OpsSystemLogSinkHealth struct { type OpsSystemLogSink struct { opsRepo OpsRepository + host string queue chan *logger.LogEvent @@ -45,10 +46,14 @@ type OpsSystemLogSink struct { lastError atomic.Value } +const maxSystemLogHostLength = 255 + func NewOpsSystemLogSink(opsRepo OpsRepository) *OpsSystemLogSink { ctx, cancel := context.WithCancel(context.Background()) + rawHost, err := os.Hostname() s := &OpsSystemLogSink{ opsRepo: opsRepo, + host: normalizeSystemLogHost(rawHost, err), queue: make(chan *logger.LogEvent, 5000), batchSize: 200, flushInterval: time.Second, @@ -59,6 +64,18 @@ func NewOpsSystemLogSink(opsRepo OpsRepository) *OpsSystemLogSink { return s } +func normalizeSystemLogHost(host string, err error) string { + host = strings.TrimSpace(host) + if err != nil || host == "" { + return "unknown" + } + runes := []rune(host) + if len(runes) > maxSystemLogHostLength { + return string(runes[:maxSystemLogHostLength]) + } + return host +} + func (s *OpsSystemLogSink) Start() { if s == nil || s.opsRepo == nil { return @@ -220,6 +237,7 @@ func (s *OpsSystemLogSink) flushBatch(baseCtx context.Context, batch []*logger.L inputs = append(inputs, &OpsInsertSystemLogInput{ CreatedAt: createdAt, + Host: s.host, Level: strings.ToLower(strings.TrimSpace(event.Level)), Component: component, Message: message, diff --git a/backend/internal/service/ops_system_log_sink_test.go b/backend/internal/service/ops_system_log_sink_test.go index b43d44c32e..0d15f1a662 100644 --- a/backend/internal/service/ops_system_log_sink_test.go +++ b/backend/internal/service/ops_system_log_sink_test.go @@ -140,6 +140,7 @@ func TestOpsSystemLogSink_StartStopAndFlushSuccess(t *testing.T) { } sink := NewOpsSystemLogSink(repo) + sink.host = "api-node-1" sink.batchSize = 1 sink.flushInterval = 10 * time.Millisecond sink.Start() @@ -172,6 +173,9 @@ func TestOpsSystemLogSink_StartStopAndFlushSuccess(t *testing.T) { t.Fatalf("captured len = %d, want 1", len(captured)) } item := captured[0] + if item.Host != "api-node-1" { + t.Fatalf("host = %q, want api-node-1", item.Host) + } if item.RequestID != "req-1" || item.ClientRequestID != "creq-1" { t.Fatalf("unexpected request ids: %+v", item) } @@ -324,3 +328,20 @@ func TestOpsSystemLogSink_HelperFunctions(t *testing.T) { } } } + +func TestNormalizeSystemLogHost(t *testing.T) { + if got := normalizeSystemLogHost(" api-node-1 ", nil); got != "api-node-1" { + t.Fatalf("trimmed host = %q, want api-node-1", got) + } + if got := normalizeSystemLogHost("", nil); got != "unknown" { + t.Fatalf("empty host = %q, want unknown", got) + } + if got := normalizeSystemLogHost("api-node-1", errors.New("hostname unavailable")); got != "unknown" { + t.Fatalf("errored host = %q, want unknown", got) + } + longHost := strings.Repeat("节", maxSystemLogHostLength+1) + got := normalizeSystemLogHost(longHost, nil) + if runeCount := len([]rune(got)); runeCount != maxSystemLogHostLength { + t.Fatalf("truncated host rune count = %d, want %d", runeCount, maxSystemLogHostLength) + } +} diff --git a/backend/internal/service/payment_order.go b/backend/internal/service/payment_order.go index 04feb8002a..9b7ee08990 100644 --- a/backend/internal/service/payment_order.go +++ b/backend/internal/service/payment_order.go @@ -16,6 +16,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/payment" "github.com/Wei-Shaw/sub2api/internal/payment/provider" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/shopspring/decimal" ) @@ -445,7 +446,9 @@ func (s *PaymentService) invokeProvider(ctx context.Context, order *dbent.Paymen IsMobile: req.IsMobile, ReturnURL: providerReturnURL, }, sel, outTradeNo, payAmountStr, subject) + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") pr, err := prov.CreatePayment(ctx, providerReq) + finishProviderCall() if err != nil { slog.Error("[PaymentService] CreatePayment failed", "provider", sel.ProviderKey, "instance", sel.InstanceID, "error", err) if appErr := new(infraerrors.ApplicationError); errors.As(err, &appErr) { diff --git a/backend/internal/service/payment_order_lifecycle.go b/backend/internal/service/payment_order_lifecycle.go index 8ed18797dd..46a2e00605 100644 --- a/backend/internal/service/payment_order_lifecycle.go +++ b/backend/internal/service/payment_order_lifecycle.go @@ -13,6 +13,7 @@ import ( "github.com/Wei-Shaw/sub2api/ent/paymentorder" "github.com/Wei-Shaw/sub2api/internal/payment" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" ) // --- Cancel & Expire --- @@ -157,7 +158,9 @@ func (s *PaymentService) checkPaidWithOptions(ctx context.Context, o *dbent.Paym if queryRef == "" { return "" } + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") resp, err := prov.QueryOrder(ctx, queryRef) + finishProviderCall() if err != nil { slog.Warn("query upstream failed", "orderID", o.ID, "error", err) return "" @@ -199,7 +202,9 @@ func (s *PaymentService) checkPaidWithOptions(ctx context.Context, o *dbent.Paym return "" } if cp, ok := prov.(payment.CancelableProvider); ok { + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") _ = cp.CancelPayment(ctx, queryRef) + finishProviderCall() } return "" } @@ -208,7 +213,9 @@ func requeryPaidOrderOnce(ctx context.Context, prov payment.Provider, queryRef s if prov == nil || strings.TrimSpace(queryRef) == "" { return nil, false } + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") resp, err := prov.QueryOrder(ctx, queryRef) + finishProviderCall() if err != nil { slog.Warn("query upstream retry failed", "queryRef", queryRef, "error", err) return nil, false diff --git a/backend/internal/service/payment_refund.go b/backend/internal/service/payment_refund.go index 91822680ed..bc073a2c34 100644 --- a/backend/internal/service/payment_refund.go +++ b/backend/internal/service/payment_refund.go @@ -19,6 +19,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/payment" "github.com/Wei-Shaw/sub2api/internal/payment/provider" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" ) // --- Refund Flow --- @@ -347,12 +348,14 @@ func (s *PaymentService) gwRefund(ctx context.Context, p *RefundPlan) (*payment. }) return nil, err } + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") resp, err := prov.Refund(ctx, payment.RefundRequest{ TradeNo: p.Order.PaymentTradeNo, OrderID: p.Order.OutTradeNo, Amount: formatGatewayRefundAmount(p.GatewayAmount, p.Order), Reason: p.Reason, }) + finishProviderCall() if err != nil { if resp != nil && strings.TrimSpace(resp.Status) == payment.ProviderStatusPending { return resp, nil @@ -417,12 +420,14 @@ func (s *PaymentService) QueryAndFinalizeRefund(ctx context.Context, oid int64) } pendingDetail := s.latestRefundPendingDetail(ctx, oid) + finishProviderCall := servertiming.ObserveDependency(ctx, "payment") resp, err := queryProvider.QueryRefund(ctx, payment.RefundQueryRequest{ TradeNo: o.PaymentTradeNo, OrderID: o.OutTradeNo, RefundID: pendingDetail.RefundID, Amount: formatGatewayRefundAmount(o.RefundAmount, o), }) + finishProviderCall() if err != nil { return nil, fmt.Errorf("query refund: %w", err) } diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index 100b240785..507e02d877 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -2006,9 +2006,19 @@ func parseOpenAIImageTryAgainCooldown(body []byte) time.Duration { const upstreamModelNotFoundCooldown = 30 * time.Minute const upstreamModelNotFoundReason = "upstream_404_model_not_found" +const upstreamCodexPlanGatedModelCooldown = 30 * time.Minute +const upstreamCodexPlanGatedModelReason = "upstream_400_codex_plan_gated_model" const tempUnschedBodyMaxBytes = 64 << 10 const tempUnschedMessageMaxBytes = 2048 +// HandleUpstreamModelNotFound marks the requested model as temporarily +// unavailable on the account when the upstream deterministically reports it +// cannot serve that model: a 404 model-not-found, or the Codex 400 rejecting a +// plan-gated model on a ChatGPT OAuth account. Returning true tells the caller +// to fail the current attempt over to another account; the scheduler skips the +// (account, model) pair via IsSchedulableForModelWithContext until the +// cooldown expires, instead of re-selecting an account that can never serve +// the model. func (s *RateLimitService) HandleUpstreamModelNotFound(ctx context.Context, account *Account, requestedModel string, statusCode int, responseBody []byte) bool { if s == nil || account == nil || s.accountRepo == nil { return false @@ -2016,19 +2026,26 @@ func (s *RateLimitService) HandleUpstreamModelNotFound(ctx context.Context, acco if !account.ShouldHandleErrorCode(statusCode) { return false } - if !isUpstreamModelNotFoundError(statusCode, responseBody) { + var cooldown time.Duration + var reason string + switch { + case isUpstreamModelNotFoundError(statusCode, responseBody): + cooldown, reason = upstreamModelNotFoundCooldown, upstreamModelNotFoundReason + case isOpenAIOAuthAccount(account) && isOpenAICodexPlanGatedModelError(statusCode, responseBody): + cooldown, reason = upstreamCodexPlanGatedModelCooldown, upstreamCodexPlanGatedModelReason + default: return false } modelKey := modelRateLimitKeyForUpstreamModelNotFound(ctx, account, requestedModel) if modelKey == "" { return false } - resetAt := time.Now().Add(upstreamModelNotFoundCooldown) - if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, resetAt, upstreamModelNotFoundReason); err != nil { - slog.Warn("upstream_model_not_found_set_model_rate_limit_failed", "account_id", account.ID, "model", modelKey, "error", err) + resetAt := time.Now().Add(cooldown) + if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, resetAt, reason); err != nil { + slog.Warn("upstream_model_not_found_set_model_rate_limit_failed", "account_id", account.ID, "model", modelKey, "reason", reason, "error", err) return true } - slog.Info("upstream_model_not_found_model_rate_limited", "account_id", account.ID, "model", modelKey, "reset_at", resetAt) + slog.Info("upstream_model_not_found_model_rate_limited", "account_id", account.ID, "model", modelKey, "reason", reason, "reset_at", resetAt) return true } diff --git a/backend/internal/service/ratelimit_service_model_not_found_test.go b/backend/internal/service/ratelimit_service_model_not_found_test.go index dfd18c5f69..51bd8a607e 100644 --- a/backend/internal/service/ratelimit_service_model_not_found_test.go +++ b/backend/internal/service/ratelimit_service_model_not_found_test.go @@ -125,3 +125,77 @@ func openAIModelNotFoundTempAccount() *Account { }, } } + +func TestRateLimitService_HandleUpstreamError_CodexPlanGatedModelUsesModelRateLimit(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAICodexPlanGatedOAuthAccount() + + handled := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusBadRequest, + http.Header{}, + []byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`), + "gpt-5.6-sol", + ) + + require.True(t, handled) + require.Zero(t, repo.tempCalls) + require.Len(t, repo.modelRateLimitCalls, 1) + call := repo.modelRateLimitCalls[0] + require.Equal(t, account.ID, call.accountID) + require.Equal(t, "gpt-5.6-sol", call.scope) + require.Equal(t, upstreamCodexPlanGatedModelReason, call.reason) + require.WithinDuration(t, time.Now().Add(upstreamCodexPlanGatedModelCooldown), call.resetAt, 5*time.Second) +} + +func TestRateLimitService_HandleUpstreamError_CodexPlanGatedModelRespectsModelMapping(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAICodexPlanGatedOAuthAccount() + account.Credentials["model_mapping"] = map[string]any{"gpt-5.6-sol": "gpt-5.6-sol-upstream"} + + handled := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusBadRequest, + http.Header{}, + []byte(`{"detail":"The 'gpt-5.6-sol-upstream' model is not supported when using Codex with a ChatGPT account."}`), + "gpt-5.6-sol", + ) + + require.True(t, handled) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, "gpt-5.6-sol-upstream", repo.modelRateLimitCalls[0].scope) +} + +func TestRateLimitService_HandleUpstreamError_CodexPlanGatedModelIgnoresAPIKeyAccount(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAICodexPlanGatedOAuthAccount() + account.Type = AccountTypeAPIKey + + handled := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusBadRequest, + http.Header{}, + []byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`), + "gpt-5.6-sol", + ) + + require.False(t, handled) + require.Empty(t, repo.modelRateLimitCalls) +} + +func openAICodexPlanGatedOAuthAccount() *Account { + return &Account{ + ID: 202, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Credentials: map[string]any{}, + } +} diff --git a/backend/internal/service/scheduler_outbox.go b/backend/internal/service/scheduler_outbox.go index 2b7665ad78..a44f2a3a30 100644 --- a/backend/internal/service/scheduler_outbox.go +++ b/backend/internal/service/scheduler_outbox.go @@ -17,6 +17,8 @@ type SchedulerOutboxEvent struct { // SchedulerOutboxRepository 提供调度 outbox 的读取接口。 type SchedulerOutboxRepository interface { ListAfterAndReleaseDedup(ctx context.Context, afterID int64, limit int) ([]SchedulerOutboxEvent, error) + // FirstCreatedAtAfter 返回指定水位之后第一条待消费事件的创建时间，不领取事件或修改去重键。 + FirstCreatedAtAfter(ctx context.Context, afterID int64) (time.Time, bool, error) MaxID(ctx context.Context) (int64, error) DeleteConsumedUpTo(ctx context.Context, watermark int64, limit int) (int64, error) TryAcquireCleanupLock(ctx context.Context) (SchedulerOutboxCleanupLease, bool, error) diff --git a/backend/internal/service/scheduler_snapshot_full_rebuild_test.go b/backend/internal/service/scheduler_snapshot_full_rebuild_test.go new file mode 100644 index 0000000000..09ec2790ce --- /dev/null +++ b/backend/internal/service/scheduler_snapshot_full_rebuild_test.go @@ -0,0 +1,145 @@ +package service + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type schedulerFullRebuildTestCache struct { + SchedulerCache + + mu sync.Mutex + listErr error + listCalls int + lockCalls int +} + +func (c *schedulerFullRebuildTestCache) ListBuckets(context.Context) ([]SchedulerBucket, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.listCalls++ + return nil, c.listErr +} + +func (c *schedulerFullRebuildTestCache) TryLockBucket(context.Context, SchedulerBucket, time.Duration) (bool, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.lockCalls++ + return false, nil +} + +func TestSchedulerSnapshotServiceFullRebuildCoalescesConcurrentRequestsIntoTrailingRun(t *testing.T) { + svc := &SchedulerSnapshotService{} + wantTrailingErr := errors.New("trailing rebuild failed") + firstStarted := make(chan struct{}) + releaseFirst := make(chan struct{}) + var releaseOnce sync.Once + release := func() { + releaseOnce.Do(func() { close(releaseFirst) }) + } + defer release() + + var calls atomic.Int32 + var active atomic.Int32 + var maxActive atomic.Int32 + run := func() error { + call := calls.Add(1) + currentActive := active.Add(1) + defer active.Add(-1) + for { + previousMax := maxActive.Load() + if currentActive <= previousMax || maxActive.CompareAndSwap(previousMax, currentActive) { + break + } + } + if call == 1 { + close(firstStarted) + <-releaseFirst + return nil + } + return wantTrailingErr + } + + firstResult := make(chan error, 1) + go func() { + firstResult <- svc.coalesceFullRebuild(run) + }() + + select { + case <-firstStarted: + case <-time.After(time.Second): + t.Fatal("first rebuild did not start") + } + + const followers = 20 + followerResults := make(chan error, followers) + for range followers { + go func() { + followerResults <- svc.coalesceFullRebuild(run) + }() + } + + require.Eventually(t, func() bool { + requested, _ := schedulerFullRebuildState(svc) + return requested == followers+1 + }, time.Second, time.Millisecond) + release() + + require.NoError(t, <-firstResult) + for range followers { + require.ErrorIs(t, <-followerResults, wantTrailingErr) + } + require.EqualValues(t, 2, calls.Load()) + require.EqualValues(t, 1, maxActive.Load()) + requested, completed := schedulerFullRebuildState(svc) + require.EqualValues(t, followers+1, requested) + require.Equal(t, requested, completed) +} + +func TestSchedulerSnapshotServiceFullRebuildRunsAgainForSequentialRequest(t *testing.T) { + svc := &SchedulerSnapshotService{} + wantSecondErr := errors.New("second rebuild failed") + var calls atomic.Int32 + run := func() error { + if calls.Add(1) == 2 { + return wantSecondErr + } + return nil + } + + require.NoError(t, svc.coalesceFullRebuild(run)) + require.ErrorIs(t, svc.coalesceFullRebuild(run), wantSecondErr) + require.EqualValues(t, 2, calls.Load()) + requested, completed := schedulerFullRebuildState(svc) + require.EqualValues(t, 2, requested) + require.Equal(t, requested, completed) +} + +func TestSchedulerSnapshotServiceInitialFullRebuildFallsBackWhenListBucketsFails(t *testing.T) { + cache := &schedulerFullRebuildTestCache{listErr: errors.New("list buckets failed")} + svc := NewSchedulerSnapshotService(cache, nil, nil, nil, nil) + + svc.runInitialRebuild() + + cache.mu.Lock() + listCalls := cache.listCalls + lockCalls := cache.lockCalls + cache.mu.Unlock() + require.Equal(t, 1, listCalls) + require.Positive(t, lockCalls, "startup should rebuild default buckets after ListBuckets fails") + requested, completed := schedulerFullRebuildState(svc) + require.EqualValues(t, 1, requested) + require.Equal(t, requested, completed) +} + +func schedulerFullRebuildState(svc *SchedulerSnapshotService) (requested uint64, completed uint64) { + svc.fullRebuildStateMu.Lock() + defer svc.fullRebuildStateMu.Unlock() + return svc.fullRebuildRequested, svc.fullRebuildCompleted +} diff --git a/backend/internal/service/scheduler_snapshot_outbox_cleanup_test.go b/backend/internal/service/scheduler_snapshot_outbox_cleanup_test.go index 535f8d54e2..91e2d36f76 100644 --- a/backend/internal/service/scheduler_snapshot_outbox_cleanup_test.go +++ b/backend/internal/service/scheduler_snapshot_outbox_cleanup_test.go @@ -6,12 +6,15 @@ import ( "reflect" "testing" "time" + + "github.com/Wei-Shaw/sub2api/internal/config" ) type outboxCleanupCache struct { - watermark int64 - setWatermarks []int64 - updateErr error + watermark int64 + setWatermarks []int64 + updateErr error + listBucketCalls int } func (c *outboxCleanupCache) GetSnapshot(ctx context.Context, bucket SchedulerBucket) ([]*Account, bool, error) { @@ -47,6 +50,7 @@ func (c *outboxCleanupCache) UnlockBucket(ctx context.Context, bucket SchedulerB } func (c *outboxCleanupCache) ListBuckets(ctx context.Context) ([]SchedulerBucket, error) { + c.listBucketCalls++ return nil, nil } @@ -66,12 +70,13 @@ type outboxCleanupDeleteCall struct { } type outboxCleanupRepo struct { - events []SchedulerOutboxEvent - rows []int64 - lockAcquired bool - lockAttempts int - releaseCount int - deleteCalls []outboxCleanupDeleteCall + events []SchedulerOutboxEvent + rows []int64 + lockAcquired bool + lockAttempts int + releaseCount int + deleteCalls []outboxCleanupDeleteCall + firstCreatedAfterID []int64 } func (r *outboxCleanupRepo) ListAfterAndReleaseDedup(ctx context.Context, afterID int64, limit int) ([]SchedulerOutboxEvent, error) { @@ -88,6 +93,16 @@ func (r *outboxCleanupRepo) ListAfterAndReleaseDedup(ctx context.Context, afterI return events, nil } +func (r *outboxCleanupRepo) FirstCreatedAtAfter(ctx context.Context, afterID int64) (time.Time, bool, error) { + r.firstCreatedAfterID = append(r.firstCreatedAfterID, afterID) + for _, event := range r.events { + if event.ID > afterID { + return event.CreatedAt, true, nil + } + } + return time.Time{}, false, nil +} + func (r *outboxCleanupRepo) MaxID(ctx context.Context) (int64, error) { var maxID int64 for _, id := range r.rows { @@ -240,6 +255,44 @@ func TestSchedulerSnapshotServicePollOutboxDoesNotCleanupOnHandleFailure(t *test } } +func TestSchedulerSnapshotServicePollOutboxDoesNotUseConsumedEventForLag(t *testing.T) { + cache := &outboxCleanupCache{} + repo := &outboxCleanupRepo{ + events: []SchedulerOutboxEvent{ + { + ID: 7, + EventType: SchedulerOutboxEventAccountLastUsed, + CreatedAt: time.Now().Add(-time.Hour), + }, + }, + } + cfg := &config.Config{ + Gateway: config.GatewayConfig{ + Scheduling: config.GatewaySchedulingConfig{ + OutboxLagWarnSeconds: 1, + OutboxLagRebuildSeconds: 1, + OutboxLagRebuildFailures: 1, + }, + }, + } + svc := NewSchedulerSnapshotService(cache, repo, nil, nil, cfg) + + svc.pollOutbox() + + if cache.watermark != 7 { + t.Fatalf("expected watermark 7, got %d", cache.watermark) + } + if !reflect.DeepEqual(repo.firstCreatedAfterID, []int64{7}) { + t.Fatalf("expected lag check after consumed watermark, got %#v", repo.firstCreatedAfterID) + } + if cache.listBucketCalls != 0 { + t.Fatalf("expected consumed event not to trigger full rebuild, got %d attempts", cache.listBucketCalls) + } + if svc.lagFailures != 0 { + t.Fatalf("expected lag failures to remain reset, got %d", svc.lagFailures) + } +} + func TestSchedulerSnapshotServiceCleanupSkipsNonPositiveWatermark(t *testing.T) { repo := &outboxCleanupRepo{ rows: []int64{1, 2, 3}, diff --git a/backend/internal/service/scheduler_snapshot_service.go b/backend/internal/service/scheduler_snapshot_service.go index dc514bc851..77bef26edf 100644 --- a/backend/internal/service/scheduler_snapshot_service.go +++ b/backend/internal/service/scheduler_snapshot_service.go @@ -43,6 +43,12 @@ type SchedulerSnapshotService struct { fallbackLimit *fallbackLimiter lagMu sync.Mutex lagFailures int + + fullRebuildRunMu sync.Mutex + fullRebuildStateMu sync.Mutex + fullRebuildRequested uint64 + fullRebuildCompleted uint64 + fullRebuildLastErr error } func NewSchedulerSnapshotService( @@ -183,22 +189,26 @@ func (s *SchedulerSnapshotService) runInitialRebuild() { if s.cache == nil { return } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() - buckets, err := s.cache.ListBuckets(ctx) - if err != nil { - logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] list buckets failed: %v", err) - } - if len(buckets) == 0 { - buckets, err = s.defaultBuckets(ctx) + _ = s.coalesceFullRebuild(func() error { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + buckets, err := s.cache.ListBuckets(ctx) if err != nil { - logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] default buckets failed: %v", err) - return + logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] list buckets failed: %v", err) } - } - if err := s.rebuildBuckets(ctx, buckets, "startup"); err != nil { - logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild startup failed: %v", err) - } + if len(buckets) == 0 { + buckets, err = s.defaultBuckets(ctx) + if err != nil { + logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] default buckets failed: %v", err) + return err + } + } + if err := s.rebuildBuckets(ctx, buckets, "startup"); err != nil { + logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild startup failed: %v", err) + return err + } + return nil + }) } func (s *SchedulerSnapshotService) runOutboxWorker(interval time.Duration) { @@ -254,7 +264,6 @@ func (s *SchedulerSnapshotService) pollOutbox() { return } - watermarkForCheck := watermark seen := make(map[batchSeenKey]struct{}) for _, event := range events { eventCtx, cancel := context.WithTimeout(context.Background(), outboxEventTimeout) @@ -281,12 +290,15 @@ func (s *SchedulerSnapshotService) pollOutbox() { } if wmErr != nil { logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox watermark write failed: %v", wmErr) - } else { - watermarkForCheck = lastID - s.cleanupConsumedOutbox(lastID) + return } + s.cleanupConsumedOutbox(lastID) - s.checkOutboxLag(ctx, events[0], watermarkForCheck) + // 只有 watermark 成功推进后，当前批次才算已消费。延迟必须按下一条待消费事件计算， + // 否则本批次处理越慢，越容易误触发一次更慢的全量重建，形成正反馈。 + lagCtx, lagCancel := context.WithTimeout(context.Background(), 5*time.Second) + s.checkOutboxLag(lagCtx, lastID) + lagCancel() } func (s *SchedulerSnapshotService) cleanupConsumedOutbox(watermark int64) { @@ -602,30 +614,72 @@ func (s *SchedulerSnapshotService) triggerFullRebuild(reason string) error { if s.cache == nil { return ErrSchedulerCacheNotReady } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() + return s.coalesceFullRebuild(func() error { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() - buckets, err := s.cache.ListBuckets(ctx) - if err != nil { - logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] list buckets failed: %v", err) - return err - } - if len(buckets) == 0 { - buckets, err = s.defaultBuckets(ctx) + buckets, err := s.cache.ListBuckets(ctx) if err != nil { - logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] default buckets failed: %v", err) + logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] list buckets failed: %v", err) return err } - } - return s.rebuildBuckets(ctx, buckets, reason) + if len(buckets) == 0 { + buckets, err = s.defaultBuckets(ctx) + if err != nil { + logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] default buckets failed: %v", err) + return err + } + } + return s.rebuildBuckets(ctx, buckets, reason) + }) } -func (s *SchedulerSnapshotService) checkOutboxLag(ctx context.Context, oldest SchedulerOutboxEvent, watermark int64) { - if oldest.CreatedAt.IsZero() || s.cfg == nil { +func (s *SchedulerSnapshotService) coalesceFullRebuild(run func() error) error { + s.fullRebuildStateMu.Lock() + s.fullRebuildRequested++ + requestID := s.fullRebuildRequested + s.fullRebuildStateMu.Unlock() + + s.fullRebuildRunMu.Lock() + defer s.fullRebuildRunMu.Unlock() + + s.fullRebuildStateMu.Lock() + if s.fullRebuildCompleted >= requestID { + err := s.fullRebuildLastErr + s.fullRebuildStateMu.Unlock() + return err + } + // 当前轮重建可能早于新 outbox 事件对应事务的提交，不能让后到请求直接复用当前轮。 + // 每轮开始前记录可覆盖的请求代次，执行期间登记的请求统一合并到下一轮。 + coveredThrough := s.fullRebuildRequested + s.fullRebuildStateMu.Unlock() + + err := run() + + s.fullRebuildStateMu.Lock() + s.fullRebuildCompleted = coveredThrough + s.fullRebuildLastErr = err + s.fullRebuildStateMu.Unlock() + return err +} + +func (s *SchedulerSnapshotService) checkOutboxLag(ctx context.Context, watermark int64) { + if s.cfg == nil || s.outboxRepo == nil { + return + } + oldestCreatedAt, ok, err := s.outboxRepo.FirstCreatedAtAfter(ctx, watermark) + if err != nil { + logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox pending event read failed: %v", err) + return + } + if !ok || oldestCreatedAt.IsZero() { + s.lagMu.Lock() + s.lagFailures = 0 + s.lagMu.Unlock() return } - lag := time.Since(oldest.CreatedAt) + lag := time.Since(oldestCreatedAt) if lagSeconds := int(lag.Seconds()); lagSeconds >= s.cfg.Gateway.Scheduling.OutboxLagWarnSeconds && s.cfg.Gateway.Scheduling.OutboxLagWarnSeconds > 0 { logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox lag warning: %ds", lagSeconds) } @@ -652,7 +706,7 @@ func (s *SchedulerSnapshotService) checkOutboxLag(ctx context.Context, oldest Sc } threshold := s.cfg.Gateway.Scheduling.OutboxBacklogRebuildRows - if threshold <= 0 || s.outboxRepo == nil { + if threshold <= 0 { return } maxID, err := s.outboxRepo.MaxID(ctx) diff --git a/backend/internal/service/upstream_models.go b/backend/internal/service/upstream_models.go index 4f6a305b25..9b4fd6f532 100644 --- a/backend/internal/service/upstream_models.go +++ b/backend/internal/service/upstream_models.go @@ -131,6 +131,8 @@ func (s *AccountTestService) buildUpstreamModelsRequest(ctx context.Context, acc switch { case account.Platform == PlatformAntigravity: return s.buildAntigravityAPIKeyModelsRequest(ctx, account) + case account.IsGrok(): + return s.buildGrokUpstreamModelsRequest(ctx, account) case account.IsOpenAI(): return s.buildOpenAIUpstreamModelsRequest(ctx, account) case account.IsGemini(): @@ -144,6 +146,36 @@ func (s *AccountTestService) buildUpstreamModelsRequest(ctx context.Context, acc } } +func (s *AccountTestService) buildGrokUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) { + if account.Type != AccountTypeAPIKey { + return nil, newUpstreamModelSyncUnsupportedError( + fmt.Sprintf("Unsupported Grok account type for upstream model sync: %s", account.Type), nil, + ) + } + apiKey := strings.TrimSpace(account.GetCredential("api_key")) + if apiKey == "" { + return nil, newUpstreamModelSyncConfigError("No Grok API key is available", nil) + } + + baseURL := strings.TrimSpace(account.GetCredential("base_url")) + if baseURL == "" { + baseURL = "https://api.x.ai" + } + normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL) + if err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Grok base URL", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, buildOpenAIModelsURL(normalizedBaseURL), nil) + if err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Grok model list URL", err) + } + req.Header.Set("Accept", "application/json") + req.Header.Set("Authorization", "Bearer "+apiKey) + account.ApplyHeaderOverrides(req.Header) + return req, nil +} + func (s *AccountTestService) buildAnthropicUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) { if account.IsBedrock() || account.Type == AccountTypeServiceAccount { return nil, newUpstreamModelSyncUnsupportedError( diff --git a/backend/internal/service/upstream_models_test.go b/backend/internal/service/upstream_models_test.go index 3904194ffa..5b5c5e9835 100644 --- a/backend/internal/service/upstream_models_test.go +++ b/backend/internal/service/upstream_models_test.go @@ -177,6 +177,18 @@ func TestBuildUpstreamModelsRequestsForAPIKeyAccounts(t *testing.T) { require.Equal(t, "https://openai.example.com/v1/models", openAIReq.URL.String()) require.Equal(t, "Bearer openai-key", openAIReq.Header.Get("Authorization")) + grokReq, err := svc.buildUpstreamModelsRequest(ctx, &Account{ + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "xai-key", + "base_url": "https://xai.example.com/v1", + }, + }) + require.NoError(t, err) + require.Equal(t, "https://xai.example.com/v1/models", grokReq.URL.String()) + require.Equal(t, "Bearer xai-key", grokReq.Header.Get("Authorization")) + geminiReq, err := svc.buildGeminiUpstreamModelsRequest(ctx, &Account{ Platform: PlatformGemini, Type: AccountTypeAPIKey, @@ -202,6 +214,22 @@ func TestBuildUpstreamModelsRequestsForAPIKeyAccounts(t *testing.T) { require.Equal(t, "antigravity-key", antigravityReq.Header.Get("x-api-key")) } +func TestBuildUpstreamModelsRequestRejectsGrokOAuth(t *testing.T) { + t.Parallel() + + svc := &AccountTestService{cfg: upstreamModelSyncTestConfig()} + _, err := svc.buildUpstreamModelsRequest(context.Background(), &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + }) + require.Error(t, err) + + var syncErr *UpstreamModelSyncError + require.True(t, errors.As(err, &syncErr)) + require.Equal(t, UpstreamModelSyncErrorUnsupported, syncErr.Kind) + require.Contains(t, syncErr.SafeMessage(), "Unsupported Grok account type") +} + func TestBuildAntigravityAPIKeyModelsRequestRejectsOfficialCloudCodeBase(t *testing.T) { t.Parallel() @@ -265,6 +293,34 @@ func TestFetchUpstreamSupportedModelsParsesOpenAIResponse(t *testing.T) { require.Equal(t, "Bearer openai-key", upstream.lastReq.Header.Get("Authorization")) } +func TestFetchUpstreamSupportedModelsParsesGrokAPIKeyResponse(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"data":[{"id":"grok-4.5"},{"id":"grok-4.5"},{"id":"grok-imagine"}]}`)), + }} + svc := &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + } + + models, err := svc.FetchUpstreamSupportedModels(context.Background(), &Account{ + ID: 9, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "xai-key", + "base_url": "https://xai.example.com/v1", + }, + }) + require.NoError(t, err) + require.Equal(t, []string{"grok-4.5", "grok-imagine"}, models) + require.Equal(t, "https://xai.example.com/v1/models", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer xai-key", upstream.lastReq.Header.Get("Authorization")) +} + func TestFetchUpstreamSupportedModelsDoesNotExposeUpstreamBody(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/usage_log.go b/backend/internal/service/usage_log.go index 62e48fc8f9..0adcc04a94 100644 --- a/backend/internal/service/usage_log.go +++ b/backend/internal/service/usage_log.go @@ -142,13 +142,14 @@ type UsageLog struct { ImageOutputTokens int ImageOutputCost float64 - InputCost float64 - OutputCost float64 - CacheCreationCost float64 - CacheReadCost float64 - TotalCost float64 - ActualCost float64 - RateMultiplier float64 + InputCost float64 + OutputCost float64 + CacheCreationCost float64 + CacheReadCost float64 + TotalCost float64 + ActualCost float64 + RateMultiplier float64 + LongContextBillingApplied bool // AccountRateMultiplier 账号计费倍率快照（nil 表示历史数据，按 1.0 处理） AccountRateMultiplier *float64 // AccountStatsCost 账号统计定价预计算费用（nil = 使用默认公式 total_cost × account_rate_multiplier） diff --git a/backend/internal/service/vertex_service_account.go b/backend/internal/service/vertex_service_account.go index 256695ded5..7ccbeee43c 100644 --- a/backend/internal/service/vertex_service_account.go +++ b/backend/internal/service/vertex_service_account.go @@ -18,6 +18,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/golang-jwt/jwt/v5" ) @@ -195,7 +196,7 @@ func vertexServiceAccountProxyURL(account *Account) string { func newVertexServiceAccountHTTPClient(proxyURL string) (*http.Client, error) { proxyURL = strings.TrimSpace(proxyURL) if proxyURL == "" { - return &http.Client{Timeout: 15 * time.Second}, nil + return servertiming.InstrumentClient(&http.Client{Timeout: 15 * time.Second}), nil } _, parsedProxy, err := proxyurl.Parse(proxyURL) @@ -211,7 +212,7 @@ func newVertexServiceAccountHTTPClient(proxyURL string) (*http.Client, error) { if err := proxyutil.ConfigureTransportProxy(transport, parsedProxy); err != nil { return nil, err } - return &http.Client{Timeout: 15 * time.Second, Transport: transport}, nil + return servertiming.InstrumentClient(&http.Client{Timeout: 15 * time.Second, Transport: transport}), nil } func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAccountKey, proxyURL string) (string, time.Duration, error) { diff --git a/backend/internal/service/vertex_service_account_test.go b/backend/internal/service/vertex_service_account_test.go index d77a1988e9..68c756eaa2 100644 --- a/backend/internal/service/vertex_service_account_test.go +++ b/backend/internal/service/vertex_service_account_test.go @@ -13,6 +13,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" ) @@ -101,6 +102,24 @@ func TestVertexServiceAccountProxyURL(t *testing.T) { require.Empty(t, vertexServiceAccountProxyURL(&Account{ProxyID: &proxyID})) } +func TestVertexServiceAccountHTTPClientRecordsDependency(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + defer server.Close() + + client, err := newVertexServiceAccountHTTPClient("") + require.NoError(t, err) + collector := servertiming.New(time.Now()) + ctx := servertiming.WithCollector(context.Background(), collector) + request, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + require.NoError(t, err) + response, err := client.Do(request) + require.NoError(t, err) + require.NoError(t, response.Body.Close()) + require.Contains(t, collector.HeaderValue(time.Now(), "bypass"), "dep_http;dur=") +} + func TestExchangeVertexServiceAccountTokenUsesProxy(t *testing.T) { privateKey, err := rsa.GenerateKey(rand.Reader, 2048) require.NoError(t, err) diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 7258ff05a3..d5d9124ec2 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -140,8 +140,9 @@ func ProvideGrokQuotaService( proxyRepo ProxyRepository, tokenProvider *GrokTokenProvider, httpUpstream HTTPUpstream, + usageLogRepo UsageLogRepository, ) *GrokQuotaService { - return NewGrokQuotaService(accountRepo, proxyRepo, tokenProvider, httpUpstream) + return NewGrokQuotaService(accountRepo, proxyRepo, tokenProvider, httpUpstream, usageLogRepo) } // ProvideGeminiTokenProvider creates GeminiTokenProvider with OAuthRefreshAPI injection diff --git a/backend/internal/web/embed_on.go b/backend/internal/web/embed_on.go index 41738e7a5d..716fb77e75 100644 --- a/backend/internal/web/embed_on.go +++ b/backend/internal/web/embed_on.go @@ -109,7 +109,8 @@ func (s *FrontendServer) Middleware() gin.HandlerFunc { return } - // Serve static files normally + // Serve static files normally (hashed assets get long-lived cache headers) + applyStaticAssetCacheHeaders(c.Writer.Header(), cleanPath) s.fileServer.ServeHTTP(c.Writer, c.Request) c.Abort() } @@ -135,6 +136,7 @@ func (s *FrontendServer) tryServeOverride(c *gin.Context, cleanPath string) bool if err != nil || info.IsDir() { return false } + applyStaticAssetCacheHeaders(c.Writer.Header(), cleanPath) c.File(filePath) c.Abort() return true @@ -273,6 +275,7 @@ func ServeEmbeddedFrontend() gin.HandlerFunc { if tryServeOverrideFile(c, overrideDir, cleanPath) { return } + applyStaticAssetCacheHeaders(c.Writer.Header(), cleanPath) fileServer.ServeHTTP(c.Writer, c.Request) c.Abort() return @@ -292,6 +295,7 @@ func tryServeOverrideFile(c *gin.Context, overrideDir, cleanPath string) bool { if err != nil || info.IsDir() { return false } + applyStaticAssetCacheHeaders(c.Writer.Header(), cleanPath) c.File(filePath) c.Abort() return true @@ -308,7 +312,9 @@ func shouldBypassEmbeddedFrontend(path string) bool { trimmed == "/health" || trimmed == "/responses" || strings.HasPrefix(trimmed, "/responses/") || - strings.HasPrefix(trimmed, "/images/") + trimmed == "/alpha/search" || + strings.HasPrefix(trimmed, "/images/") || + strings.HasPrefix(trimmed, "/videos/") } func serveIndexHTML(c *gin.Context, fsys fs.FS) { diff --git a/backend/internal/web/embed_test.go b/backend/internal/web/embed_test.go index 27e15ef166..b27bbfc9dc 100644 --- a/backend/internal/web/embed_test.go +++ b/backend/internal/web/embed_test.go @@ -507,6 +507,32 @@ func TestFrontendServer_Middleware(t *testing.T) { assert.JSONEq(t, `{"ok":true}`, w.Body.String()) }) + t.Run("skips_alpha_search_post_route", func(t *testing.T) { + provider := &mockSettingsProvider{ + settings: map[string]string{"test": "value"}, + } + + server, err := NewFrontendServer(provider) + require.NoError(t, err) + + router := gin.New() + router.Use(server.Middleware()) + nextCalled := false + router.POST("/alpha/search", func(c *gin.Context) { + nextCalled = true + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/alpha/search", strings.NewReader(`{"model":"gpt-5.6-sol"}`)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(w, req) + + assert.True(t, nextCalled, "next handler should be called for alpha search API route") + assert.Equal(t, http.StatusOK, w.Code) + assert.JSONEq(t, `{"ok":true}`, w.Body.String()) + }) + t.Run("serves_index_for_spa_routes", func(t *testing.T) { provider := &mockSettingsProvider{ settings: map[string]string{"test": "value"}, @@ -562,6 +588,17 @@ func TestFrontendServer_Middleware(t *testing.T) { }) } +func TestEmbeddedFrontendBypassesBareVideoAPIRoutes(t *testing.T) { + for _, path := range []string{ + "/videos/generations", + "/videos/edits", + "/videos/extensions", + "/videos/request-123", + } { + require.True(t, shouldBypassEmbeddedFrontend(path), "path=%s", path) + } +} + func TestNewFrontendServer(t *testing.T) { t.Run("creates_server_successfully", func(t *testing.T) { provider := &mockSettingsProvider{ diff --git a/backend/internal/web/static_cache.go b/backend/internal/web/static_cache.go new file mode 100644 index 0000000000..09abc8642a --- /dev/null +++ b/backend/internal/web/static_cache.go @@ -0,0 +1,31 @@ +//go:build embed || unit + +package web + +import ( + "net/http" + "strings" +) + +// staticAssetsCacheControl matches deploy/Caddyfile for hashed frontend assets. +// Vite emits content-hashed filenames under assets/, so long-lived immutable +// caching is safe without relying on a reverse proxy. +const staticAssetsCacheControl = "public, max-age=31536000, immutable" + +// isLongCacheStaticPath reports whether a cleaned URL path (no leading slash) +// should receive long-lived Cache-Control headers. Aligned with deploy/Caddyfile. +func isLongCacheStaticPath(cleanPath string) bool { + cleanPath = strings.TrimPrefix(cleanPath, "/") + return strings.HasPrefix(cleanPath, "assets/") || + cleanPath == "logo.png" || + cleanPath == "favicon.ico" +} + +// applyStaticAssetCacheHeaders sets Cache-Control for long-cacheable static paths. +// index.html / SPA routes must keep no-cache and are not handled here. +func applyStaticAssetCacheHeaders(header http.Header, cleanPath string) { + if header == nil || !isLongCacheStaticPath(cleanPath) { + return + } + header.Set("Cache-Control", staticAssetsCacheControl) +} diff --git a/backend/internal/web/static_cache_test.go b/backend/internal/web/static_cache_test.go new file mode 100644 index 0000000000..130347c41e --- /dev/null +++ b/backend/internal/web/static_cache_test.go @@ -0,0 +1,71 @@ +//go:build unit + +package web + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestIsLongCacheStaticPath(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + path string + want bool + }{ + {name: "hashed_js", path: "assets/index-abc123.js", want: true}, + {name: "hashed_css", path: "assets/app-def456.css", want: true}, + {name: "nested_asset", path: "assets/vendor/chunk.js", want: true}, + {name: "leading_slash_asset", path: "/assets/index.js", want: true}, + {name: "logo", path: "logo.png", want: true}, + {name: "favicon", path: "favicon.ico", want: true}, + {name: "index_html", path: "index.html", want: false}, + {name: "spa_route", path: "dashboard", want: false}, + {name: "assets_prefix_only", path: "assets", want: false}, + {name: "similar_name", path: "assets-backup/x.js", want: false}, + {name: "empty", path: "", want: false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tc.want, isLongCacheStaticPath(tc.path)) + }) + } +} + +func TestApplyStaticAssetCacheHeaders(t *testing.T) { + t.Parallel() + + t.Run("sets_immutable_cache_for_assets", func(t *testing.T) { + t.Parallel() + header := make(http.Header) + applyStaticAssetCacheHeaders(header, "assets/index-abc.js") + assert.Equal(t, staticAssetsCacheControl, header.Get("Cache-Control")) + }) + + t.Run("sets_immutable_cache_for_logo", func(t *testing.T) { + t.Parallel() + header := make(http.Header) + applyStaticAssetCacheHeaders(header, "logo.png") + assert.Equal(t, staticAssetsCacheControl, header.Get("Cache-Control")) + }) + + t.Run("skips_index_html", func(t *testing.T) { + t.Parallel() + header := make(http.Header) + applyStaticAssetCacheHeaders(header, "index.html") + assert.Empty(t, header.Get("Cache-Control")) + }) + + t.Run("nil_header_is_noop", func(t *testing.T) { + t.Parallel() + assert.NotPanics(t, func() { + applyStaticAssetCacheHeaders(nil, "assets/x.js") + }) + }) +} diff --git a/backend/migrations/174_add_usage_log_long_context_billing.sql b/backend/migrations/174_add_usage_log_long_context_billing.sql new file mode 100644 index 0000000000..090403c310 --- /dev/null +++ b/backend/migrations/174_add_usage_log_long_context_billing.sql @@ -0,0 +1,4 @@ +-- Snapshot whether long-context pricing changed token prices for a request so +-- usage history can explain the applied charge without inferring from totals. +ALTER TABLE usage_logs + ADD COLUMN IF NOT EXISTS long_context_billing_applied BOOLEAN NOT NULL DEFAULT FALSE; diff --git a/backend/migrations/174_add_usage_logs_api_key_latest_ip_index_notx.sql b/backend/migrations/174_add_usage_logs_api_key_latest_ip_index_notx.sql new file mode 100644 index 0000000000..261698f8cc --- /dev/null +++ b/backend/migrations/174_add_usage_logs_api_key_latest_ip_index_notx.sql @@ -0,0 +1,5 @@ +-- Support the per-key latest non-empty source IP lookup without scanning full key history. +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 <> ''; diff --git a/backend/migrations/175_add_ops_system_logs_host.sql b/backend/migrations/175_add_ops_system_logs_host.sql new file mode 100644 index 0000000000..e5f9f7299c --- /dev/null +++ b/backend/migrations/175_add_ops_system_logs_host.sql @@ -0,0 +1,3 @@ +-- Track the application host that emitted each indexed system log. +ALTER TABLE ops_system_logs + ADD COLUMN IF NOT EXISTS host VARCHAR(255); diff --git a/backend/migrations/175_default_openai_long_context_billing.sql b/backend/migrations/175_default_openai_long_context_billing.sql new file mode 100644 index 0000000000..cccbea4108 --- /dev/null +++ b/backend/migrations/175_default_openai_long_context_billing.sql @@ -0,0 +1,162 @@ +-- Keep mixed-version writers consistent before backfilling rows that already exist. +CREATE OR REPLACE FUNCTION public.enforce_openai_long_context_billing_extra() +RETURNS TRIGGER +LANGUAGE plpgsql +AS $$ +DECLARE + parent_effective_value JSONB; +BEGIN + IF NEW.platform IS DISTINCT FROM 'openai' THEN + RETURN NEW; + END IF; + + NEW.extra := COALESCE(NEW.extra, '{}'::jsonb); + IF NEW.parent_account_id IS NOT NULL AND NEW.quota_dimension = 'spark' THEN + SELECT CASE + WHEN parent.platform IS DISTINCT FROM 'openai' THEN 'false'::jsonb + WHEN NOT (COALESCE(parent.extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled') THEN 'false'::jsonb + WHEN jsonb_typeof(parent.extra->'openai_long_context_billing_enabled') = 'boolean' + THEN parent.extra->'openai_long_context_billing_enabled' + ELSE 'false'::jsonb + END + INTO parent_effective_value + FROM accounts AS parent + WHERE parent.id = NEW.parent_account_id; + + NEW.extra := jsonb_set( + NEW.extra, + '{openai_long_context_billing_enabled}', + COALESCE(parent_effective_value, 'false'::jsonb), + true + ); + ELSIF NOT (NEW.extra ? 'openai_long_context_billing_enabled') + AND TG_OP = 'UPDATE' + AND OLD.platform = 'openai' + AND jsonb_typeof(OLD.extra->'openai_long_context_billing_enabled') = 'boolean' THEN + NEW.extra := jsonb_set( + NEW.extra, + '{openai_long_context_billing_enabled}', + OLD.extra->'openai_long_context_billing_enabled', + true + ); + ELSIF NOT (NEW.extra ? 'openai_long_context_billing_enabled') THEN + NEW.extra := jsonb_set( + NEW.extra, + '{openai_long_context_billing_enabled}', + 'false'::jsonb, + true + ); + END IF; + + IF jsonb_typeof(NEW.extra->'openai_long_context_billing_enabled') IS DISTINCT FROM 'boolean' THEN + RAISE EXCEPTION 'openai_long_context_billing_enabled must be a boolean' + USING ERRCODE = '22023'; + END IF; + RETURN NEW; +END; +$$; + +CREATE OR REPLACE FUNCTION public.propagate_openai_long_context_billing_extra_to_shadows() +RETURNS TRIGGER +LANGUAGE plpgsql +AS $$ +BEGIN + WITH updated_shadows AS ( + UPDATE accounts AS shadow + SET extra = jsonb_set( + COALESCE(shadow.extra, '{}'::jsonb), + '{openai_long_context_billing_enabled}', + NEW.extra->'openai_long_context_billing_enabled', + true + ) + WHERE shadow.parent_account_id = NEW.id + AND shadow.platform = 'openai' + AND shadow.quota_dimension = 'spark' + AND shadow.extra->'openai_long_context_billing_enabled' + IS DISTINCT FROM NEW.extra->'openai_long_context_billing_enabled' + RETURNING shadow.id + ) + INSERT INTO scheduler_outbox (event_type, account_id) + SELECT 'account_changed', id + FROM updated_shadows; + RETURN NULL; +END; +$$; + +DROP TRIGGER IF EXISTS accounts_enforce_openai_long_context_billing_extra ON accounts; +CREATE TRIGGER accounts_enforce_openai_long_context_billing_extra +BEFORE INSERT OR UPDATE OF platform, extra, parent_account_id, quota_dimension +ON accounts +FOR EACH ROW +EXECUTE FUNCTION public.enforce_openai_long_context_billing_extra(); + +DROP TRIGGER IF EXISTS accounts_propagate_openai_long_context_billing_extra ON accounts; +CREATE TRIGGER accounts_propagate_openai_long_context_billing_extra +AFTER UPDATE OF platform, extra +ON accounts +FOR EACH ROW +WHEN ( + NEW.platform = 'openai' + AND NEW.parent_account_id IS NULL + AND ( + OLD.platform IS DISTINCT FROM NEW.platform + OR OLD.extra->'openai_long_context_billing_enabled' + IS DISTINCT FROM NEW.extra->'openai_long_context_billing_enabled' + ) +) +EXECUTE FUNCTION public.propagate_openai_long_context_billing_extra_to_shadows(); + +UPDATE accounts +SET extra = jsonb_set( + COALESCE(extra, '{}'::jsonb), + '{openai_long_context_billing_enabled}', + 'false'::jsonb, + true +) +WHERE platform = 'openai' + AND COALESCE(extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled' + AND jsonb_typeof(extra->'openai_long_context_billing_enabled') IS DISTINCT FROM 'boolean'; + +UPDATE accounts +SET extra = jsonb_set( + COALESCE(extra, '{}'::jsonb), + '{openai_long_context_billing_enabled}', + 'false'::jsonb, + true +) +WHERE platform = 'openai' + AND parent_account_id IS NULL + AND NOT (COALESCE(extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled'); + +WITH shadow_values AS ( + SELECT + shadow.id, + CASE + WHEN parent.platform IS DISTINCT FROM 'openai' THEN 'false'::jsonb + WHEN NOT (COALESCE(parent.extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled') THEN 'false'::jsonb + WHEN jsonb_typeof(parent.extra->'openai_long_context_billing_enabled') = 'boolean' + THEN parent.extra->'openai_long_context_billing_enabled' + ELSE 'false'::jsonb + END AS effective_value + FROM accounts AS shadow + JOIN accounts AS parent ON parent.id = shadow.parent_account_id + WHERE shadow.platform = 'openai' + AND shadow.quota_dimension = 'spark' +), +updated_shadows AS ( + UPDATE accounts AS shadow + SET extra = jsonb_set( + COALESCE(shadow.extra, '{}'::jsonb), + '{openai_long_context_billing_enabled}', + shadow_values.effective_value, + true + ) + FROM shadow_values + WHERE shadow.id = shadow_values.id + AND shadow.extra->'openai_long_context_billing_enabled' + IS DISTINCT FROM shadow_values.effective_value + RETURNING shadow.id +) +INSERT INTO scheduler_outbox (event_type, account_id) +SELECT 'account_changed', id +FROM updated_shadows; diff --git a/backend/migrations/175a_add_ops_system_logs_host_index_notx.sql b/backend/migrations/175a_add_ops_system_logs_host_index_notx.sql new file mode 100644 index 0000000000..ec2705e49b --- /dev/null +++ b/backend/migrations/175a_add_ops_system_logs_host_index_notx.sql @@ -0,0 +1,2 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_ops_system_logs_host_created_at + ON ops_system_logs (host, created_at DESC); diff --git a/backend/migrations/176_channel_monitor_grok_provider.sql b/backend/migrations/176_channel_monitor_grok_provider.sql new file mode 100644 index 0000000000..b1bad754a4 --- /dev/null +++ b/backend/migrations/176_channel_monitor_grok_provider.sql @@ -0,0 +1,39 @@ +-- Migration: 176_channel_monitor_grok_provider +-- Allow Grok as a channel-monitor provider. Grok checks use the existing +-- OpenAI-compatible chat completions protocol with model grok-4.5 by default. + +DO $$ +DECLARE + monitor_constraint_def TEXT; + template_constraint_def TEXT; +BEGIN + SELECT pg_get_constraintdef(c.oid) + INTO monitor_constraint_def + FROM pg_constraint c + JOIN pg_class t ON t.oid = c.conrelid + WHERE t.relname = 'channel_monitors' + AND c.conname = 'channel_monitors_provider_check'; + + IF monitor_constraint_def IS NULL OR position('grok' IN monitor_constraint_def) = 0 THEN + ALTER TABLE channel_monitors + DROP CONSTRAINT IF EXISTS channel_monitors_provider_check; + ALTER TABLE channel_monitors + ADD CONSTRAINT channel_monitors_provider_check + CHECK (provider IN ('openai', 'anthropic', 'gemini', 'grok')); + END IF; + + SELECT pg_get_constraintdef(c.oid) + INTO template_constraint_def + FROM pg_constraint c + JOIN pg_class t ON t.oid = c.conrelid + WHERE t.relname = 'channel_monitor_request_templates' + AND c.conname = 'channel_monitor_request_templates_provider_check'; + + IF template_constraint_def IS NULL OR position('grok' IN template_constraint_def) = 0 THEN + ALTER TABLE channel_monitor_request_templates + DROP CONSTRAINT IF EXISTS channel_monitor_request_templates_provider_check; + ALTER TABLE channel_monitor_request_templates + ADD CONSTRAINT channel_monitor_request_templates_provider_check + CHECK (provider IN ('openai', 'anthropic', 'gemini', 'grok')); + END IF; +END $$; diff --git a/backend/migrations/channel_monitor_grok_provider_migration_test.go b/backend/migrations/channel_monitor_grok_provider_migration_test.go new file mode 100644 index 0000000000..2545173f0e --- /dev/null +++ b/backend/migrations/channel_monitor_grok_provider_migration_test.go @@ -0,0 +1,20 @@ +package migrations + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestChannelMonitorGrokProviderMigration(t *testing.T) { + content, err := FS.ReadFile("176_channel_monitor_grok_provider.sql") + require.NoError(t, err) + + sql := strings.Join(strings.Fields(string(content)), " ") + require.Contains(t, sql, "channel_monitors_provider_check") + require.Contains(t, sql, "channel_monitor_request_templates_provider_check") + require.Contains(t, sql, "CHECK (provider IN ('openai', 'anthropic', 'gemini', 'grok'))") + require.Contains(t, sql, "position('grok' IN monitor_constraint_def) = 0") + require.Contains(t, sql, "position('grok' IN template_constraint_def) = 0") +} diff --git a/backend/migrations/latest_api_key_ip_index_test.go b/backend/migrations/latest_api_key_ip_index_test.go new file mode 100644 index 0000000000..1de64a9ff5 --- /dev/null +++ b/backend/migrations/latest_api_key_ip_index_test.go @@ -0,0 +1,19 @@ +package migrations + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestLatestAPIKeyIPIndexMigration(t *testing.T) { + content, err := FS.ReadFile("174_add_usage_logs_api_key_latest_ip_index_notx.sql") + require.NoError(t, err) + + sql := strings.Join(strings.Fields(string(content)), " ") + require.Contains(t, sql, "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_api_key_latest_ip") + require.Contains(t, sql, "ON usage_logs (api_key_id, created_at DESC, id DESC)") + require.Contains(t, sql, "INCLUDE (ip_address)") + require.Contains(t, sql, "WHERE ip_address IS NOT NULL AND ip_address <> ''") +} diff --git a/backend/migrations/openai_long_context_billing_migration_test.go b/backend/migrations/openai_long_context_billing_migration_test.go new file mode 100644 index 0000000000..212ac15d9d --- /dev/null +++ b/backend/migrations/openai_long_context_billing_migration_test.go @@ -0,0 +1,36 @@ +package migrations + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestMigration175DefaultsOrdinaryOpenAIAndInheritsForSparkShadows(t *testing.T) { + content, err := FS.ReadFile("175_default_openai_long_context_billing.sql") + require.NoError(t, err) + + sql := string(content) + require.Contains(t, sql, "parent_account_id IS NULL") + require.Contains(t, sql, "quota_dimension = 'spark'") + require.Contains(t, sql, "parent.extra") + require.Contains(t, sql, "jsonb_typeof") + require.Contains(t, sql, "openai_long_context_billing_enabled") +} + +func TestMigration175GuardsMixedVersionAccountWrites(t *testing.T) { + content, err := FS.ReadFile("175_default_openai_long_context_billing.sql") + require.NoError(t, err) + + sql := string(content) + require.Contains(t, sql, "RETURNS TRIGGER") + require.Contains(t, sql, "BEFORE INSERT OR UPDATE") + require.Contains(t, sql, "CREATE TRIGGER") + require.Contains(t, sql, "must be a boolean") + require.Contains(t, sql, "INSERT INTO scheduler_outbox") + require.Contains(t, sql, "'account_changed'") + require.Contains(t, sql, "jsonb_typeof(extra->'openai_long_context_billing_enabled') IS DISTINCT FROM 'boolean'") + require.Contains(t, sql, "WITH shadow_values AS") + require.Contains(t, sql, "TG_OP = 'UPDATE'") + require.Contains(t, sql, "OLD.extra->'openai_long_context_billing_enabled'") +} diff --git a/deploy/.env.example b/deploy/.env.example index 5925f0abb4..57056b4907 100644 --- a/deploy/.env.example +++ b/deploy/.env.example @@ -1,25 +1,37 @@ # ============================================================================= -# Sub2API Docker Environment Configuration +# Sub2API Container Environment Configuration # ============================================================================= # Copy this file to .env and modify as needed: # cp .env.example .env +# chmod 600 .env # nano .env # -# Then start with: docker-compose up -d +# Then start with Docker Compose or Apple container: +# docker compose up -d +# ./apple-container.sh up # ============================================================================= # ----------------------------------------------------------------------------- # Server Configuration # ----------------------------------------------------------------------------- -# Bind address for host port mapping +# IPv4 bind address for host port mapping BIND_HOST=0.0.0.0 -# Server port (exposed on host) +# Server port exposed on the host (Apple container requires 1025-65535) SERVER_PORT=8080 # Server mode: release or debug SERVER_MODE=release +# Return Server-Timing for authenticated requests made by the Admin web UI +ENABLE_SERVER_TIMING=false + +# Apple container image overrides (ignored by Docker Compose). Pin release tags +# or digests for repeatable operator-managed deployments. +APPLE_CONTAINER_SUB2API_IMAGE=weishaw/sub2api:latest +APPLE_CONTAINER_POSTGRES_IMAGE=postgres:18-alpine +APPLE_CONTAINER_REDIS_IMAGE=redis:8-alpine + # ----------------------------------------------------------------------------- # Logging Configuration # 日志配置 @@ -307,13 +319,15 @@ GATEWAY_SCHEDULING_OUTBOX_BACKLOG_REBUILD_ROWS=10000 GATEWAY_SCHEDULING_FULL_REBUILD_INTERVAL_SECONDS=300 # ----------------------------------------------------------------------------- -# Image Generation Stream & Concurrency (Optional) -# 图片生成流式与并发隔离配置（可选） +# Image Generation Keepalive & Concurrency (Optional) +# 图片生成保活与并发隔离配置（可选） # ----------------------------------------------------------------------------- # 图片流式上游数据间隔超时（秒）。0 表示禁用；非 0 时必须为 60-1800。 GATEWAY_IMAGE_STREAM_DATA_INTERVAL_TIMEOUT=900 # 图片流式 keepalive 间隔（秒）。0 表示禁用；非 0 时必须为 5-60。 GATEWAY_IMAGE_STREAM_KEEPALIVE_INTERVAL=10 +# 图片非流式 JSON keepalive 间隔（秒）。默认 0 禁用；首个心跳后 HTTP 状态会固化为 200。 +GATEWAY_IMAGE_NONSTREAM_KEEPALIVE_INTERVAL=0 # 是否启用进程级图片生成并发限制。默认 false，保持历史行为。 GATEWAY_IMAGE_CONCURRENCY_ENABLED=false # 当前进程允许同时处理的图片生成请求数。0 表示不限制。 diff --git a/deploy/APPLE_CONTAINER.md b/deploy/APPLE_CONTAINER.md new file mode 100644 index 0000000000..1133464a81 --- /dev/null +++ b/deploy/APPLE_CONTAINER.md @@ -0,0 +1,221 @@ +# Apple container Deployment + +Sub2API can run as a native three-service stack with Apple's `container` CLI. This workflow runs the published Sub2API, PostgreSQL, and Redis OCI images without Docker Desktop or a Docker-compatible daemon. + +## Support Level + +Apple `container` support is intended for local development and operator-managed deployments on a Mac. Docker Compose remains the recommended production deployment path. + +Apple `container` 1.1 does not provide restart policies, automatic startup, workload health scheduling, a Docker API socket, or full Compose orchestration. `apple-container.sh` supplies ordered startup and readiness checks when you invoke it, but it is not a continuously running supervisor. + +## Requirements + +- A Mac with Apple silicon +- macOS 26 or newer +- Apple `container` 1.1.0 or newer +- `openssl` for generating initial secrets +- Local Network access for `container-runtime-linux` when macOS prompts during the first published-container startup + +Install Apple `container` from its [official releases](https://github.com/apple/container/releases), then verify it: + +```bash +container --version +``` + +## Quick Start + +```bash +git clone https://github.com/Wei-Shaw/sub2api.git +cd sub2api/deploy + +# Creates .env with random PostgreSQL, JWT, and TOTP secrets. +./apple-container.sh init + +# Review optional settings before startup. +nano .env + +# Creates volumes/network/containers, waits for dependencies, and starts Sub2API. +./apple-container.sh up + +# Verifies PostgreSQL, Redis, and the application endpoint. +./apple-container.sh status +``` + +Open `http://localhost:8080`. If `ADMIN_PASSWORD` is empty, retrieve the generated password with: + +```bash +./apple-container.sh logs app +``` + +The env file uses literal `KEY=value` syntax. Do not use Compose expressions such as `${VALUE:-default}`, and do not quote values unless the quote characters are part of the intended value. `BIND_HOST` must be an IPv4 address, and `SERVER_PORT` must be between 1025 and 65535. + +## Commands + +```bash +# Start dependencies and recreate the lightweight app container with current IPs. +./apple-container.sh up + +# Also recreate PostgreSQL and Redis containers, preserving their volumes. +./apple-container.sh up --recreate + +# Stop containers while preserving all resources and data. +./apple-container.sh down + +# Restart PostgreSQL, Redis, and Sub2API in dependency order. +./apple-container.sh restart + +# Show resource state and run live health probes. +./apple-container.sh status + +# Follow one service's logs. +./apple-container.sh logs app -f +./apple-container.sh logs postgres -f +./apple-container.sh logs redis -f + +# Pull all configured images for linux/arm64, then recreate containers. +./apple-container.sh pull +./apple-container.sh up --recreate + +# Delete containers and the network, preserving named volumes. +./apple-container.sh destroy --yes + +# Permanently delete the stack and all application/database/cache data. +./apple-container.sh destroy --volumes --yes +``` + +`destroy --volumes` does not remove `.env`, backup files, or pulled images. Delete credentials and backups separately when decommissioning a deployment. Use `container image delete ` only after confirming no other Apple containers use that image. + +After a host reboot or `container system stop`, run `./apple-container.sh up` again. Apple `container` does not automatically restart persisted containers. + +## Configuration + +The script uses `deploy/.env`, the same source file used by Docker Compose. Export `SUB2API_ENV_FILE` to use another file for every command in the current shell: + +```bash +export SUB2API_ENV_FILE=/absolute/path/to/sub2api.env +./apple-container.sh init +./apple-container.sh up +``` + +Apple-specific image overrides are available: + +```dotenv +APPLE_CONTAINER_SUB2API_IMAGE=weishaw/sub2api:latest +APPLE_CONTAINER_POSTGRES_IMAGE=postgres:18-alpine +APPLE_CONTAINER_REDIS_IMAGE=redis:8-alpine +``` + +The normal `up` command recreates the application container, so application environment changes are applied immediately. Use `up --recreate` when changing PostgreSQL or Redis container images or Redis runtime configuration. Persistent data remains in named volumes. + +`POSTGRES_USER`, `POSTGRES_PASSWORD`, and `POSTGRES_DB` are applied only when PostgreSQL initializes an empty data volume. Changing them in `.env` and recreating the container does not change an existing database. Rotate a password with `ALTER ROLE`, and plan explicit migrations for user or database changes. To intentionally initialize a new empty database, first back up the old one and use `destroy --volumes`. + +Apple-specific handling of shared settings: + +| Setting | Apple workflow behavior | +|---|---| +| Application and gateway variables | Passed to Sub2API from `.env` | +| `BIND_HOST`, `SERVER_PORT` | Used for the macOS published port | +| `POSTGRES_USER`, `POSTGRES_PASSWORD`, `POSTGRES_DB` | PostgreSQL first initialization only | +| `REDIS_PASSWORD` | Applied to Redis and Sub2API | +| `DATABASE_PORT`, `REDIS_PORT` | Internal ports are fixed to 5432 and 6379 | +| `POSTGRES_MAX_*`, `REDIS_MAXCLIENTS` | Not currently applied to the database/cache server | + +## Managed Resources + +The script creates only resources carrying the `org.sub2api.stack=apple-container` label: + +| Type | Names | +|---|---| +| Containers | `sub2api-apple`, `sub2api-apple-postgres`, `sub2api-apple-redis` | +| Network | `sub2api-apple` | +| Volumes | `sub2api-apple-data`, `sub2api-apple-postgres-data`, `sub2api-apple-redis-data` | + +The PostgreSQL volume is mounted at `/var/lib/postgresql`, retaining PostgreSQL 18's default child data directory. Sub2API and Redis also store data in child directories below their Apple volume mount points. This is required because Apple named volumes do not have Docker's copy-up and mount-point ownership behavior. + +## Networking + +Apple `container` 1.1 does not provide Compose-style network-scoped service aliases. After PostgreSQL and Redis start, the script reads their current private-network IPv4 addresses from `container inspect`, injects those addresses into a newly created application container, and then starts Sub2API. The script does not modify `~/.config/container/config.toml` or the macOS host resolver. + +All three services attach only to the private `sub2api-apple` network. Only the application publishes a host port; database and Redis ports remain unpublished. + +The application container is intentionally recreated by every `up` and `restart` operation because dependency VM addresses can change after they stop. Application data remains in `sub2api-apple-data`. + +The script checks the published `/health` endpoint from macOS before reporting success. Approve the Local Network prompt on first startup. If the internal probe succeeds but the host-port probe fails with a connection reset, enable Local Network access for `container-runtime-linux`, run `container system stop` followed by `container system start`, and then run `up` again. Runtime upgrades may prompt for permission again. + +## Backup and Upgrade + +Pin image release tags or digests in `.env` before using this workflow for persistent data. Before an application or database image upgrade, create backups while the stack is healthy: + +```bash +umask 077 +mkdir -p backups + +# Logical PostgreSQL backup. +container exec sub2api-apple sh -c \ + 'PGPASSWORD="$DATABASE_PASSWORD" pg_dump -h "$DATABASE_HOST" -U "$DATABASE_USER" "$DATABASE_DBNAME"' \ + > backups/sub2api.sql + +# Application configuration and local files. +container exec sub2api-apple sh -c 'tar -C "$DATA_DIR" -czf - .' \ + > backups/sub2api-data.tar.gz + +./apple-container.sh pull +./apple-container.sh up --recreate +./apple-container.sh status +``` + +Database migrations are forward-only. Keep the previous image reference and both backups until the upgraded stack has been validated; image rollback alone cannot reverse a migrated database. Test restore procedures before relying on this workflow for important data. + +To restore these backups into an existing stack, first ensure the image versions are compatible with the backup, then stop writers and replace both data sets: + +```bash +# Ensure empty/current resources exist, then stop the stack. +./apple-container.sh up +./apple-container.sh down + +# Remove only the app container so a helper can mount its named volume. +container delete sub2api-apple +SUB2API_IMAGE=weishaw/sub2api:latest # Match APPLE_CONTAINER_SUB2API_IMAGE in .env. +container run --rm --name sub2api-apple-data-restore \ + --entrypoint /bin/sh \ + --volume sub2api-apple-data:/restore \ + --volume "$PWD/backups:/backup:ro" \ + "$SUB2API_IMAGE" \ + -c 'rm -rf /restore/data && mkdir -p /restore/data && tar -xzf /backup/sub2api-data.tar.gz -C /restore/data' + +# Restore the logical database while the application is absent. +container start sub2api-apple-postgres +until container exec sub2api-apple-postgres sh -c 'pg_isready -U "$POSTGRES_USER" -d "$POSTGRES_DB"'; do sleep 1; done +container copy backups/sub2api.sql sub2api-apple-postgres:/tmp/sub2api.sql +container exec sub2api-apple-postgres sh -c ' + export PGPASSWORD="$POSTGRES_PASSWORD" + dropdb -h 127.0.0.1 -U "$POSTGRES_USER" --if-exists --force "$POSTGRES_DB" + createdb -h 127.0.0.1 -U "$POSTGRES_USER" "$POSTGRES_DB" + psql -h 127.0.0.1 -U "$POSTGRES_USER" -d "$POSTGRES_DB" -v ON_ERROR_STOP=1 -f /tmp/sub2api.sql + rm /tmp/sub2api.sql +' + +./apple-container.sh up +./apple-container.sh status +``` + +For disaster recovery after deleting the named volumes, run `up` once to create a fresh stack before following the restore sequence. Perform restore drills with non-production data first. + +To upgrade the Apple runtime itself: + +```bash +./apple-container.sh down +container system stop +# Install/update Apple container 1.1.0 or newer. +container system start +./apple-container.sh up +``` + +## Operational Limitations + +- There is no `restart: unless-stopped` equivalent. Run `up` after reboot, or add your own launchd supervisor. +- Health probes run during `up`, `restart`, and `status`; Apple `container` does not continuously schedule them. +- Docker Compose, Testcontainers, Buildx, and tools requiring `/var/run/docker.sock` cannot use this runtime directly. +- Named volume backup and restore must be tested before using this workflow for important data. +- The script targets native `linux/arm64` images. The normal Sub2API release publishes an arm64 variant. +- Runtime environment values, including credentials, are retained in Apple container configuration and are visible to users who can inspect the local runtime. diff --git a/deploy/README.md b/deploy/README.md index dd311721d9..d4fcc133b2 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -1,12 +1,13 @@ # Sub2API Deployment Files -This directory contains files for deploying Sub2API on Linux servers. +This directory contains files for deploying Sub2API on Linux servers and Apple-silicon Macs. ## Deployment Methods | Method | Best For | Setup Wizard | |--------|----------|--------------| | **Docker Compose** | Quick setup, all-in-one | Not needed (auto-setup) | +| **Apple container** | Native local stack on macOS 26 | Not needed (auto-setup) | | **Binary Install** | Production servers, systemd | Web-based wizard | ## Files @@ -16,7 +17,9 @@ This directory contains files for deploying Sub2API on Linux servers. | `docker-compose.yml` | Docker Compose configuration (named volumes) | | `docker-compose.local.yml` | Docker Compose configuration (local directories, easy migration) | | `docker-deploy.sh` | **One-click Docker deployment script (recommended)** | -| `.env.example` | Docker environment variables template | +| `apple-container.sh` | Native Apple `container` lifecycle script | +| `APPLE_CONTAINER.md` | Apple `container` deployment and operations guide | +| `.env.example` | Container environment variables template | | `DOCKER.md` | Docker Hub documentation | | `install.sh` | One-click binary installation script | | `install-datamanagementd.sh` | datamanagementd 一键安装脚本 | @@ -27,6 +30,23 @@ This directory contains files for deploying Sub2API on Linux servers. --- +## Apple container Deployment + +Apple-silicon Macs running macOS 26 can run the complete Sub2API, PostgreSQL, and Redis stack with Apple `container` 1.1.0 or newer: + +```bash +./apple-container.sh init +./apple-container.sh up +./apple-container.sh status +./apple-container.sh logs app -f +``` + +The script uses Apple named volumes, starts dependencies in order, and performs live readiness checks. It does not provide a continuous restart supervisor; run `./apple-container.sh up` after a host reboot. Docker Compose remains the recommended production deployment path. + +See [APPLE_CONTAINER.md](./APPLE_CONTAINER.md) for configuration, upgrades, persistence, networking behavior, and limitations. + +--- + ## Docker Deployment (Recommended) ### Method 1: One-Click Deployment (Recommended) @@ -76,6 +96,7 @@ cd sub2api/deploy # Configure environment cp .env.example .env +chmod 600 .env nano .env # Set POSTGRES_PASSWORD and other required variables # Generate secure secrets (recommended) diff --git a/deploy/apple-container.sh b/deploy/apple-container.sh new file mode 100755 index 0000000000..5de4d5cab1 --- /dev/null +++ b/deploy/apple-container.sh @@ -0,0 +1,926 @@ +#!/usr/bin/env bash + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENV_FILE="${SUB2API_ENV_FILE:-${SCRIPT_DIR}/.env}" + +STACK_LABEL_KEY="org.sub2api.stack" +STACK_LABEL_VALUE="apple-container" +NETWORK_NAME="sub2api-apple" +APP_CONTAINER="sub2api-apple" +POSTGRES_CONTAINER="sub2api-apple-postgres" +REDIS_CONTAINER="sub2api-apple-redis" +APP_VOLUME="sub2api-apple-data" +POSTGRES_VOLUME="sub2api-apple-postgres-data" +REDIS_VOLUME="sub2api-apple-redis-data" +PLATFORM="linux/arm64" + +TEMP_DIR="" +LOCK_DIR="${TMPDIR:-/tmp}/sub2api-apple-container.lock" +LOCK_ACQUIRED=false + +APP_IMAGE="" +POSTGRES_IMAGE="" +REDIS_IMAGE="" +BIND_HOST="" +HOST_PORT="" +ACCESS_HOST="" +POSTGRES_USER="" +POSTGRES_PASSWORD="" +POSTGRES_DB="" +REDIS_PASSWORD="" +TZ_VALUE="" +POSTGRES_ADDRESS="" +REDIS_ADDRESS="" +APP_ENV_FILE="" +POSTGRES_ENV_FILE="" +POSTGRES_PROBE_ENV_FILE="" +REDIS_ENV_FILE="" + +info() { + printf '[INFO] %s\n' "$*" +} + +warn() { + printf '[WARN] %s\n' "$*" >&2 +} + +die() { + printf '[ERROR] %s\n' "$*" >&2 + exit 1 +} + +usage() { + cat <<'EOF' +Usage: ./apple-container.sh [options] + +Commands: + init Create .env and generate required secrets + up [--recreate] Create and start the complete Sub2API stack + down Stop the stack and preserve all data + restart Restart the stack in dependency order + status Show container and workload health + logs [-f] Show logs for app, postgres, or redis + pull Pull all stack images for linux/arm64 + destroy [options] Delete stack containers and network + +Destroy options: + --volumes Also delete all persistent data volumes + --yes Skip the confirmation prompt + +Environment: + SUB2API_ENV_FILE Path to the deployment env file (default: deploy/.env) +EOF +} + +cleanup() { + local exit_code=$? + + if [[ -n "${TEMP_DIR}" && -d "${TEMP_DIR}" ]]; then + rm -rf "${TEMP_DIR}" + fi + if [[ "${LOCK_ACQUIRED}" == true && -d "${LOCK_DIR}" ]]; then + rm -f "${LOCK_DIR}/pid" + rmdir "${LOCK_DIR}" 2>/dev/null || true + fi + + exit "${exit_code}" +} + +acquire_lock() { + if ! mkdir "${LOCK_DIR}" 2>/dev/null; then + local owner_pid="" + if [[ -f "${LOCK_DIR}/pid" ]]; then + owner_pid="$(<"${LOCK_DIR}/pid")" + fi + if [[ "${owner_pid}" =~ ^[0-9]+$ ]] && ! kill -0 "${owner_pid}" 2>/dev/null; then + rm -rf "${LOCK_DIR}" + mkdir "${LOCK_DIR}" || die "Failed to reclaim stale operation lock." + else + die "Another Sub2API Apple container operation is already running." + fi + fi + printf '%s\n' "$$" >"${LOCK_DIR}/pid" + LOCK_ACQUIRED=true + trap cleanup EXIT + trap 'exit 130' INT + trap 'exit 143' TERM + trap 'exit 129' HUP +} + +require_command() { + command -v "$1" >/dev/null 2>&1 || die "Required command not found: $1" +} + +require_container_version() { + local version_output major minor + + require_command container + require_command plutil + version_output="$(container --version)" + if [[ ! "${version_output}" =~ ([0-9]+)\.([0-9]+)\.([0-9]+) ]]; then + die "Unable to parse Apple container version: ${version_output}" + fi + + major="${BASH_REMATCH[1]}" + minor="${BASH_REMATCH[2]}" + if (( major < 1 || (major == 1 && minor < 1) )); then + die "Apple container 1.1.0 or newer is required; found ${version_output}." + fi +} + +system_is_running() { + container system status >/dev/null 2>&1 +} + +start_system() { + if ! system_is_running; then + info "Starting Apple container services..." + container system start --enable-kernel-install + fi +} + +list_resource_ids() { + case "$1" in + container) container list --all --quiet ;; + network) container network list --quiet ;; + volume) container volume list --quiet ;; + *) die "Unknown resource type: $1" ;; + esac +} + +resource_exists() { + local resource_type=$1 + local resource_name=$2 + local output line + + if ! output="$(list_resource_ids "${resource_type}")"; then + die "Failed to list Apple container ${resource_type} resources." + fi + + while IFS= read -r line; do + if [[ "${line}" == "${resource_name}" ]]; then + return 0 + fi + done <<<"${output}" + + return 1 +} + +inspect_resource() { + case "$1" in + container) container inspect "$2" ;; + network) container network inspect "$2" ;; + volume) container volume inspect "$2" ;; + *) die "Unknown resource type: $1" ;; + esac +} + +assert_resource_owned() { + local resource_type=$1 + local resource_name=$2 + local inspection compact + + inspection="$(inspect_resource "${resource_type}" "${resource_name}" | \ + plutil -extract 0.configuration.labels json -o - -)" || \ + die "Failed to inspect ${resource_type} ${resource_name}." + compact="$(printf '%s' "${inspection}" | tr -d '[:space:]')" + if [[ "${compact}" != *"\"${STACK_LABEL_KEY}\":\"${STACK_LABEL_VALUE}\""* ]]; then + die "Refusing to manage existing ${resource_type} '${resource_name}' because it is not owned by this stack." + fi +} + +preflight_stack_ownership() { + local resource_name + + for resource_name in "${APP_CONTAINER}" "${REDIS_CONTAINER}" "${POSTGRES_CONTAINER}"; do + if resource_exists container "${resource_name}"; then + assert_resource_owned container "${resource_name}" + fi + done + if resource_exists network "${NETWORK_NAME}"; then + assert_resource_owned network "${NETWORK_NAME}" + fi + for resource_name in "${APP_VOLUME}" "${REDIS_VOLUME}" "${POSTGRES_VOLUME}"; do + if resource_exists volume "${resource_name}"; then + assert_resource_owned volume "${resource_name}" + fi + done +} + +ensure_network() { + if resource_exists network "${NETWORK_NAME}"; then + assert_resource_owned network "${NETWORK_NAME}" + return + fi + + info "Creating network ${NETWORK_NAME}..." + container network create \ + --label "${STACK_LABEL_KEY}=${STACK_LABEL_VALUE}" \ + "${NETWORK_NAME}" >/dev/null +} + +ensure_volume() { + local volume_name=$1 + + if resource_exists volume "${volume_name}"; then + assert_resource_owned volume "${volume_name}" + return + fi + + info "Creating volume ${volume_name}..." + container volume create \ + --label "${STACK_LABEL_KEY}=${STACK_LABEL_VALUE}" \ + "${volume_name}" >/dev/null +} + +ensure_image_available() { + local image=$1 + + if container image inspect "${image}" >/dev/null 2>&1; then + return + fi + info "Pulling ${image}..." + container image pull --platform "${PLATFORM}" "${image}" +} + +container_is_running() { + local container_name=$1 + local output line + + output="$(container list --quiet)" || die "Failed to list running Apple containers." + while IFS= read -r line; do + if [[ "${line}" == "${container_name}" ]]; then + return 0 + fi + done <<<"${output}" + + return 1 +} + +ensure_system() { + require_container_version + require_command curl + start_system +} + +container_ipv4_address() { + local container_name=$1 + local address + + address="$(container inspect "${container_name}" | \ + plutil -extract 0.status.networks.0.ipv4Address raw -o - -)" || \ + die "Unable to read the network address for ${container_name}." + address="${address%%/*}" + [[ "${address}" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]] || \ + die "Apple container returned an invalid IPv4 address for ${container_name}: ${address}" + printf '%s\n' "${address}" +} + +read_env_value() { + local key=$1 + local fallback=${2-} + + awk -v wanted="${key}" -v fallback="${fallback}" ' + BEGIN { found = 0 } + /^[[:space:]]*#/ || /^[[:space:]]*$/ { next } + { + separator = index($0, "=") + if (separator == 0) { next } + key = substr($0, 1, separator - 1) + if (key == wanted) { + value = substr($0, separator + 1) + sub(/\r$/, "", value) + found = 1 + } + } + END { + if (found) { print value } + else { print fallback } + } + ' "${ENV_FILE}" +} + +replace_env_value() { + local key=$1 + local value=$2 + local target_file=${3:-${ENV_FILE}} + local temp_file="${target_file}.tmp.$$" + + awk -v wanted="${key}" -v replacement="${value}" ' + BEGIN { replaced = 0 } + { + separator = index($0, "=") + key = separator == 0 ? "" : substr($0, 1, separator - 1) + if (key == wanted) { + if (!replaced) { print wanted "=" replacement } + replaced = 1 + next + } + print + } + END { + if (!replaced) { print wanted "=" replacement } + } + ' "${target_file}" >"${temp_file}" + chmod 600 "${temp_file}" + mv "${temp_file}" "${target_file}" +} + +generate_secret() { + openssl rand -hex 32 +} + +cmd_init() { + local env_dir temp_file postgres_secret jwt_secret totp_secret + + require_command openssl + + if [[ -e "${ENV_FILE}" ]]; then + die "Environment file already exists: ${ENV_FILE}" + fi + + postgres_secret="$(generate_secret)" || die "Failed to generate PostgreSQL password." + jwt_secret="$(generate_secret)" || die "Failed to generate JWT secret." + totp_secret="$(generate_secret)" || die "Failed to generate TOTP encryption key." + [[ -n "${postgres_secret}" && -n "${jwt_secret}" && -n "${totp_secret}" ]] || \ + die "Secret generation returned an empty value." + + env_dir="$(dirname "${ENV_FILE}")" + temp_file="${ENV_FILE}.init.tmp.$$" + mkdir -p "${env_dir}" + cp "${SCRIPT_DIR}/.env.example" "${temp_file}" + chmod 600 "${temp_file}" + replace_env_value POSTGRES_PASSWORD "${postgres_secret}" "${temp_file}" + replace_env_value JWT_SECRET "${jwt_secret}" "${temp_file}" + replace_env_value TOTP_ENCRYPTION_KEY "${totp_secret}" "${temp_file}" + mv "${temp_file}" "${ENV_FILE}" + + info "Created ${ENV_FILE} with generated secrets." + info "Review the file, then run: SUB2API_ENV_FILE='${ENV_FILE}' ${SCRIPT_DIR}/apple-container.sh up" +} + +validate_port() { + local port=$1 + local decimal_port + + [[ "${port}" =~ ^[0-9]+$ ]] || die "SERVER_PORT must be numeric: ${port}" + decimal_port=$((10#${port})) + (( decimal_port >= 1025 && decimal_port <= 65535 )) || \ + die "SERVER_PORT must be between 1025 and 65535 for Apple container port forwarding." +} + +validate_ipv4_address() { + local address=$1 + local first second third fourth extra octet + + IFS=. read -r first second third fourth extra <<<"${address}" + [[ -n "${first}" && -n "${second}" && -n "${third}" && -n "${fourth}" && -z "${extra}" ]] || \ + die "BIND_HOST must be a valid IPv4 address: ${address}" + for octet in "${first}" "${second}" "${third}" "${fourth}"; do + [[ "${octet}" =~ ^[0-9]+$ ]] || die "BIND_HOST must be a valid IPv4 address: ${address}" + (( 10#${octet} <= 255 )) || die "BIND_HOST must be a valid IPv4 address: ${address}" + done +} + +validate_env_file_security() { + local owner mode permissions + + [[ -f "${ENV_FILE}" ]] || die "Environment file not found: ${ENV_FILE}. Run '$0 init' first." + owner="$(stat -f '%u' "${ENV_FILE}")" || die "Unable to read owner for ${ENV_FILE}." + mode="$(stat -f '%Lp' "${ENV_FILE}")" || die "Unable to read permissions for ${ENV_FILE}." + [[ "${owner}" == "${EUID}" ]] || die "Environment file must be owned by the current user: ${ENV_FILE}" + [[ "${mode}" =~ ^[0-7]+$ ]] || die "Unable to parse permissions for ${ENV_FILE}: ${mode}" + permissions=$((8#${mode})) + (( (permissions & 077) == 0 )) || \ + die "Environment file must not be readable by group or others. Run: chmod 600 '${ENV_FILE}'" +} + +prepare_environment() { + validate_env_file_security + + APP_IMAGE="$(read_env_value APPLE_CONTAINER_SUB2API_IMAGE weishaw/sub2api:latest)" + POSTGRES_IMAGE="$(read_env_value APPLE_CONTAINER_POSTGRES_IMAGE postgres:18-alpine)" + REDIS_IMAGE="$(read_env_value APPLE_CONTAINER_REDIS_IMAGE redis:8-alpine)" + BIND_HOST="$(read_env_value BIND_HOST 0.0.0.0)" + HOST_PORT="$(read_env_value SERVER_PORT 8080)" + POSTGRES_USER="$(read_env_value POSTGRES_USER sub2api)" + POSTGRES_PASSWORD="$(read_env_value POSTGRES_PASSWORD)" + POSTGRES_DB="$(read_env_value POSTGRES_DB sub2api)" + REDIS_PASSWORD="$(read_env_value REDIS_PASSWORD)" + TZ_VALUE="$(read_env_value TZ Asia/Shanghai)" + + [[ -n "${BIND_HOST}" ]] || die "BIND_HOST must not be empty." + validate_ipv4_address "${BIND_HOST}" + validate_port "${HOST_PORT}" + if [[ "${BIND_HOST}" == "0.0.0.0" ]]; then + ACCESS_HOST="127.0.0.1" + else + ACCESS_HOST="${BIND_HOST}" + fi + [[ -n "${POSTGRES_USER}" ]] || die "POSTGRES_USER must not be empty." + [[ -n "${POSTGRES_DB}" ]] || die "POSTGRES_DB must not be empty." + if [[ -z "${POSTGRES_PASSWORD}" || "${POSTGRES_PASSWORD}" == "change_this_secure_password" ]]; then + die "Set a secure POSTGRES_PASSWORD in ${ENV_FILE}." + fi + + TEMP_DIR="$(mktemp -d "${TMPDIR:-/tmp}/sub2api-apple.XXXXXX")" + APP_ENV_FILE="${TEMP_DIR}/app.env" + POSTGRES_ENV_FILE="${TEMP_DIR}/postgres.env" + POSTGRES_PROBE_ENV_FILE="${TEMP_DIR}/postgres-probe.env" + REDIS_ENV_FILE="${TEMP_DIR}/redis.env" + + cat >"${POSTGRES_ENV_FILE}" <"${POSTGRES_PROBE_ENV_FILE}" <"${REDIS_ENV_FILE}" <>"${REDIS_ENV_FILE}" + fi + + chmod 600 "${POSTGRES_ENV_FILE}" "${POSTGRES_PROBE_ENV_FILE}" "${REDIS_ENV_FILE}" +} + +prepare_app_environment() { + [[ -n "${POSTGRES_ADDRESS}" && -n "${REDIS_ADDRESS}" ]] || \ + die "Dependency network addresses are not available." + + cp "${ENV_FILE}" "${APP_ENV_FILE}" + cat >>"${APP_ENV_FILE}" </dev/null +} + +create_redis_container() { + info "Creating Redis container..." + container create \ + --name "${REDIS_CONTAINER}" \ + --label "${STACK_LABEL_KEY}=${STACK_LABEL_VALUE}" \ + --network "${NETWORK_NAME}" \ + --platform "${PLATFORM}" \ + --ulimit nofile=100000:100000 \ + --env-file "${REDIS_ENV_FILE}" \ + --volume "${REDIS_VOLUME}:/var/lib/redis" \ + "${REDIS_IMAGE}" \ + sh -c 'set -e; mkdir -p /var/lib/redis/data; chown redis:redis /var/lib/redis/data; exec /usr/local/bin/docker-entrypoint.sh redis-server --dir /var/lib/redis/data --save 60 1 --appendonly yes --appendfsync everysec ${REDIS_PASSWORD:+--requirepass "$REDIS_PASSWORD"}' \ + >/dev/null +} + +create_app_container() { + info "Creating Sub2API container..." + container create \ + --name "${APP_CONTAINER}" \ + --label "${STACK_LABEL_KEY}=${STACK_LABEL_VALUE}" \ + --network "${NETWORK_NAME}" \ + --platform "${PLATFORM}" \ + --ulimit nofile=100000:100000 \ + --publish "${BIND_HOST}:${HOST_PORT}:8080/tcp" \ + --env-file "${APP_ENV_FILE}" \ + --volume "${APP_VOLUME}:/app/storage" \ + --entrypoint /bin/sh \ + "${APP_IMAGE}" \ + -c 'set -e; mkdir -p "$DATA_DIR"; chown -R sub2api:sub2api "$DATA_DIR"; exec su-exec sub2api /app/sub2api' \ + >/dev/null +} + +ensure_container() { + local container_name=$1 + local create_function=$2 + + if resource_exists container "${container_name}"; then + assert_resource_owned container "${container_name}" + return + fi + + "${create_function}" +} + +start_container_if_needed() { + local container_name=$1 + + if container_is_running "${container_name}"; then + return + fi + + info "Starting ${container_name}..." + container start "${container_name}" >/dev/null +} + +stop_container_if_running() { + local container_name=$1 + + if ! resource_exists container "${container_name}"; then + return + fi + assert_resource_owned container "${container_name}" + if container_is_running "${container_name}"; then + info "Stopping ${container_name}..." + container stop --time 30 "${container_name}" >/dev/null + fi +} + +delete_container_if_present() { + local container_name=$1 + + if ! resource_exists container "${container_name}"; then + return + fi + assert_resource_owned container "${container_name}" + if container_is_running "${container_name}"; then + container stop --time 30 "${container_name}" >/dev/null + fi + info "Deleting ${container_name}..." + container delete "${container_name}" >/dev/null +} + +wait_for_probe() { + local description=$1 + local attempts=$2 + shift 2 + + local attempt + for ((attempt = 1; attempt <= attempts; attempt++)); do + if "$@" >/dev/null 2>&1; then + info "${description} is ready." + return 0 + fi + sleep 1 + done + + return 1 +} + +probe_postgres() { + container exec --env-file "${POSTGRES_PROBE_ENV_FILE}" \ + "${POSTGRES_CONTAINER}" \ + psql -h 127.0.0.1 -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" \ + -v ON_ERROR_STOP=1 -tAc 'SELECT 1' +} + +probe_redis() { + container exec --env-file "${REDIS_ENV_FILE}" \ + "${REDIS_CONTAINER}" \ + redis-cli ping +} + +probe_app() { + container exec "${APP_CONTAINER}" \ + wget -q -T 5 -O /dev/null http://localhost:8080/health +} + +probe_host_app() { + curl --fail --silent --show-error --max-time 5 \ + "http://${ACCESS_HOST}:${HOST_PORT}/health" +} + +show_failure_logs() { + local container_name=$1 + + warn "Last logs from ${container_name}:" + container logs -n 50 "${container_name}" >&2 || true +} + +start_dependencies() { + start_container_if_needed "${POSTGRES_CONTAINER}" + if ! wait_for_probe "PostgreSQL" 90 probe_postgres; then + show_failure_logs "${POSTGRES_CONTAINER}" + die "PostgreSQL did not become ready." + fi + + start_container_if_needed "${REDIS_CONTAINER}" + if ! wait_for_probe "Redis" 60 probe_redis; then + show_failure_logs "${REDIS_CONTAINER}" + die "Redis did not become ready." + fi +} + +start_app() { + start_container_if_needed "${APP_CONTAINER}" + if ! wait_for_probe "Sub2API" 180 probe_app; then + show_failure_logs "${APP_CONTAINER}" + die "Sub2API did not become ready." + fi + if ! wait_for_probe "Sub2API host port" 15 probe_host_app; then + die "Host port forwarding failed. In System Settings > Privacy & Security > Local Network, allow container-runtime-linux; restart Apple container services; then run 'apple-container.sh up' again." + fi +} + +cmd_up() { + local recreate=false + + if [[ $# -gt 1 || ($# -eq 1 && "${1-}" != "--recreate") ]]; then + usage + exit 2 + fi + if [[ $# -eq 1 ]]; then + recreate=true + fi + + ensure_system + prepare_environment + preflight_stack_ownership + ensure_network + ensure_volume "${APP_VOLUME}" + ensure_volume "${POSTGRES_VOLUME}" + ensure_volume "${REDIS_VOLUME}" + ensure_image_available "${APP_IMAGE}" + ensure_image_available "${POSTGRES_IMAGE}" + ensure_image_available "${REDIS_IMAGE}" + + if [[ "${recreate}" == true ]]; then + delete_container_if_present "${APP_CONTAINER}" + delete_container_if_present "${REDIS_CONTAINER}" + delete_container_if_present "${POSTGRES_CONTAINER}" + fi + + ensure_container "${POSTGRES_CONTAINER}" create_postgres_container + ensure_container "${REDIS_CONTAINER}" create_redis_container + start_dependencies + POSTGRES_ADDRESS="$(container_ipv4_address "${POSTGRES_CONTAINER}")" + REDIS_ADDRESS="$(container_ipv4_address "${REDIS_CONTAINER}")" + prepare_app_environment + # The dependency IPs may change whenever their lightweight VMs restart. + delete_container_if_present "${APP_CONTAINER}" + create_app_container + start_app + + info "Sub2API is available at http://${ACCESS_HOST}:${HOST_PORT}" +} + +cmd_down() { + require_container_version + if ! system_is_running; then + info "Apple container services are already stopped." + return + fi + preflight_stack_ownership + stop_container_if_running "${APP_CONTAINER}" + stop_container_if_running "${REDIS_CONTAINER}" + stop_container_if_running "${POSTGRES_CONTAINER}" + info "Sub2API stack stopped; persistent volumes were preserved." +} + +cmd_restart() { + cmd_down + cmd_up +} + +print_container_status() { + local service=$1 + local container_name=$2 + + if ! resource_exists container "${container_name}"; then + printf '%-12s %s\n' "${service}" "missing" + elif container_is_running "${container_name}"; then + printf '%-12s %s\n' "${service}" "running" + else + printf '%-12s %s\n' "${service}" "stopped" + fi +} + +cmd_status() { + local failed=0 + + require_container_version + if ! system_is_running; then + printf '%-12s %s\n' "system" "stopped" + return 1 + fi + + printf '%-12s %s\n' "system" "running" + preflight_stack_ownership + print_container_status app "${APP_CONTAINER}" + print_container_status postgres "${POSTGRES_CONTAINER}" + print_container_status redis "${REDIS_CONTAINER}" + + if [[ -f "${ENV_FILE}" ]]; then + prepare_environment + if container_is_running "${POSTGRES_CONTAINER}" && probe_postgres >/dev/null 2>&1; then + printf '%-12s %s\n' "postgres" "healthy" + else + printf '%-12s %s\n' "postgres" "unhealthy" + failed=1 + fi + if container_is_running "${REDIS_CONTAINER}" && probe_redis >/dev/null 2>&1; then + printf '%-12s %s\n' "redis" "healthy" + else + printf '%-12s %s\n' "redis" "unhealthy" + failed=1 + fi + if container_is_running "${APP_CONTAINER}" && probe_app >/dev/null 2>&1; then + printf '%-12s %s\n' "app" "healthy" + else + printf '%-12s %s\n' "app" "unhealthy" + failed=1 + fi + if container_is_running "${APP_CONTAINER}" && probe_host_app >/dev/null 2>&1; then + printf '%-12s %s\n' "host-port" "healthy" + else + printf '%-12s %s\n' "host-port" "unhealthy" + failed=1 + fi + else + warn "Health probes require ${ENV_FILE}." + failed=1 + fi + + return "${failed}" +} + +cmd_logs() { + local service=${1-} + local follow=${2-} + local container_name + + [[ $# -ge 1 && $# -le 2 ]] || { usage; exit 2; } + if [[ -n "${follow}" && "${follow}" != "-f" && "${follow}" != "--follow" ]]; then + usage + exit 2 + fi + + case "${service}" in + app|sub2api) container_name="${APP_CONTAINER}" ;; + postgres) container_name="${POSTGRES_CONTAINER}" ;; + redis) container_name="${REDIS_CONTAINER}" ;; + *) die "Unknown service '${service}'. Use app, postgres, or redis." ;; + esac + + require_container_version + system_is_running || die "Apple container services are not running." + resource_exists container "${container_name}" || die "Container not found: ${container_name}" + assert_resource_owned container "${container_name}" + if [[ -n "${follow}" ]]; then + container logs --follow "${container_name}" + else + container logs "${container_name}" + fi +} + +cmd_pull() { + ensure_system + prepare_environment + info "Pulling ${APP_IMAGE}..." + container image pull --platform "${PLATFORM}" "${APP_IMAGE}" + info "Pulling ${POSTGRES_IMAGE}..." + container image pull --platform "${PLATFORM}" "${POSTGRES_IMAGE}" + info "Pulling ${REDIS_IMAGE}..." + container image pull --platform "${PLATFORM}" "${REDIS_IMAGE}" +} + +confirm_destroy() { + local include_volumes=$1 + local answer + + if [[ "${include_volumes}" == true ]]; then + printf 'Delete the Sub2API stack and all persistent data? [y/N] ' + else + printf 'Delete the Sub2API containers and network, preserving volumes? [y/N] ' + fi + read -r answer + [[ "${answer}" == "y" || "${answer}" == "Y" ]] +} + +delete_volume_if_present() { + local volume_name=$1 + + if resource_exists volume "${volume_name}"; then + assert_resource_owned volume "${volume_name}" + info "Deleting volume ${volume_name}..." + container volume delete "${volume_name}" >/dev/null + fi +} + +cmd_destroy() { + local include_volumes=false + local assume_yes=false + local argument + + for argument in "$@"; do + case "${argument}" in + --volumes) include_volumes=true ;; + --yes) assume_yes=true ;; + *) usage; exit 2 ;; + esac + done + + require_container_version + start_system + preflight_stack_ownership + if [[ "${assume_yes}" != true ]] && ! confirm_destroy "${include_volumes}"; then + info "Cancelled." + return + fi + + delete_container_if_present "${APP_CONTAINER}" + delete_container_if_present "${REDIS_CONTAINER}" + delete_container_if_present "${POSTGRES_CONTAINER}" + + if resource_exists network "${NETWORK_NAME}"; then + assert_resource_owned network "${NETWORK_NAME}" + info "Deleting network ${NETWORK_NAME}..." + container network delete "${NETWORK_NAME}" >/dev/null + fi + + if [[ "${include_volumes}" == true ]]; then + delete_volume_if_present "${APP_VOLUME}" + delete_volume_if_present "${REDIS_VOLUME}" + delete_volume_if_present "${POSTGRES_VOLUME}" + info "Sub2API stack and persistent data deleted." + else + info "Sub2API stack deleted; persistent volumes were preserved." + fi +} + +main() { + local command=${1-} + if [[ $# -gt 0 ]]; then + shift + fi + + case "${command}" in + init) + [[ $# -eq 0 ]] || { usage; exit 2; } + acquire_lock + cmd_init + ;; + up) + acquire_lock + cmd_up "$@" + ;; + down) + [[ $# -eq 0 ]] || { usage; exit 2; } + acquire_lock + cmd_down + ;; + restart) + [[ $# -eq 0 ]] || { usage; exit 2; } + acquire_lock + cmd_restart + ;; + status) + [[ $# -eq 0 ]] || { usage; exit 2; } + trap cleanup EXIT + cmd_status + ;; + logs) + cmd_logs "$@" + ;; + pull) + [[ $# -eq 0 ]] || { usage; exit 2; } + acquire_lock + cmd_pull + ;; + destroy) + acquire_lock + cmd_destroy "$@" + ;; + help|-h|--help) + usage + ;; + *) + usage + exit 2 + ;; + esac +} + +main "$@" diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index eff4bfb598..4cf759497f 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -20,6 +20,9 @@ server: # Mode: "debug" for development, "release" for production # 运行模式："debug" 用于开发，"release" 用于生产环境 mode: "release" + # Return Server-Timing for authenticated requests made by the Admin web UI + # 为管理端 Web 页面发出的已认证请求返回 Server-Timing + enable_server_timing: false # Frontend base URL used to generate external links in emails (e.g. password reset) # 用于生成邮件中的外部链接（例如：重置密码链接）的前端基础地址 # Example: "https://example.com" @@ -254,6 +257,10 @@ gateway: # ingress 默认模式：off|ctx_pool|passthrough|http_bridge（仅 mode_router_v2_enabled=true 生效） # 兼容旧值：shared/dedicated 会按 ctx_pool 处理。 ingress_mode_default: ctx_pool + # Close a client WebSocket that stays idle between completed turns (seconds). Set 0 to disable. + ingress_inter_turn_idle_timeout_seconds: 300 + # Limit live client WebSocket ingress sessions per API key across all instances. Set 0 to disable. + max_ingress_connections_per_api_key: 64 # 全局总开关，默认 true；关闭时所有请求保持原有 HTTP/SSE 路由 enabled: true # 按账号类型细分开关 @@ -383,6 +390,9 @@ gateway: # Image stream keepalive interval (seconds), 0=disable; independent from ordinary text streams # 图片流式 keepalive 间隔（秒），0=禁用；独立于普通文本流式 image_stream_keepalive_interval: 10 + # Non-streaming Images JSON keepalive interval (seconds), 0=disable; commits HTTP 200 after the first heartbeat + # 图片非流式 JSON keepalive 间隔（秒），0=禁用；首个心跳后 HTTP 状态会固化为 200 + image_nonstream_keepalive_interval: 0 # Image generation independent concurrency limiter (process-local, default disabled) # 图片生成独立并发限制（进程级，默认关闭；多实例总上限约为实例数×该值） image_concurrency: diff --git a/deploy/docker-compose.dev.yml b/deploy/docker-compose.dev.yml index 6f5b3f56f3..43f5dd3f60 100644 --- a/deploy/docker-compose.dev.yml +++ b/deploy/docker-compose.dev.yml @@ -26,6 +26,7 @@ services: - SERVER_HOST=0.0.0.0 - SERVER_PORT=8080 - SERVER_MODE=debug + - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} - DATABASE_HOST=postgres - DATABASE_PORT=5432 diff --git a/deploy/docker-compose.local.yml b/deploy/docker-compose.local.yml index 042752e857..5fb161603b 100644 --- a/deploy/docker-compose.local.yml +++ b/deploy/docker-compose.local.yml @@ -51,6 +51,7 @@ services: - SERVER_HOST=0.0.0.0 - SERVER_PORT=8080 - SERVER_MODE=${SERVER_MODE:-release} + - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} # ======================================================================= diff --git a/deploy/docker-compose.standalone.yml b/deploy/docker-compose.standalone.yml index 2e1d335624..40ed4751d6 100644 --- a/deploy/docker-compose.standalone.yml +++ b/deploy/docker-compose.standalone.yml @@ -37,6 +37,7 @@ services: - SERVER_HOST=0.0.0.0 - SERVER_PORT=8080 - SERVER_MODE=${SERVER_MODE:-release} + - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} # ======================================================================= diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 22713c59aa..6aecdcfa5a 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -47,6 +47,7 @@ services: - SERVER_HOST=0.0.0.0 - SERVER_PORT=8080 - SERVER_MODE=${SERVER_MODE:-release} + - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} # ======================================================================= diff --git a/deploy/tests/apple-container-test.sh b/deploy/tests/apple-container-test.sh new file mode 100755 index 0000000000..a12582104f --- /dev/null +++ b/deploy/tests/apple-container-test.sh @@ -0,0 +1,78 @@ +#!/bin/bash + +set -euo pipefail + +TEST_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +DEPLOY_DIR="$(cd "${TEST_DIR}/.." && pwd)" +SCRIPT="${DEPLOY_DIR}/apple-container.sh" +TEST_ROOT="$(mktemp -d "${TMPDIR:-/tmp}/sub2api-apple-test.XXXXXX")" +STATE_DIR="${TEST_ROOT}/state" +ENV_FILE="${TEST_ROOT}/sub2api.env" + +cleanup() { + rm -rf "${TEST_ROOT}" +} +trap cleanup EXIT + +fail() { + printf 'FAIL: %s\n' "$*" >&2 + exit 1 +} + +assert_exists() { + [[ -e "$1" ]] || fail "Expected path to exist: $1" +} + +assert_missing() { + [[ ! -e "$1" ]] || fail "Expected path to be absent: $1" +} + +export FAKE_CONTAINER_STATE="${STATE_DIR}" +export PATH="${TEST_DIR}/fixtures/bin:${PATH}" +export SUB2API_ENV_FILE="${ENV_FILE}" + +mkdir -p "${STATE_DIR}" + +"${SCRIPT}" init +[[ "$(stat -f '%Lp' "${ENV_FILE}")" == "600" ]] || fail "init did not create a mode-600 env file" +grep -q '^POSTGRES_PASSWORD=change_this_secure_password$' "${ENV_FILE}" && fail "init retained the placeholder password" + +chmod 644 "${ENV_FILE}" +if "${SCRIPT}" up >/dev/null 2>&1; then + fail "up accepted an insecure env file" +fi +chmod 600 "${ENV_FILE}" + +"${SCRIPT}" up +assert_exists "${STATE_DIR}/containers/sub2api-apple" +assert_exists "${STATE_DIR}/containers/sub2api-apple-postgres" +assert_exists "${STATE_DIR}/containers/sub2api-apple-redis" +assert_exists "${STATE_DIR}/running/sub2api-apple" +"${SCRIPT}" status >/dev/null + +"${SCRIPT}" up --recreate +assert_exists "${STATE_DIR}/running/sub2api-apple" +"${SCRIPT}" down +assert_missing "${STATE_DIR}/running/sub2api-apple" +assert_missing "${STATE_DIR}/running/sub2api-apple-postgres" +assert_missing "${STATE_DIR}/running/sub2api-apple-redis" + +"${SCRIPT}" destroy --yes +assert_missing "${STATE_DIR}/containers/sub2api-apple" +assert_missing "${STATE_DIR}/networks/sub2api-apple" +assert_exists "${STATE_DIR}/volumes/sub2api-apple-data" + +"${SCRIPT}" up +"${SCRIPT}" destroy --volumes --yes +assert_missing "${STATE_DIR}/volumes/sub2api-apple-data" +assert_missing "${STATE_DIR}/volumes/sub2api-apple-postgres-data" +assert_missing "${STATE_DIR}/volumes/sub2api-apple-redis-data" + +touch "${STATE_DIR}/system-running" +touch "${STATE_DIR}/containers/sub2api-apple" +touch "${STATE_DIR}/unowned/container/sub2api-apple" +if "${SCRIPT}" status >/dev/null 2>&1; then + fail "status accepted an unowned same-name container" +fi + +printf 'Apple container lifecycle tests passed.\n' diff --git a/deploy/tests/fixtures/bin/container b/deploy/tests/fixtures/bin/container new file mode 100755 index 0000000000..a864111f27 --- /dev/null +++ b/deploy/tests/fixtures/bin/container @@ -0,0 +1,164 @@ +#!/bin/bash + +set -eu + +STATE_DIR="${FAKE_CONTAINER_STATE:?FAKE_CONTAINER_STATE is required}" +mkdir -p \ + "${STATE_DIR}/containers" \ + "${STATE_DIR}/running" \ + "${STATE_DIR}/networks" \ + "${STATE_DIR}/volumes" \ + "${STATE_DIR}/unowned/container" \ + "${STATE_DIR}/unowned/network" \ + "${STATE_DIR}/unowned/volume" + +list_names() { + local directory=$1 + local path + + for path in "${directory}"/*; do + [[ -e "${path}" ]] || continue + basename "${path}" + done +} + +last_argument() { + local value="" + + for value in "$@"; do :; done + printf '%s\n' "${value}" +} + +inspect_resource() { + local resource_type=$1 + local resource_name=$2 + local label_value="apple-container" + local address="192.168.65.4/24" + + if [[ -e "${STATE_DIR}/unowned/${resource_type}/${resource_name}" ]]; then + label_value="other" + fi + case "${resource_name}" in + sub2api-apple-postgres) address="192.168.65.2/24" ;; + sub2api-apple-redis) address="192.168.65.3/24" ;; + esac + + printf '[{"configuration":{"labels":{"org.sub2api.stack":"%s"}},"status":{"networks":[{"ipv4Address":"%s"}]}}]\n' \ + "${label_value}" "${address}" +} + +command=${1-} +if [[ $# -gt 0 ]]; then shift; fi + +case "${command}" in + --version) + echo "container CLI version 1.1.0 (build: release, commit: fake)" + ;; + system) + subcommand=${1-} + case "${subcommand}" in + status) [[ -e "${STATE_DIR}/system-running" ]] ;; + start) touch "${STATE_DIR}/system-running" ;; + stop) rm -f "${STATE_DIR}/system-running" "${STATE_DIR}/running"/* ;; + *) exit 1 ;; + esac + ;; + list) + include_all=false + for argument in "$@"; do + [[ "${argument}" == "--all" || "${argument}" == "-a" ]] && include_all=true + done + if [[ "${include_all}" == true ]]; then + list_names "${STATE_DIR}/containers" + else + list_names "${STATE_DIR}/running" + fi + ;; + network) + subcommand=${1-} + shift || true + case "${subcommand}" in + list) + echo default + list_names "${STATE_DIR}/networks" + ;; + create) touch "${STATE_DIR}/networks/$(last_argument "$@")" ;; + inspect) inspect_resource network "${1}" ;; + delete) rm -f "${STATE_DIR}/networks/${1}" ;; + *) exit 1 ;; + esac + ;; + volume) + subcommand=${1-} + shift || true + case "${subcommand}" in + list) list_names "${STATE_DIR}/volumes" ;; + create) touch "${STATE_DIR}/volumes/$(last_argument "$@")" ;; + inspect) inspect_resource volume "${1}" ;; + delete) rm -f "${STATE_DIR}/volumes/${1}" ;; + *) exit 1 ;; + esac + ;; + image) + subcommand=${1-} + case "${subcommand}" in + inspect|pull) exit 0 ;; + *) exit 1 ;; + esac + ;; + create) + name="" + while [[ $# -gt 0 ]]; do + case "$1" in + --name) + name=$2 + shift 2 + ;; + --label|--network|--platform|--ulimit|--env-file|--volume|--entrypoint|--publish) + shift 2 + ;; + *) + shift + ;; + esac + done + [[ -n "${name}" ]] + touch "${STATE_DIR}/containers/${name}" + ;; + inspect) + inspect_resource container "${1}" + ;; + start) + touch "${STATE_DIR}/running/${1}" + ;; + stop) + for argument in "$@"; do + case "${argument}" in + --time|--signal) skip_next=true ;; + [0-9]*|SIG*) ;; + *) rm -f "${STATE_DIR}/running/${argument}" ;; + esac + done + ;; + delete) + for argument in "$@"; do + case "${argument}" in + --force|-f) ;; + *) + rm -f "${STATE_DIR}/running/${argument}" + rm -f "${STATE_DIR}/containers/${argument}" + ;; + esac + done + ;; + exec) + echo 1 + ;; + logs|copy) + exit 0 + ;; + *) + echo "Unsupported fake container command: ${command} $*" >&2 + exit 1 + ;; +esac diff --git a/deploy/tests/fixtures/bin/curl b/deploy/tests/fixtures/bin/curl new file mode 100755 index 0000000000..5e611b474f --- /dev/null +++ b/deploy/tests/fixtures/bin/curl @@ -0,0 +1,4 @@ +#!/bin/bash + +set -eu +printf '{"status":"ok"}\n' diff --git a/frontend/src/api/__tests__/admin.grok.spec.ts b/frontend/src/api/__tests__/admin.grok.spec.ts new file mode 100644 index 0000000000..7560ce443f --- /dev/null +++ b/frontend/src/api/__tests__/admin.grok.spec.ts @@ -0,0 +1,37 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { post } = vi.hoisted(() => ({ + post: vi.fn(), +})) + +vi.mock('@/api/client', () => ({ + apiClient: { post }, +})) + +import { createFromSSO, getGrokSSOImportTimeout } from '@/api/admin/grok' + +describe('admin Grok SSO import API', () => { + beforeEach(() => { + post.mockReset() + post.mockResolvedValue({ data: { created: [], failed: [] } }) + }) + + it.each([ + [1, 180_000], + [3, 180_000], + [4, 270_000], + [7, 360_000], + ])('uses a timeout sized for %i keys', async (keyCount, expectedTimeout) => { + expect(getGrokSSOImportTimeout(keyCount)).toBe(expectedTimeout) + + await createFromSSO({ + sso_tokens: Array.from({ length: keyCount }, (_, index) => `sso-${index + 1}`), + }) + + expect(post).toHaveBeenCalledWith( + '/admin/grok/sso-to-oauth', + expect.objectContaining({ sso_tokens: expect.any(Array) }), + { timeout: expectedTimeout }, + ) + }) +}) diff --git a/frontend/src/api/__tests__/adminUIRequest.spec.ts b/frontend/src/api/__tests__/adminUIRequest.spec.ts new file mode 100644 index 0000000000..9064a52f1a --- /dev/null +++ b/frontend/src/api/__tests__/adminUIRequest.spec.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from 'vitest' + +import { + ADMIN_UI_REQUEST_HEADER, + shouldMarkAdminUIRequest, +} from '@/api/adminUIRequest' + +describe('Admin UI request marker', () => { + it('uses the stable request header name', () => { + expect(ADMIN_UI_REQUEST_HEADER).toBe('X-Admin-UI-Request') + }) + + it.each([ + '/admin', + '/admin/users', + '/api/v1/admin', + '/api/v1/admin/accounts?status=active', + 'https://api.example.test/api/v1/admin/dashboard', + ])('marks Admin API request %s before page navigation', (requestURL) => { + expect(shouldMarkAdminUIRequest(requestURL, '/login')).toBe(true) + }) + + it.each(['/keys', '/groups/available', '/auth/me', '/announcements'])( + 'marks shared request %s while an Admin page is active', + (requestURL) => { + expect(shouldMarkAdminUIRequest(requestURL, '/admin/dashboard')).toBe(true) + } + ) + + it.each([ + ['/keys', '/dashboard'], + ['/api/v1/administer', '/dashboard'], + ['/keys', '/administrator'], + ['', '/'], + ])('does not mark request %s on page %s', (requestURL, pagePath) => { + expect(shouldMarkAdminUIRequest(requestURL, pagePath)).toBe(false) + }) +}) diff --git a/frontend/src/api/__tests__/client.spec.ts b/frontend/src/api/__tests__/client.spec.ts index a0a05410d4..b275cca34b 100644 --- a/frontend/src/api/__tests__/client.spec.ts +++ b/frontend/src/api/__tests__/client.spec.ts @@ -12,6 +12,7 @@ describe('API Client', () => { beforeEach(async () => { localStorage.clear() + window.history.replaceState({}, '', '/') // 每次测试重新导入以获取干净的模块状态 vi.resetModules() const mod = await import('@/api/client') @@ -120,6 +121,55 @@ describe('API Client', () => { const config = adapter.mock.calls[0][0] expect(config.withCredentials).toBe(true) }) + + it('Admin API 在进入管理页面前也带 Admin UI 标记', async () => { + const adapter = vi.fn().mockResolvedValue({ + status: 200, + data: { code: 0, data: {} }, + headers: {}, + config: {}, + statusText: 'OK', + }) + apiClient.defaults.adapter = adapter + + await apiClient.get('/admin/users') + + const config = adapter.mock.calls[0][0] + expect(config.headers.get('X-Admin-UI-Request')).toBe('1') + }) + + it('管理页面调用共享 API 时带 Admin UI 标记', async () => { + window.history.replaceState({}, '', '/admin/dashboard') + const adapter = vi.fn().mockResolvedValue({ + status: 200, + data: { code: 0, data: {} }, + headers: {}, + config: {}, + statusText: 'OK', + }) + apiClient.defaults.adapter = adapter + + await apiClient.get('/groups/available') + + const config = adapter.mock.calls[0][0] + expect(config.headers.get('X-Admin-UI-Request')).toBe('1') + }) + + it('普通用户页面调用共享 API 时不带 Admin UI 标记', async () => { + const adapter = vi.fn().mockResolvedValue({ + status: 200, + data: { code: 0, data: {} }, + headers: {}, + config: {}, + statusText: 'OK', + }) + apiClient.defaults.adapter = adapter + + await apiClient.get('/groups/available') + + const config = adapter.mock.calls[0][0] + expect(config.headers.get('X-Admin-UI-Request')).toBeFalsy() + }) }) // --- 响应拦截器 --- diff --git a/frontend/src/api/admin/channelMonitor.ts b/frontend/src/api/admin/channelMonitor.ts index 0b9c62231c..de605351e3 100644 --- a/frontend/src/api/admin/channelMonitor.ts +++ b/frontend/src/api/admin/channelMonitor.ts @@ -5,7 +5,7 @@ import { apiClient } from '../client' -export type Provider = 'openai' | 'anthropic' | 'gemini' +export type Provider = 'openai' | 'anthropic' | 'gemini' | 'grok' export type MonitorStatus = 'operational' | 'degraded' | 'failed' | 'error' export type BodyOverrideMode = 'off' | 'merge' | 'replace' export type APIMode = 'chat_completions' | 'responses' diff --git a/frontend/src/api/admin/grok.ts b/frontend/src/api/admin/grok.ts index c0055d4dcc..f59c05c1d7 100644 --- a/frontend/src/api/admin/grok.ts +++ b/frontend/src/api/admin/grok.ts @@ -4,6 +4,9 @@ */ import { apiClient } from '../client' +import type { GrokBillingSummary, GrokQuotaWindow, WindowStats } from '@/types' + +export type { GrokBillingSummary, GrokQuotaWindow } from '@/types' export interface GrokAuthUrlResponse { auth_url: string @@ -34,16 +37,49 @@ export interface GrokTokenInfo { scope?: string client_id?: string email?: string + sub?: string + team_id?: string subscription_tier?: string entitlement_status?: string [key: string]: unknown } -export interface GrokQuotaWindow { - limit?: number | null - remaining?: number | null - reset_unix?: number | null - reset_at?: string | null +export interface GrokSSOToOAuthRequest { + sso_tokens: string[] + name?: string + notes?: string | null + proxy_id?: number | null + group_ids?: number[] + credentials?: Record + extra?: Record + concurrency?: number + load_factor?: number + priority?: number + rate_multiplier?: number + expires_at?: number | null + auto_pause_on_expired?: boolean +} + +export interface GrokSSOToOAuthItemResult { + index: number + name?: string + email?: string + account?: unknown + error?: string +} + +export interface GrokSSOToOAuthResponse { + created: GrokSSOToOAuthItemResult[] + failed: GrokSSOToOAuthItemResult[] +} + +const GROK_SSO_IMPORT_CONCURRENCY = 3 +const GROK_SSO_IMPORT_TIMEOUT_PER_BATCH_MS = 90_000 +const GROK_SSO_IMPORT_TIMEOUT_BUFFER_MS = 90_000 + +export function getGrokSSOImportTimeout(keyCount: number): number { + const batches = Math.ceil(Math.max(1, keyCount) / GROK_SSO_IMPORT_CONCURRENCY) + return batches * GROK_SSO_IMPORT_TIMEOUT_PER_BATCH_MS + GROK_SSO_IMPORT_TIMEOUT_BUFFER_MS } export interface GrokQuotaSnapshot { @@ -62,13 +98,19 @@ export interface GrokQuotaSnapshot { } export interface GrokQuotaProbeResult { - source: 'active_probe' - model: string + source: 'active_probe' | 'billing_probe' | 'hybrid_probe' + model?: string + billing?: GrokBillingSummary | null snapshot?: GrokQuotaSnapshot | null + local_usage_24h?: WindowStats | null + local_usage_7d?: WindowStats | null + local_usage_monthly?: WindowStats | null status_code?: number headers_observed: boolean reset_supported: boolean fetched_at: number + persisted?: boolean + probe_error?: string } export interface GrokQuotaResetResult { @@ -119,4 +161,13 @@ export async function resetQuota(id: number): Promise { return data } -export default { generateAuthUrl, exchangeCode, refreshGrokToken, queryQuota, resetQuota } +export async function createFromSSO(payload: GrokSSOToOAuthRequest): Promise { + const { data } = await apiClient.post( + '/admin/grok/sso-to-oauth', + payload, + { timeout: getGrokSSOImportTimeout(payload.sso_tokens.length) } + ) + return data +} + +export default { generateAuthUrl, exchangeCode, refreshGrokToken, queryQuota, resetQuota, createFromSSO } diff --git a/frontend/src/api/admin/ops.ts b/frontend/src/api/admin/ops.ts index c7cbc64a4b..3284ef7e15 100644 --- a/frontend/src/api/admin/ops.ts +++ b/frontend/src/api/admin/ops.ts @@ -828,6 +828,7 @@ export interface OpsRuntimeLogConfig { export interface OpsSystemLog { id: number created_at: string + host: string level: string component: string message: string @@ -849,6 +850,7 @@ export interface OpsSystemLogQuery { time_range?: '5m' | '30m' | '1h' | '6h' | '24h' | '7d' | '30d' start_time?: string end_time?: string + host?: string level?: string component?: string request_id?: string @@ -864,6 +866,7 @@ export interface OpsSystemLogQuery { export interface OpsSystemLogCleanupRequest { start_time?: string end_time?: string + host?: string level?: string component?: string request_id?: string diff --git a/frontend/src/api/adminUIRequest.ts b/frontend/src/api/adminUIRequest.ts new file mode 100644 index 0000000000..2d60e2987d --- /dev/null +++ b/frontend/src/api/adminUIRequest.ts @@ -0,0 +1,27 @@ +export const ADMIN_UI_REQUEST_HEADER = 'X-Admin-UI-Request' + +function isAdminPath(path: string): boolean { + return ( + path === '/admin' || + path.startsWith('/admin/') || + path === '/api/v1/admin' || + path.startsWith('/api/v1/admin/') + ) +} + +function requestPath(rawURL: string): string { + const value = rawURL.trim() + if (!value) return '' + try { + const origin = typeof window !== 'undefined' ? window.location.origin : 'http://localhost' + return new URL(value, origin).pathname + } catch { + return value.split(/[?#]/, 1)[0] + } +} + +export function shouldMarkAdminUIRequest(requestURL: string, pagePath?: string): boolean { + const currentPath = + pagePath ?? (typeof window !== 'undefined' ? window.location.pathname : '') + return isAdminPath(requestPath(requestURL)) || isAdminPath(currentPath) +} diff --git a/frontend/src/api/client.ts b/frontend/src/api/client.ts index 5df969f188..a2b4d2f650 100644 --- a/frontend/src/api/client.ts +++ b/frontend/src/api/client.ts @@ -6,6 +6,7 @@ import axios, { AxiosInstance, AxiosError, InternalAxiosRequestConfig, AxiosResponse } from 'axios' import type { ApiResponse } from '@/types' import { getLocale } from '@/i18n' +import { ADMIN_UI_REQUEST_HEADER, shouldMarkAdminUIRequest } from './adminUIRequest' import { getAPIBaseURL } from './url' export { buildApiUrl, buildGatewayUrl } from './url' @@ -74,6 +75,10 @@ apiClient.interceptors.request.use( config.params.timezone = getUserTimezone() } + if (config.headers && shouldMarkAdminUIRequest(String(config.url || ''))) { + config.headers[ADMIN_UI_REQUEST_HEADER] = '1' + } + return config }, (error) => { diff --git a/frontend/src/api/payment.ts b/frontend/src/api/payment.ts index ab83c55d94..e18508c535 100644 --- a/frontend/src/api/payment.ts +++ b/frontend/src/api/payment.ts @@ -7,7 +7,6 @@ import { apiClient } from './client' import type { PaymentConfig, SubscriptionPlan, - PaymentChannel, MethodLimitsResponse, CheckoutInfoResponse, CreateOrderRequest, @@ -35,11 +34,6 @@ export const paymentAPI = { return apiClient.get('/payment/plans') }, - /** Get available payment channels */ - getChannels() { - return apiClient.get('/payment/channels') - }, - /** Get all checkout page data in a single call */ getCheckoutInfo() { return apiClient.get('/payment/checkout-info') diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index 5d7af74fc6..301b512c25 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -382,7 +382,15 @@ + +
{{ t('admin.accounts.usageWindow.grokRetryAfter', { time: grokRetryAfterLabel }) }}
@@ -409,7 +425,7 @@
{{ grokQuotaStatusLine }}
- +
-
@@ -602,6 +618,7 @@ import { ref, computed, onMounted, onBeforeUnmount, onUnmounted, watch } from 'vue' import { useI18n } from 'vue-i18n' import { adminAPI } from '@/api/admin' +import type { GrokQuotaProbeResult } from '@/api/admin/grok' import type { Account, AccountUsageInfo, GeminiCredentials, WindowStats } from '@/types' import { buildOpenAIUsageRefreshKey } from '@/utils/accountUsageRefresh' import { enqueueUsageRequest } from '@/utils/usageLoadQueue' @@ -614,6 +631,8 @@ import GrokQuotaProbeCell from './GrokQuotaProbeCell.vue' // Module-level cache shared across all AccountUsageCell instances const _usageCache = new Map() const USAGE_CACHE_TTL = 5 * 60 * 1000 // 5 minutes +// xAI Free billing exposes a window without usage_percent, so estimate it from local tokens. +const GROK_FREE_TOKEN_LIMIT = 2_000_000 const props = withDefaults( defineProps<{ @@ -1047,9 +1066,58 @@ const makeGrokQuotaBar = (quota?: { limit?: number | null; remaining?: number | const grokRequestQuotaBar = computed(() => makeGrokQuotaBar(usageInfo.value?.grok_request_quota)) const grokTokenQuotaBar = computed(() => makeGrokQuotaBar(usageInfo.value?.grok_token_quota)) +const grokBilling = computed(() => usageInfo.value?.grok_billing || null) +const grokWeeklyBillingBar = computed((): GrokQuotaBarInfo | null => { + const billing = grokBilling.value + if (billing?.period_type?.toLowerCase() !== 'weekly' || billing.usage_percent == null) { + return null + } + return { + utilization: Math.min(100, Math.max(0, billing.usage_percent)), + resetsAt: billing.period_end || null + } +}) +const grokPlanLabelIsFree = (value: string) => value.includes('free') || value.includes('basic') +const grokPlanLabelIsPaid = (value: string) => { + return value !== '' && !grokPlanLabelIsFree(value) && !value.includes('unknown') +} +const grokIsFree = computed(() => { + if (props.account.platform !== 'grok' || props.account.type !== 'oauth') return false + const billing = grokBilling.value + if ( + billing?.usage_percent != null || + billing?.used_percent != null || + (billing?.monthly_limit_cents != null && billing.monthly_limit_cents > 0) + ) return false + + const plan = (billing?.plan || '').trim().toLowerCase() + const tier = (usageInfo.value?.subscription_tier || '').trim().toLowerCase() + const entitlement = (usageInfo.value?.grok_entitlement_status || '').toLowerCase() + if (grokPlanLabelIsPaid(plan) || grokPlanLabelIsPaid(tier)) return false + if ( + grokPlanLabelIsFree(plan) || + grokPlanLabelIsFree(tier) || + grokPlanLabelIsFree(entitlement) + ) return true + return billing != null +}) +const grokFreeQuotaUsage = computed(() => usageInfo.value?.grok_local_usage_24h || null) +const grokLocalUsage = computed(() => { + if (grokIsFree.value) return grokFreeQuotaUsage.value + return props.todayStats || + usageInfo.value?.grok_local_usage || + usageInfo.value?.grok_local_usage_7d || + usageInfo.value?.grok_local_usage_monthly || + null +}) +const grokFreeTokenBar = computed(() => { + if (!grokIsFree.value || !grokFreeQuotaUsage.value) return null + const used = Math.max(0, grokFreeQuotaUsage.value.tokens || 0) + return { utilization: Math.min(100, (used / GROK_FREE_TOKEN_LIMIT) * 100) } +}) const grokQuotaUnknown = computed(() => { if (props.account.platform !== 'grok') return false - if (grokRequestQuotaBar.value || grokTokenQuotaBar.value) return false + if (grokBilling.value || grokFreeTokenBar.value || grokRequestQuotaBar.value || grokTokenQuotaBar.value) return false return usageInfo.value?.grok_quota_snapshot_state !== 'observed' }) const grokQuotaUnknownLabel = computed(() => { @@ -1080,7 +1148,6 @@ const grokQuotaStatusLine = computed(() => { } return parts.length > 0 ? parts.join(' | ') : null }) -const grokLocalUsage = computed(() => usageInfo.value?.grok_local_usage || props.todayStats || null) const grokEntitlementLabel = computed(() => { const status = (usageInfo.value?.grok_entitlement_status || '').trim() return status || null @@ -1283,6 +1350,35 @@ const loadActiveUsage = async () => { } } +const handleGrokProbed = (result: GrokQuotaProbeResult) => { + const current = usageInfo.value + if (!current) return + const snapshot = result.snapshot + const merged: AccountUsageInfo = { + ...current, + grok_billing: result.billing ?? current.grok_billing, + grok_local_usage_24h: result.local_usage_24h ?? current.grok_local_usage_24h, + grok_local_usage_7d: result.local_usage_7d ?? current.grok_local_usage_7d, + grok_local_usage_monthly: result.local_usage_monthly ?? current.grok_local_usage_monthly, + grok_request_quota: snapshot?.requests ?? current.grok_request_quota, + grok_token_quota: snapshot?.tokens ?? current.grok_token_quota, + grok_retry_after_seconds: snapshot?.retry_after_seconds ?? current.grok_retry_after_seconds, + grok_entitlement_status: snapshot?.entitlement_status || current.grok_entitlement_status, + grok_quota_snapshot_state: result.billing + ? 'billing_observed' + : snapshot?.headers_observed + ? 'observed' + : current.grok_quota_snapshot_state, + grok_last_quota_probe_at: result.billing?.fetched_at ?? snapshot?.last_probe_at ?? current.grok_last_quota_probe_at, + grok_last_headers_seen_at: snapshot?.last_headers_seen_at ?? current.grok_last_headers_seen_at, + grok_last_status_code: result.status_code ?? snapshot?.status_code ?? current.grok_last_status_code, + error: result.billing || snapshot ? undefined : current.error, + error_code: result.billing || snapshot ? undefined : current.error_code + } + usageInfo.value = merged + _usageCache.set(props.account.id, { data: merged, ts: Date.now() }) +} + // ===== API Key quota progress bars ===== interface QuotaBarInfo { diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index dad7c7a26f..ea9deb7ba6 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -50,7 +50,7 @@
@@ -2824,6 +2824,38 @@
+
+
+
+ +

+ {{ t('admin.accounts.openai.longContextBillingDesc') }} +

+
+ +
+
+
@@ -3492,6 +3528,7 @@ interface OAuthFlowExposed { sessionToken: string codexSession: string codexPAT: string + ssoCookie: string inputMethod: AuthInputMethod reset: () => void } @@ -3676,6 +3713,8 @@ const fillHeaderOverrideTemplate = () => { const interceptWarmupRequests = ref(false) const autoPauseOnExpired = ref(true) const openaiPassthroughEnabled = ref(false) +const openAILongContextBillingEnabled = ref(false) +const openAILongContextBillingTouched = ref(false) const openAICompactMode = ref('auto') const openAIResponsesMode = ref('auto') const openAIEndpointCapabilities = ref(['chat_completions', 'embeddings']) @@ -3688,6 +3727,11 @@ const anthropicPassthroughEnabled = ref(false) const anthropicAPIKeyAuthScheme = ref('x_api_key') const webSearchEmulationMode = ref('default') const webSearchGlobalEnabled = ref(false) + +const toggleOpenAILongContextBilling = () => { + openAILongContextBillingEnabled.value = !openAILongContextBillingEnabled.value + openAILongContextBillingTouched.value = true +} const { globalEnabled: quotaNotifyGlobalEnabled, state: quotaNotifyState, @@ -3976,6 +4020,8 @@ const isOAuthFlow = computed(() => { return accountCategory.value === 'oauth-based' }) +const isGrokSSOInputMethod = computed(() => form.platform === 'grok' && oauthFlowRef.value?.inputMethod === 'sso_cookie') + const isManualInputMethod = computed(() => { return oauthFlowRef.value?.inputMethod === 'manual' }) @@ -4530,6 +4576,8 @@ const resetForm = () => { interceptWarmupRequests.value = false autoPauseOnExpired.value = true openaiPassthroughEnabled.value = false + openAILongContextBillingEnabled.value = false + openAILongContextBillingTouched.value = false openAICompactMode.value = 'auto' openAIResponsesMode.value = 'auto' openAIEndpointCapabilities.value = ['chat_completions', 'embeddings'] @@ -4612,6 +4660,7 @@ const buildOpenAIExtra = (base?: Record): Record): Record 0 ? extra : undefined } +const buildOpenAICodexImportExtra = (): Record | undefined => { + const extra = buildOpenAIExtra() + if (!extra) { + return undefined + } + if (!openAILongContextBillingTouched.value) { + delete extra.openai_long_context_billing_enabled + } + return Object.keys(extra).length > 0 ? extra : undefined +} + const buildAnthropicExtra = (base?: Record): Record | undefined => { if (form.platform !== 'anthropic' || accountCategory.value !== 'apikey') { return base @@ -4766,7 +4826,7 @@ const handleVertexServiceAccountDrop = async (event: DragEvent) => { const handleSubmit = async () => { // For OAuth-based type, handle OAuth flow (goes to step 2) if (isOAuthFlow.value) { - if (!form.name.trim()) { + if (!isGrokSSOInputMethod.value && !form.name.trim()) { appStore.showError(t('admin.accounts.pleaseEnterAccountName')) return } @@ -5204,6 +5264,76 @@ const handleGrokValidateRT = async (refreshTokenInput: string) => { } } +const handleGrokImportSSO = async (ssoInput: string) => { + // Align with OpenAI/Grok RT batch import: one token per line, no client-side dedupe. + const ssoTokens = ssoInput + .split('\n') + .map((token) => token.trim()) + .filter((token) => token) + if (ssoTokens.length === 0) return + + grokOAuth.loading.value = true + grokOAuth.error.value = '' + + const credentials: Record = {} + const modelMapping = buildModelMappingObject(modelRestrictionMode.value, allowedModels.value, modelMappings.value) + if (modelMapping) { + credentials.model_mapping = modelMapping + } + if (!applyTempUnschedConfig(credentials)) { + grokOAuth.loading.value = false + return + } + + try { + const result = await adminAPI.grok.createFromSSO({ + sso_tokens: ssoTokens, + name: form.name || undefined, + notes: form.notes || undefined, + proxy_id: form.proxy_id, + group_ids: form.group_ids, + credentials, + concurrency: form.concurrency, + load_factor: form.load_factor ?? undefined, + priority: form.priority, + rate_multiplier: form.rate_multiplier, + expires_at: form.expires_at, + auto_pause_on_expired: autoPauseOnExpired.value + }) + + const successCount = result.created?.length || 0 + const failedCount = result.failed?.length || 0 + if (successCount > 0 && failedCount === 0) { + appStore.showSuccess( + ssoTokens.length > 1 + ? t('admin.accounts.oauth.batchSuccess', { count: successCount }) + : t('admin.accounts.accountCreated') + ) + emit('created') + handleClose() + } else if (successCount > 0 && failedCount > 0) { + // Same as OpenAI/Grok RT: keep input, show failures, refresh list. + appStore.showWarning( + t('admin.accounts.oauth.batchPartialSuccess', { success: successCount, failed: failedCount }) + ) + grokOAuth.error.value = (result.failed || []) + .map((item) => `#${item.index}: ${item.error || 'Unknown error'}`) + .join('\n') + emit('created') + } else { + grokOAuth.error.value = (result.failed || []) + .map((item) => `#${item.index}: ${item.error || 'Unknown error'}`) + .join('\n') || t('admin.accounts.oauth.grok.failedToConvertSSO') + appStore.showError(t('admin.accounts.oauth.batchFailed')) + } + } catch (error: any) { + grokOAuth.error.value = error.response?.data?.detail || error.message || t('admin.accounts.oauth.grok.failedToConvertSSO') + appStore.showError(grokOAuth.error.value) + } finally { + grokOAuth.loading.value = false + } +} + // OpenAI OAuth 授权码兑换 const handleOpenAIExchange = async (authCode: string) => { const oauthClient = openaiOAuth @@ -5332,7 +5462,7 @@ const handleOpenAIImportCodexSession = async (content: string) => { oauthClient.error.value = '' try { - const extra = buildOpenAIExtra() + const extra = buildOpenAICodexImportExtra() const result = await adminAPI.accounts.importCodexSession({ content: trimmed, name: form.name, @@ -5410,7 +5540,7 @@ const handleOpenAIImportCodexPAT = async (accessToken: string) => { oauthClient.error.value = '' try { - const extra = buildOpenAIExtra() + const extra = buildOpenAICodexImportExtra() await adminAPI.accounts.createOpenAICodexPAT({ access_token: trimmed, name: form.name, diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 3d16a901fe..cfc2fed151 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -1786,7 +1786,39 @@ /> - + +
+
+
+ +

+ {{ t('admin.accounts.openai.longContextBillingDesc') }} +

+
+ +
+
+
+ +
+
+
+ +

+ {{ t('admin.accounts.openai.planTypeDesc') }} +

+
+
+ +
+ +
+
+

+ {{ t(getOAuthKey('ssoCookieDesc')) }} +

+ +
+ + +

+ {{ t(getOAuthKey('ssoCookieHint')) }} +

+
+ +
+

+ {{ error }} +

+
+ + +
+
+
(), { showAccessTokenOption: false, showCodexSessionImportOption: false, showCodexPatOption: false, + showSsoOption: false, + showManualOption: true, + initialInputMethod: 'manual', platform: 'anthropic', showProjectId: true }) @@ -771,6 +863,7 @@ const emit = defineEmits<{ 'import-access-token': [accessToken: string] 'import-codex-session': [content: string] 'import-codex-pat': [accessToken: string] + 'import-sso': [content: string] 'update:inputMethod': [method: AuthInputMethod] }>() @@ -807,19 +900,31 @@ const oauthImportantNotice = computed(() => { }) // Local state -const inputMethod = ref(props.showCookieOption ? 'manual' : 'manual') +const inputMethod = ref(props.initialInputMethod) const authCodeInput = ref('') const sessionKeyInput = ref('') const refreshTokenInput = ref('') const sessionTokenInput = ref('') const codexSessionInput = ref('') const codexPATInput = ref('') +const ssoCookieInput = ref('') const showHelpDialog = ref(false) const oauthState = ref('') const projectId = ref('') -// Computed: show method selection when either cookie or refresh token option is enabled -const showMethodSelection = computed(() => props.showCookieOption || props.showRefreshTokenOption || props.showMobileRefreshTokenOption || props.showSessionTokenOption || props.showAccessTokenOption || props.showCodexSessionImportOption || props.showCodexPatOption) +// Computed: show method selection only when there is something to choose. +const methodOptionCount = computed(() => [ + props.showManualOption, + props.showCookieOption, + props.showRefreshTokenOption, + props.showMobileRefreshTokenOption, + props.showSessionTokenOption, + props.showAccessTokenOption, + props.showCodexSessionImportOption, + props.showCodexPatOption, + props.showSsoOption +].filter(Boolean).length) +const showMethodSelection = computed(() => methodOptionCount.value > 1) // Clipboard const { copied, copyToClipboard } = useClipboard() @@ -850,7 +955,18 @@ const parsedCodexSessionCount = computed(() => { .filter((item) => item).length }) +const parsedSSOCount = computed(() => { + return ssoCookieInput.value + .split('\n') + .map((item) => item.trim()) + .filter((item) => item).length +}) + // Watchers +watch(() => props.initialInputMethod, (newVal) => { + inputMethod.value = newVal +}) + watch(inputMethod, (newVal) => { emit('update:inputMethod', newVal) }) @@ -933,6 +1049,12 @@ const handleImportCodexPAT = () => { } } +const handleImportSSO = () => { + if (ssoCookieInput.value.trim()) { + emit('import-sso', ssoCookieInput.value.trim()) + } +} + // Expose methods and state defineExpose({ authCode: authCodeInput, @@ -943,6 +1065,7 @@ defineExpose({ sessionToken: sessionTokenInput, codexSession: codexSessionInput, codexPAT: codexPATInput, + ssoCookie: ssoCookieInput, inputMethod, reset: () => { authCodeInput.value = '' @@ -953,7 +1076,8 @@ defineExpose({ sessionTokenInput.value = '' codexSessionInput.value = '' codexPATInput.value = '' - inputMethod.value = 'manual' + ssoCookieInput.value = '' + inputMethod.value = props.initialInputMethod showHelpDialog.value = false } }) diff --git a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts index 2abf6513ca..7c2fe48a56 100644 --- a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts +++ b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts @@ -660,6 +660,391 @@ describe('AccountUsageCell', () => { expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokTokens|25|true') }) + it('Grok OAuth uses the official weekly billing percentage when available', async () => { + getUsage.mockResolvedValue({ + grok_billing: { + period_type: 'weekly', + usage_percent: 37, + period_end: '2026-07-16T03:25:00Z', + plan: 'SuperGrok' + }, + grok_local_usage: { + requests: 5, + tokens: 2_200_000, + cost: 4.42, + standard_cost: 4.42, + user_cost: 0.44 + }, + grok_request_quota: { limit: 100, remaining: 100 }, + grok_token_quota: { limit: 2_000_000, remaining: 2_000_000 } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4201, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization', 'resetsAt', 'remainingCapacity'], + template: '
{{ label }}|{{ utilization }}|{{ resetsAt }}|{{ remainingCapacity }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('7d|37|2026-07-16T03:25:00Z') + expect(wrapper.text()).not.toContain('admin.accounts.usageWindow.grokRequests|') + expect(wrapper.text()).not.toContain('admin.accounts.usageWindow.grokTokens|') + expect(wrapper.text()).not.toContain('2M|') + }) + + it.each([ + { tokens: 0, expected: 0, compact: '0' }, + { tokens: 1_000_000, expected: 50, compact: '1.0M' }, + { tokens: 2_000_000, expected: 100, compact: '2.0M' }, + { tokens: 2_200_000, expected: 100, compact: '2.2M' } + ])('Grok Free derives its 2M quota from local tokens: $tokens -> $expected%', async ({ tokens, expected, compact }) => { + getUsage.mockResolvedValue({ + grok_billing: { + period_type: 'weekly', + usage_percent: null, + plan: '' + }, + grok_local_usage_24h: { + requests: 5, + tokens, + cost: 0, + standard_cost: 0, + user_cost: 0 + }, + grok_request_quota: { limit: 100, remaining: 100 }, + grok_token_quota: { limit: 2_000_000, remaining: 2_000_000 } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4300 + expected, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain(`24h|${expected}`) + expect(wrapper.findAll('span').filter((node) => node.text() === compact)).toHaveLength(1) + expect(wrapper.findAll('.usage-bar')).toHaveLength(1) + expect(wrapper.text()).not.toContain('admin.accounts.usageWindow.grokRequests|') + expect(wrapper.text()).not.toContain('admin.accounts.usageWindow.grokTokens|') + }) + + it('Grok Free uses rolling 24h usage instead of today-only usage', async () => { + getUsage.mockResolvedValue({ + grok_billing: { period_type: 'weekly', usage_percent: null, plan: '' }, + grok_local_usage: { + requests: 2, + tokens: 250_000, + cost: 0, + standard_cost: 0 + }, + grok_local_usage_24h: { + requests: 12, + tokens: 1_500_000, + cost: 0, + standard_cost: 0 + } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4398, platform: 'grok', type: 'oauth', extra: {} }), + todayStats: { + requests: 2, + tokens: 200_000, + cost: 0, + standard_cost: 0 + } + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization', 'title'], + template: '
{{ label }}|{{ utilization }}|{{ title }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('24h|75|admin.accounts.usageWindow.grokFreeQuota24hHint') + expect(wrapper.text()).toContain('1.5M') + expect(wrapper.text()).not.toContain('7d|') + expect(wrapper.text()).not.toContain('200.0K') + expect(wrapper.text()).not.toContain('250.0K') + }) + + it('Grok Free does not substitute today stats when rolling 24h usage is unavailable', async () => { + getUsage.mockResolvedValue({ + grok_billing: { period_type: 'weekly', usage_percent: null, plan: '' }, + grok_local_usage: { + requests: 1, + tokens: 250_000, + cost: 0, + standard_cost: 0, + user_cost: 0 + } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4399, platform: 'grok', type: 'oauth', extra: {} }), + todayStats: { + requests: 4, + tokens: 1_000_000, + cost: 0, + standard_cost: 0, + user_cost: 0 + } + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.findAll('.usage-bar')).toHaveLength(0) + expect(wrapper.text()).not.toContain('24h|') + expect(wrapper.text()).not.toContain('1.0M') + expect(wrapper.text()).not.toContain('250.0K') + }) + + it('Grok paid plans are not mistaken for Free when weekly usage is temporarily missing', async () => { + getUsage.mockResolvedValue({ + grok_billing: { + period_type: 'weekly', + usage_percent: null, + plan: 'SuperGrok Heavy' + }, + grok_entitlement_status: 'free', + grok_local_usage: { + requests: 2, + tokens: 2_000_000, + cost: 1, + standard_cost: 1 + }, + grok_token_quota: { limit: 1_000, remaining: 250 } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4401, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokTokens|25') + expect(wrapper.text()).not.toContain('2M|') + }) + + it('Grok custom paid monthly limits override stale Free entitlement', async () => { + getUsage.mockResolvedValue({ + grok_billing: { + period_type: 'weekly', + usage_percent: null, + monthly_limit_cents: 25_000, + plan: '' + }, + grok_entitlement_status: 'free', + grok_local_usage: { + requests: 2, + tokens: 2_000_000, + cost: 1, + standard_cost: 1 + }, + grok_token_quota: { limit: 1_000, remaining: 250 } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4402, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokTokens|25') + expect(wrapper.text()).not.toContain('2M|') + }) + + it('Grok credential Free tier keeps the 2M fallback when billing is unavailable', async () => { + getUsage.mockResolvedValue({ + subscription_tier: 'FREE', + grok_local_usage_24h: { + requests: 3, + tokens: 1_000_000, + cost: 0, + standard_cost: 0 + } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4403, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('24h|50') + }) + + it('Grok paid manual probes keep the weekly/local summary when 24h usage is returned', async () => { + getUsage.mockResolvedValue({ + grok_quota_snapshot_state: 'no_headers', + error: 'stale error', + error_code: 'quota_unknown' + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4501, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization', 'resetsAt'], + template: '
{{ label }}|{{ utilization }}|{{ resetsAt }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: { + emits: ['probed'], + template: `` + } + } + } + }) + + await flushPromises() + await wrapper.get('.probe').trigger('click') + + expect(wrapper.text()).toContain('7d|42|2026-07-17T00:00:00Z') + expect(wrapper.text()).toContain('1.0M') + expect(wrapper.text()).not.toContain('750.0K') + expect(wrapper.text()).toContain('ACTIVE') + expect(wrapper.text()).not.toContain('stale error') + }) + + it('Grok Free manual probes merge rolling 24h usage', async () => { + getUsage.mockResolvedValue({ + subscription_tier: 'FREE', + grok_quota_snapshot_state: 'no_headers' + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4502, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: { + emits: ['probed'], + template: `` + } + } + } + }) + + await flushPromises() + await wrapper.get('.probe').trigger('click') + + expect(wrapper.text()).toContain('24h|75') + expect(wrapper.text()).toContain('1.5M') + expect(wrapper.text()).not.toContain('7d|') + }) + it('Key 账号在 today stats loading 时显示骨架屏', async () => { const wrapper = mount(AccountUsageCell, { props: { diff --git a/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts new file mode 100644 index 0000000000..62c97d35a4 --- /dev/null +++ b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts @@ -0,0 +1,213 @@ +import { defineComponent } from 'vue' +import { flushPromises, mount } from '@vue/test-utils' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { + createAccountMock, + importCodexSessionMock, + createOpenAICodexPATMock, +} = vi.hoisted(() => ({ + createAccountMock: vi.fn(), + importCodexSessionMock: vi.fn(), + createOpenAICodexPATMock: vi.fn(), +})) + +vi.mock('@/stores/app', () => ({ + useAppStore: () => ({ + showError: vi.fn(), + showSuccess: vi.fn(), + showWarning: vi.fn(), + }), +})) + +vi.mock('@/stores/auth', () => ({ + useAuthStore: () => ({ isSimpleMode: true }), +})) + +vi.mock('@/api/admin', () => ({ + adminAPI: { + accounts: { + create: createAccountMock, + checkMixedChannelRisk: vi.fn().mockResolvedValue({ has_risk: false }), + importCodexSession: importCodexSessionMock, + createOpenAICodexPAT: createOpenAICodexPATMock, + }, + settings: { + getWebSearchEmulationConfig: vi.fn().mockResolvedValue({ enabled: false, providers: [] }), + getSettings: vi.fn().mockResolvedValue({}), + }, + tlsFingerprintProfiles: { + list: vi.fn().mockResolvedValue([]), + }, + }, +})) + +vi.mock('@/api/admin/accounts', () => ({ + getAntigravityDefaultModelMapping: vi.fn().mockResolvedValue([]), +})) + +vi.mock('vue-i18n', async () => { + const actual = await vi.importActual('vue-i18n') + return { + ...actual, + useI18n: () => ({ t: (key: string) => key }), + } +}) + +import CreateAccountModal from '../CreateAccountModal.vue' + +const BaseDialogStub = defineComponent({ + name: 'BaseDialog', + props: { show: { type: Boolean, default: false } }, + template: '
', +}) + +const OAuthAuthorizationFlowStub = defineComponent({ + name: 'OAuthAuthorizationFlow', + emits: ['import-codex-session', 'import-codex-pat'], + template: ` +
+ + +
+ `, +}) + +function mountModal() { + return mount(CreateAccountModal, { + props: { show: true, proxies: [], groups: [] }, + global: { + stubs: { + BaseDialog: BaseDialogStub, + OAuthAuthorizationFlow: OAuthAuthorizationFlowStub, + ConfirmDialog: true, + Select: true, + Icon: true, + PlatformIcon: true, + ProxySelector: true, + ProxyAdBanner: true, + GroupSelector: true, + ModelWhitelistSelector: true, + QuotaLimitCard: true, + }, + }, + }) +} + +async function selectButtonByText(wrapper: ReturnType, text: string) { + const button = wrapper.findAll('button').find((candidate) => candidate.text().includes(text)) + expect(button).toBeDefined() + await button?.trigger('click') +} + +async function submitApiKeyAccount(platform: 'openai' | 'anthropic', enableLongContextBilling = false) { + const wrapper = mountModal() + await selectButtonByText(wrapper, platform === 'openai' ? 'OpenAI' : 'admin.accounts.claudeConsole') + if (platform === 'openai') { + await selectButtonByText(wrapper, 'API Key') + } + await wrapper.get('form#create-account-form input[type="text"]').setValue(`${platform} account`) + await wrapper.get('form#create-account-form input[type="password"]').setValue('test-api-key') + if (enableLongContextBilling) { + await wrapper.get('[data-testid="openai-long-context-billing-toggle"]').trigger('click') + } + await wrapper.get('form#create-account-form').trigger('submit.prevent') + await flushPromises() +} + +async function openCodexImportStep(toggleClicks = 0) { + const wrapper = mountModal() + await selectButtonByText(wrapper, 'OpenAI') + for (let click = 0; click < toggleClicks; click += 1) { + await wrapper.get('[data-testid="openai-long-context-billing-toggle"]').trigger('click') + } + await wrapper.get('form#create-account-form input[type="text"]').setValue('Codex import') + await wrapper.get('form#create-account-form').trigger('submit.prevent') + return wrapper +} + +describe('CreateAccountModal OpenAI long-context billing', () => { + beforeEach(() => { + createAccountMock.mockReset().mockResolvedValue({}) + importCodexSessionMock.mockReset().mockResolvedValue({ + created: 1, + updated: 0, + skipped: 0, + failed: 0, + errors: [], + warnings: [], + }) + createOpenAICodexPATMock.mockReset().mockResolvedValue({}) + }) + + it('sends false explicitly for normal OpenAI account creation by default', async () => { + await submitApiKeyAccount('openai') + + expect(createAccountMock).toHaveBeenCalledTimes(1) + expect(createAccountMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBe(false) + }) + + it('sends true explicitly when OpenAI long-context billing is enabled', async () => { + await submitApiKeyAccount('openai', true) + + expect(createAccountMock).toHaveBeenCalledTimes(1) + expect(createAccountMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBe(true) + }) + + it('omits the OpenAI setting for non-OpenAI account creation', async () => { + await submitApiKeyAccount('anthropic') + + expect(createAccountMock).toHaveBeenCalledTimes(1) + expect(createAccountMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBeUndefined() + }) + + it('leaves Codex session import billing ownership to the backend', async () => { + const wrapper = await openCodexImportStep() + await wrapper.get('[data-testid="import-codex-session"]').trigger('click') + await flushPromises() + + expect(importCodexSessionMock).toHaveBeenCalledTimes(1) + expect(importCodexSessionMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBeUndefined() + }) + + it('leaves Codex PAT import billing ownership to the backend', async () => { + const wrapper = await openCodexImportStep() + await wrapper.get('[data-testid="import-codex-pat"]').trigger('click') + await flushPromises() + + expect(createOpenAICodexPATMock).toHaveBeenCalledTimes(1) + expect(createOpenAICodexPATMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBeUndefined() + }) + + it('sends explicit true for Codex session import after the toggle is enabled', async () => { + const wrapper = await openCodexImportStep(1) + await wrapper.get('[data-testid="import-codex-session"]').trigger('click') + await flushPromises() + + expect(importCodexSessionMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBe(true) + }) + + it('sends explicit false for Codex session import after the toggle is changed back', async () => { + const wrapper = await openCodexImportStep(2) + await wrapper.get('[data-testid="import-codex-session"]').trigger('click') + await flushPromises() + + expect(importCodexSessionMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBe(false) + }) + + it('sends explicit true for Codex PAT import after the toggle is enabled', async () => { + const wrapper = await openCodexImportStep(1) + await wrapper.get('[data-testid="import-codex-pat"]').trigger('click') + await flushPromises() + + expect(createOpenAICodexPATMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBe(true) + }) + + it('sends explicit false for Codex PAT import after the toggle is changed back', async () => { + const wrapper = await openCodexImportStep(2) + await wrapper.get('[data-testid="import-codex-pat"]').trigger('click') + await flushPromises() + + expect(createOpenAICodexPATMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBe(false) + }) +}) diff --git a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts index 44691200f9..b3a583d102 100644 --- a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts @@ -395,6 +395,105 @@ describe('EditAccountModal', () => { }) }) + it('loads and submits the per-account OpenAI long-context billing toggle', async () => { + const account = buildAccount() + account.extra = { + openai_long_context_billing_enabled: true + } + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + const toggle = wrapper.get('[data-testid="openai-long-context-billing-toggle"]') + expect(toggle.attributes('aria-checked')).toBe('true') + + await toggle.trigger('click') + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.openai_long_context_billing_enabled).toBe(false) + }) + + it('defaults legacy OpenAI accounts to long-context billing disabled', async () => { + const account = buildAccount() + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + const toggle = wrapper.get('[data-testid="openai-long-context-billing-toggle"]') + expect(toggle.attributes('aria-checked')).toBe('false') + + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.openai_long_context_billing_enabled).toBe(false) + }) + + it('does not render or submit the long-context billing toggle for Spark shadow accounts', async () => { + const account = buildOpenAISparkShadowAccount() + account.extra = { + openai_long_context_billing_enabled: false + } + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + const wrapper = mountModal(account) + + expect(wrapper.find('[data-testid="openai-long-context-billing-toggle"]').exists()).toBe(false) + + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra).not.toHaveProperty( + 'openai_long_context_billing_enabled' + ) + }) + + it('preserves an explicit OpenAI long-context billing opt-out', async () => { + const account = buildAccount() + account.extra = { + openai_long_context_billing_enabled: false + } + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + const toggle = wrapper.get('[data-testid="openai-long-context-billing-toggle"]') + expect(toggle.attributes('aria-checked')).toBe('false') + + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.openai_long_context_billing_enabled).toBe(false) + }) + + it('fails closed for malformed OpenAI long-context billing values', async () => { + const account = buildAccount() + account.extra = { + openai_long_context_billing_enabled: 'false' + } + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + + expect(wrapper.get('[data-testid="openai-long-context-billing-toggle"]').attributes('aria-checked')).toBe('false') + + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.openai_long_context_billing_enabled).toBe(false) + }) + it('loads and submits Grok OAuth model mapping edits', async () => { const account = buildGrokOAuthAccount() updateAccountMock.mockReset() diff --git a/frontend/src/components/account/__tests__/GrokQuotaProbeCell.spec.ts b/frontend/src/components/account/__tests__/GrokQuotaProbeCell.spec.ts new file mode 100644 index 0000000000..9b431b339e --- /dev/null +++ b/frontend/src/components/account/__tests__/GrokQuotaProbeCell.spec.ts @@ -0,0 +1,54 @@ +import { flushPromises, mount } from '@vue/test-utils' +import { beforeEach, describe, expect, it, vi } from 'vitest' +import GrokQuotaProbeCell from '../GrokQuotaProbeCell.vue' +import type { Account } from '@/types' + +const { queryQuota } = vi.hoisted(() => ({ + queryQuota: vi.fn() +})) + +vi.mock('@/api/admin', () => ({ + adminAPI: { + grok: { queryQuota } + } +})) + +vi.mock('vue-i18n', () => ({ + useI18n: () => ({ + t: (key: string, params?: Record) => + params?.percent == null ? key : `${key}:${params.percent}` + }) +})) + +const account = { + id: 99, + platform: 'grok', + type: 'oauth' +} as Account + +describe('GrokQuotaProbeCell', () => { + beforeEach(() => { + queryQuota.mockReset() + }) + + it('keeps billing data while exposing a failed Free quota fallback', async () => { + queryQuota.mockResolvedValue({ + source: 'hybrid_probe', + billing: { period_type: 'weekly', usage_percent: null }, + headers_observed: false, + reset_supported: false, + fetched_at: 1, + probe_error: 'upstream returned 402 for probe model "grok-4.5"' + }) + const wrapper = mount(GrokQuotaProbeCell, { props: { account } }) + + await wrapper.get('button').trigger('click') + await flushPromises() + + expect(wrapper.text()).toContain('upstream returned 402 for probe model "grok-4.5"') + expect(wrapper.emitted('probed')?.[0]?.[0]).toMatchObject({ + billing: { period_type: 'weekly', usage_percent: null }, + probe_error: 'upstream returned 402 for probe model "grok-4.5"' + }) + }) +}) diff --git a/frontend/src/components/account/__tests__/credentialsBuilder.spec.ts b/frontend/src/components/account/__tests__/credentialsBuilder.spec.ts index c2cb093805..ae4c6739a8 100644 --- a/frontend/src/components/account/__tests__/credentialsBuilder.spec.ts +++ b/frontend/src/components/account/__tests__/credentialsBuilder.spec.ts @@ -6,9 +6,13 @@ import { applyAntigravityProjectID, applyHeaderOverride, applyInterceptWarmup, + applyPlanType, buildHeaderOverridesObject, + buildPlanTypeOptions, getHeaderOverrideTemplate, isHeaderOverridePlatform, + planTypeDisplayLabel, + readPlanType, splitHeaderOverridesObject, validateHeaderOverrideRows } from '../credentialsBuilder' @@ -289,3 +293,88 @@ describe('validateHeaderOverrideRows session isolation headers', () => { expect(validateHeaderOverrideRows([{ name: 'x'.repeat(201), value: 'v' }])).toBe('invalidName') }) }) + +describe('plan_type helpers', () => { + describe('planTypeDisplayLabel', () => { + it('maps canonical + alias values to friendly labels', () => { + expect(planTypeDisplayLabel('plus')).toBe('Plus') + expect(planTypeDisplayLabel('pro')).toBe('Pro') + expect(planTypeDisplayLabel('chatgptpro')).toBe('Pro') + expect(planTypeDisplayLabel('free')).toBe('Free') + expect(planTypeDisplayLabel('team')).toBe('Team') + expect(planTypeDisplayLabel('CHATGPTPRO')).toBe('Pro') + }) + it('returns unknown values verbatim', () => { + expect(planTypeDisplayLabel('self_serve_business')).toBe('self_serve_business') + }) + }) + + describe('readPlanType', () => { + it('reads a string plan_type', () => { + expect(readPlanType({ plan_type: 'plus' })).toBe('plus') + }) + it('treats non-string / missing values as empty', () => { + expect(readPlanType({ plan_type: 42 })).toBe('') + expect(readPlanType({ plan_type: true })).toBe('') + expect(readPlanType({})).toBe('') + expect(readPlanType(undefined)).toBe('') + expect(readPlanType(null)).toBe('') + }) + }) + + describe('buildPlanTypeOptions', () => { + const clear = 'Clear' + it('returns clear + presets when current is empty', () => { + expect(buildPlanTypeOptions('', clear)).toEqual([ + { value: '', label: clear }, + { value: 'plus', label: 'Plus' }, + { value: 'pro', label: 'Pro' }, + { value: 'free', label: 'Free' } + ]) + }) + it('keeps canonical chatgptpro under a single friendly "Pro" option (no duplicate)', () => { + const opts = buildPlanTypeOptions('chatgptpro', clear) + const pros = opts.filter(o => o.label === 'Pro') + expect(pros).toHaveLength(1) + expect(pros[0].value).toBe('chatgptpro') + expect(opts.map(o => o.value)).toEqual(['', 'plus', 'chatgptpro', 'free']) + }) + it('appends an unknown-but-labeled value (team) as its own option', () => { + const opts = buildPlanTypeOptions('team', clear) + expect(opts.find(o => o.value === 'team')).toEqual({ value: 'team', label: 'Team' }) + // presets untouched + expect(opts.map(o => o.value)).toEqual(['', 'plus', 'pro', 'free', 'team']) + }) + it('appends a fully custom value with a raw label', () => { + const opts = buildPlanTypeOptions('weird_x', clear) + expect(opts.at(-1)).toEqual({ value: 'weird_x', label: 'weird_x' }) + }) + it('does not duplicate an exact preset value', () => { + const opts = buildPlanTypeOptions('pro', clear) + expect(opts.filter(o => o.value === 'pro')).toHaveLength(1) + expect(opts.map(o => o.value)).toEqual(['', 'plus', 'pro', 'free']) + }) + }) + + describe('applyPlanType', () => { + it('sets plan_type and preserves all other credential keys', () => { + const creds = { + chatgpt_account_id: 'acc', + email: 'a@b.c', + subscription_expires_at: '2026-01-01', + model_mapping: { x: 'y' } + } + const out = applyPlanType({ ...creds }, 'plus') + expect(out).toEqual({ ...creds, plan_type: 'plus' }) + }) + it('trims the value', () => { + expect(applyPlanType({}, ' pro ')).toEqual({ plan_type: 'pro' }) + }) + it('deletes the key when cleared (empty), keeping other keys', () => { + const out = applyPlanType({ plan_type: 'pro', email: 'a@b.c' }, '') + expect(out).toEqual({ email: 'a@b.c' }) + expect('plan_type' in out).toBe(false) + }) + }) +}) + diff --git a/frontend/src/components/account/credentialsBuilder.ts b/frontend/src/components/account/credentialsBuilder.ts index 3cdc0bdd54..e78cf41de6 100644 --- a/frontend/src/components/account/credentialsBuilder.ts +++ b/frontend/src/components/account/credentialsBuilder.ts @@ -201,3 +201,87 @@ export function applyHeaderOverride( delete credentials[HEADER_OVERRIDES_CREDENTIAL_KEY] } } + +// ===== OpenAI plan_type (ChatGPT 订阅档位) 手动覆盖 ===== + +export interface PlanTypeOption { + value: string + label: string + // 兼容 common/Select.vue 的 SelectOption(含索引签名) + [key: string]: unknown +} + +/** + * plan_type 值的友好显示标签，镜像 PlatformTypeBadge 的映射 + * （canonical 值 chatgptpro 显示为 Pro，team 显示为 Team）。未知值原样返回。 + */ +export function planTypeDisplayLabel(value: string): string { + switch (value.trim().toLowerCase()) { + case 'plus': + return 'Plus' + case 'pro': + case 'chatgptpro': + return 'Pro' + case 'free': + return 'Free' + case 'team': + return 'Team' + default: + return value + } +} + +/** + * 从凭据里读取 plan_type，仅接受字符串（脏数据 42/true 等一律视为空， + * 避免被当作合法自定义项保留）。 + */ +export function readPlanType(credentials: Record | undefined | null): string { + const v = credentials?.plan_type + return typeof v === 'string' ? v : '' +} + +/** + * 构建 plan_type 下拉选项：清空 + Plus/Pro/Free 预设。 + * 若当前值是某预设的别名（如 chatgptpro↔Pro），用当前的 canonical 值占据该 + * 标签位（保留 canonical，显示友好标签，避免重复项）；若是完全预设外的值 + * （如 team 或异常值），追加为一项，避免编辑时下拉丢失原值。 + */ +export function buildPlanTypeOptions(current: string, clearLabel: string): PlanTypeOption[] { + const cur = (current || '').trim() + const curLabel = cur ? planTypeDisplayLabel(cur) : '' + const presets: PlanTypeOption[] = [ + { value: 'plus', label: 'Plus' }, + { value: 'pro', label: 'Pro' }, + { value: 'free', label: 'Free' } + ] + const opts: PlanTypeOption[] = [{ value: '', label: clearLabel }] + for (const p of presets) { + if (cur && p.value !== cur.toLowerCase() && p.label === curLabel) { + // 当前值是该预设的别名：用 canonical 当前值占位，标签仍显示友好名 + opts.push({ value: cur, label: p.label }) + } else { + opts.push(p) + } + } + if (cur && !opts.some(o => o.value.toLowerCase() === cur.toLowerCase())) { + opts.push({ value: cur, label: planTypeDisplayLabel(cur) }) + } + return opts +} + +/** + * 把手动选择的 plan_type 写入凭据：非空则设置，空则删除该键（清空/自动识别）。 + * 直接修改传入对象并返回。 + */ +export function applyPlanType( + credentials: Record, + planType: string +): Record { + const pt = (planType || '').trim() + if (pt) { + credentials.plan_type = pt + } else { + delete credentials.plan_type + } + return credentials +} diff --git a/frontend/src/components/admin/account/AccountTestModal.vue b/frontend/src/components/admin/account/AccountTestModal.vue index 0a0e3dd9ae..0a8f853ebb 100644 --- a/frontend/src/components/admin/account/AccountTestModal.vue +++ b/frontend/src/components/admin/account/AccountTestModal.vue @@ -250,6 +250,7 @@ import TextArea from '@/components/common/TextArea.vue' import { Icon } from '@/components/icons' import { useClipboard } from '@/composables/useClipboard' import { buildApiUrl } from '@/api/client' +import { ADMIN_UI_REQUEST_HEADER } from '@/api/adminUIRequest' import { adminAPI } from '@/api/admin' import type { Account, ClaudeModel } from '@/types' @@ -438,7 +439,8 @@ const startTest = async () => { method: 'POST', headers: { Authorization: `Bearer ${localStorage.getItem('auth_token')}`, - 'Content-Type': 'application/json' + 'Content-Type': 'application/json', + [ADMIN_UI_REQUEST_HEADER]: '1' }, body: JSON.stringify(requestBody), signal: abortController.signal diff --git a/frontend/src/components/admin/monitor/MonitorAdvancedRequestConfig.vue b/frontend/src/components/admin/monitor/MonitorAdvancedRequestConfig.vue index 404b691692..c4ccdfa3ac 100644 --- a/frontend/src/components/admin/monitor/MonitorAdvancedRequestConfig.vue +++ b/frontend/src/components/admin/monitor/MonitorAdvancedRequestConfig.vue @@ -109,6 +109,8 @@ import { useI18n } from 'vue-i18n' import type { APIMode, BodyOverrideMode, Provider } from '@/api/admin/channelMonitor' import { API_MODE_RESPONSES, + DEFAULT_GROK_MODEL, + PROVIDER_GROK, PROVIDER_OPENAI, } from '@/constants/channelMonitor' @@ -305,11 +307,12 @@ const bodyPlaceholder = computed(() => { } return '{\n "model": "gpt-4o-mini",\n "instructions": "You are a health check endpoint. Reply briefly.",\n "input": "Reply with exactly: ok",\n "max_output_tokens": 20,\n "stream": false\n}' } - if (props.provider === PROVIDER_OPENAI) { + if (props.provider === PROVIDER_OPENAI || props.provider === PROVIDER_GROK) { if (props.bodyOverrideMode === 'merge') { return '{\n "max_tokens": 20\n}' } - return '{\n "model": "gpt-4o-mini",\n "messages": [{"role":"user","content":"Reply with exactly: ok"}],\n "max_tokens": 20,\n "stream": false\n}' + const model = props.provider === PROVIDER_GROK ? DEFAULT_GROK_MODEL : 'gpt-4o-mini' + return `{\n "model": "${model}",\n "messages": [{"role":"user","content":"Reply with exactly: ok"}],\n "max_tokens": 20,\n "stream": false\n}` } if (props.bodyOverrideMode === 'merge') { return '{\n "system": "You are Claude Code..."\n}' diff --git a/frontend/src/components/admin/monitor/MonitorFiltersBar.vue b/frontend/src/components/admin/monitor/MonitorFiltersBar.vue index eb2a5c7857..544238f49c 100644 --- a/frontend/src/components/admin/monitor/MonitorFiltersBar.vue +++ b/frontend/src/components/admin/monitor/MonitorFiltersBar.vue @@ -70,6 +70,7 @@ import { PROVIDER_OPENAI, PROVIDER_ANTHROPIC, PROVIDER_GEMINI, + PROVIDER_GROK, } from '@/constants/channelMonitor' defineProps<{ @@ -94,6 +95,7 @@ const providerFilterOptions = computed(() => [ { value: PROVIDER_OPENAI, label: t('monitorCommon.providers.openai') }, { value: PROVIDER_ANTHROPIC, label: t('monitorCommon.providers.anthropic') }, { value: PROVIDER_GEMINI, label: t('monitorCommon.providers.gemini') }, + { value: PROVIDER_GROK, label: t('monitorCommon.providers.grok') }, ]) const enabledFilterOptions = computed(() => [ diff --git a/frontend/src/components/admin/monitor/MonitorFormDialog.vue b/frontend/src/components/admin/monitor/MonitorFormDialog.vue index e6cab8edf1..14a9e2dd15 100644 --- a/frontend/src/components/admin/monitor/MonitorFormDialog.vue +++ b/frontend/src/components/admin/monitor/MonitorFormDialog.vue @@ -13,15 +13,16 @@
-
+
@@ -80,6 +81,7 @@ (() => [ { value: PROVIDER_ANTHROPIC, label: t('monitorCommon.providers.anthropic') }, { value: PROVIDER_OPENAI, label: t('monitorCommon.providers.openai') }, { value: PROVIDER_GEMINI, label: t('monitorCommon.providers.gemini') }, + { value: PROVIDER_GROK, label: t('monitorCommon.providers.grok') }, ]) +function selectProvider(provider: Provider) { + if (form.provider === provider) return + const previousProvider = form.provider + const clearGrokEndpoint = + previousProvider === PROVIDER_GROK && form.endpoint === DEFAULT_GROK_ENDPOINT + const clearGrokModel = + previousProvider === PROVIDER_GROK && form.primary_model === DEFAULT_GROK_MODEL + form.provider = provider + if (provider === PROVIDER_GROK) { + if (!form.endpoint.trim()) form.endpoint = DEFAULT_GROK_ENDPOINT + if (!form.primary_model.trim()) form.primary_model = DEFAULT_GROK_MODEL + return + } + if (clearGrokEndpoint) form.endpoint = '' + if (clearGrokModel) form.primary_model = '' +} + // Clear api_key whenever provider changes to avoid cross-provider key mismatch. // Editing mode loads api_key='' via loadFromMonitor and only sets it on user // typing, so clearing on provider change is always a safe no-op until the user diff --git a/frontend/src/components/admin/monitor/MonitorTemplateManagerDialog.vue b/frontend/src/components/admin/monitor/MonitorTemplateManagerDialog.vue index e54ecf631d..63b87db75a 100644 --- a/frontend/src/components/admin/monitor/MonitorTemplateManagerDialog.vue +++ b/frontend/src/components/admin/monitor/MonitorTemplateManagerDialog.vue @@ -7,7 +7,7 @@ >
-
+