diff --git a/backend/cmd/server/VERSION b/backend/cmd/server/VERSION index 7b9dfc4dbf97..56ebc9e5b0e6 100644 --- a/backend/cmd/server/VERSION +++ b/backend/cmd/server/VERSION @@ -1 +1 @@ -0.1.132 +0.1.133 diff --git a/backend/cmd/server/wire.go b/backend/cmd/server/wire.go index 9bfa27174db6..b474cfa1aad5 100644 --- a/backend/cmd/server/wire.go +++ b/backend/cmd/server/wire.go @@ -98,6 +98,7 @@ func provideCleanup( backupSvc *service.BackupService, paymentOrderExpiry *service.PaymentOrderExpiryService, channelMonitorRunner *service.ChannelMonitorRunner, + quotaFlusher *service.UserPlatformQuotaUsageFlusher, ) func() { return func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -246,6 +247,12 @@ func provideCleanup( } return nil }}, + {"UserPlatformQuotaUsageFlusher", func() error { + if quotaFlusher != nil { + quotaFlusher.Stop() + } + return nil + }}, } infraSteps := []cleanupStep{ diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 465f5e25a712..106432158a79 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -137,7 +137,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { httpUpstream := repository.NewHTTPUpstream(configConfig) deferredService := service.ProvideDeferredService(accountRepository, timingWheelService) openAIOAuthClient := repository.NewOpenAIOAuthClient() - openAIOAuthService := service.NewOpenAIOAuthService(proxyRepository, openAIOAuthClient) + openAIOAuthService := service.ProvideOpenAIOAuthService(proxyRepository, openAIOAuthClient, privacyClientFactory) oAuthRefreshAPI := service.ProvideOAuthRefreshAPI(accountRepository, geminiTokenCache) openAITokenProvider := service.ProvideOpenAITokenProvider(accountRepository, geminiTokenCache, openAIOAuthService, oAuthRefreshAPI) channelRepository := repository.NewChannelRepository(db) @@ -195,7 +195,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { gatewayService := service.NewGatewayService(accountRepository, groupRepository, usageLogRepository, usageBillingRepository, userRepository, userSubscriptionRepository, userGroupRateRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, identityService, httpUpstream, deferredService, claudeTokenProvider, sessionLimitCache, rpmCache, digestSessionStore, settingService, tlsFingerprintProfileService, channelService, modelPricingResolver, balanceNotifyService, serviceUserPlatformQuotaRepository) geminiMessagesCompatService := service.NewGeminiMessagesCompatService(accountRepository, groupRepository, gatewayCache, schedulerSnapshotService, geminiTokenProvider, rateLimitService, httpUpstream, antigravityGatewayService, configConfig) opsSystemLogSink := service.ProvideOpsSystemLogSink(opsRepository) - opsService := service.NewOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink) + opsService := service.ProvideOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink, settingService) encryptionKey, err := payment.ProvideEncryptionKey(configConfig) if err != nil { return nil, err @@ -269,7 +269,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { scheduledTestRunnerService := service.ProvideScheduledTestRunnerService(scheduledTestPlanRepository, scheduledTestService, accountTestService, rateLimitService, configConfig) paymentOrderExpiryService := service.ProvidePaymentOrderExpiryService(paymentService) channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService) - v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, schedulerSnapshotService, tokenRefreshService, accountExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner) + userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService) + v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, schedulerSnapshotService, tokenRefreshService, accountExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher) application := &Application{ Server: httpServer, Cleanup: v, @@ -324,6 +325,7 @@ func provideCleanup( backupSvc *service.BackupService, paymentOrderExpiry *service.PaymentOrderExpiryService, channelMonitorRunner *service.ChannelMonitorRunner, + quotaFlusher *service.UserPlatformQuotaUsageFlusher, ) func() { return func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -471,6 +473,12 @@ func provideCleanup( } return nil }}, + {"UserPlatformQuotaUsageFlusher", func() error { + if quotaFlusher != nil { + quotaFlusher.Stop() + } + return nil + }}, } infraSteps := []cleanupStep{ diff --git a/backend/cmd/server/wire_gen_test.go b/backend/cmd/server/wire_gen_test.go index a44b2d5caacd..7f4e47731765 100644 --- a/backend/cmd/server/wire_gen_test.go +++ b/backend/cmd/server/wire_gen_test.go @@ -77,6 +77,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) { nil, // backupSvc nil, // paymentOrderExpiry nil, // channelMonitorRunner + nil, // quotaFlusher ) require.NotPanics(t, func() { diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index f689c2a908e6..0df7d09e84a4 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -649,6 +649,9 @@ type BillingConfig struct { // - billing_cache_service.checkUserPlatformQuotaEligibility 首次缓存装载 // 读写两端必须共用同一 TTL，避免缓存生命周期不一致导致 quota 计数漂移。 UserPlatformQuotaCacheTTLSeconds int `mapstructure:"user_platform_quota_cache_ttl_seconds"` + // UserPlatformQuotaSentinelTTLSeconds sentinel(无 limit 占位)entry 的 TTL, + // 显著短于 quota cache 默认 86400s 以控 Redis 内存;默认 3600=1h。 + UserPlatformQuotaSentinelTTLSeconds int `mapstructure:"user_platform_quota_sentinel_ttl_seconds"` } type CircuitBreakerConfig struct { @@ -882,6 +885,12 @@ type GatewayOpenAIWSConfig struct { StoreDisabledForceNewConn bool `mapstructure:"store_disabled_force_new_conn"` // PrewarmGenerateEnabled: 是否启用 WSv2 generate=false 预热（默认 false） PrewarmGenerateEnabled bool `mapstructure:"prewarm_generate_enabled"` + // ClientReadLimitBytes: 入站客户端 WS 单帧读取上限。 + ClientReadLimitBytes int64 `mapstructure:"client_read_limit_bytes"` + // HTTPBridgeEnabled: 首包过大时，保持客户端 WS，改用 HTTP Responses 上游。 + HTTPBridgeEnabled bool `mapstructure:"http_bridge_enabled"` + // HTTPBridgeThresholdBytes: 触发 HTTP bridge 的入站 WS payload 阈值。 + HTTPBridgeThresholdBytes int64 `mapstructure:"http_bridge_threshold_bytes"` // Feature 开关：v2 优先于 v1 ResponsesWebsockets bool `mapstructure:"responses_websockets"` @@ -1091,6 +1100,13 @@ type DatabaseConfig struct { ConnMaxLifetimeMinutes int `mapstructure:"conn_max_lifetime_minutes"` // ConnMaxIdleTimeMinutes: 空闲连接最大存活时间，及时释放不活跃连接 ConnMaxIdleTimeMinutes int `mapstructure:"conn_max_idle_time_minutes"` + // UserPlatformQuotaFlusherEnabled: 是否启用 user×platform 配额写聚合 flusher + UserPlatformQuotaFlusherEnabled bool `mapstructure:"user_platform_quota_flusher_enabled"` + // UserPlatformQuotaFlushIntervalMs: flusher 刷写间隔（毫秒） + UserPlatformQuotaFlushIntervalMs int `mapstructure:"user_platform_quota_flush_interval_ms"` + // UserPlatformQuotaFlushBatchSize: flusher 单批最大条数 + // 建议 ≤ 6000（单条 UPSERT 原子上限） + UserPlatformQuotaFlushBatchSize int `mapstructure:"user_platform_quota_flush_batch_size"` } func (d *DatabaseConfig) DSN() string { @@ -1571,6 +1587,7 @@ func setDefaults() { viper.SetDefault("billing.circuit_breaker.reset_timeout_seconds", 30) viper.SetDefault("billing.circuit_breaker.half_open_requests", 3) viper.SetDefault("billing.user_platform_quota_cache_ttl_seconds", 86400) + viper.SetDefault("billing.user_platform_quota_sentinel_ttl_seconds", 3600) // Turnstile viper.SetDefault("turnstile.required", false) @@ -1657,6 +1674,9 @@ func setDefaults() { viper.SetDefault("database.max_idle_conns", 128) viper.SetDefault("database.conn_max_lifetime_minutes", 30) viper.SetDefault("database.conn_max_idle_time_minutes", 5) + viper.SetDefault("database.user_platform_quota_flusher_enabled", false) + viper.SetDefault("database.user_platform_quota_flush_interval_ms", 2000) + viper.SetDefault("database.user_platform_quota_flush_batch_size", 1000) // Redis viper.SetDefault("redis.host", "localhost") @@ -1792,6 +1812,9 @@ func setDefaults() { viper.SetDefault("gateway.openai_ws.store_disabled_conn_mode", "strict") viper.SetDefault("gateway.openai_ws.store_disabled_force_new_conn", true) viper.SetDefault("gateway.openai_ws.prewarm_generate_enabled", false) + viper.SetDefault("gateway.openai_ws.client_read_limit_bytes", 64*1024*1024) + viper.SetDefault("gateway.openai_ws.http_bridge_enabled", true) + viper.SetDefault("gateway.openai_ws.http_bridge_threshold_bytes", 15*1024*1024) viper.SetDefault("gateway.openai_ws.responses_websockets", false) viper.SetDefault("gateway.openai_ws.responses_websockets_v2", true) viper.SetDefault("gateway.openai_ws.max_conns_per_account", 128) @@ -2529,6 +2552,15 @@ func (c *Config) Validate() error { if c.Gateway.OpenAIWS.PrewarmCooldownMS < 0 { return fmt.Errorf("gateway.openai_ws.prewarm_cooldown_ms must be non-negative") } + if c.Gateway.OpenAIWS.ClientReadLimitBytes <= 0 { + return fmt.Errorf("gateway.openai_ws.client_read_limit_bytes must be positive") + } + if c.Gateway.OpenAIWS.HTTPBridgeThresholdBytes < 0 { + return fmt.Errorf("gateway.openai_ws.http_bridge_threshold_bytes must be non-negative") + } + if c.Gateway.OpenAIWS.HTTPBridgeEnabled && c.Gateway.OpenAIWS.HTTPBridgeThresholdBytes == 0 { + return fmt.Errorf("gateway.openai_ws.http_bridge_threshold_bytes must be positive when http_bridge_enabled is true") + } if c.Gateway.OpenAIWS.FallbackCooldownSeconds < 0 { return fmt.Errorf("gateway.openai_ws.fallback_cooldown_seconds must be non-negative") } diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index 1eae5ed95953..9478b5102dae 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -134,6 +134,15 @@ func TestLoadDefaultOpenAIWSConfig(t *testing.T) { if cfg.Gateway.OpenAIWS.PrewarmCooldownMS != 300 { t.Fatalf("Gateway.OpenAIWS.PrewarmCooldownMS = %d, want 300", cfg.Gateway.OpenAIWS.PrewarmCooldownMS) } + if cfg.Gateway.OpenAIWS.ClientReadLimitBytes != 64*1024*1024 { + t.Fatalf("Gateway.OpenAIWS.ClientReadLimitBytes = %d, want %d", cfg.Gateway.OpenAIWS.ClientReadLimitBytes, 64*1024*1024) + } + if !cfg.Gateway.OpenAIWS.HTTPBridgeEnabled { + t.Fatalf("Gateway.OpenAIWS.HTTPBridgeEnabled = false, want true") + } + if cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes != 15*1024*1024 { + t.Fatalf("Gateway.OpenAIWS.HTTPBridgeThresholdBytes = %d, want %d", cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes, 15*1024*1024) + } if cfg.Gateway.OpenAIWS.RetryBackoffInitialMS != 120 { t.Fatalf("Gateway.OpenAIWS.RetryBackoffInitialMS = %d, want 120", cfg.Gateway.OpenAIWS.RetryBackoffInitialMS) } diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index 4f566a8be9f0..5719534230b6 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -2131,6 +2131,56 @@ func (h *AccountHandler) SyncUpstreamModels(c *gin.Context) { response.Success(c, gin.H{"models": models}) } +// SyncUpstreamModelsPreview handles syncing live supported models using provided credentials (no account ID needed). +// POST /api/v1/admin/accounts/models/sync-upstream-preview +func (h *AccountHandler) SyncUpstreamModelsPreview(c *gin.Context) { + var req struct { + Platform string `json:"platform" binding:"required"` + Type string `json:"type" binding:"required"` + BaseURL string `json:"base_url"` + APIKey string `json:"api_key" binding:"required"` + } + if err := c.ShouldBindJSON(&req); err != nil { + response.BadRequest(c, "Invalid request: "+err.Error()) + return + } + + tempAccount := &service.Account{ + Platform: req.Platform, + Type: req.Type, + Credentials: map[string]any{ + "api_key": req.APIKey, + "base_url": req.BaseURL, + }, + } + + if h.accountTestService == nil { + response.InternalError(c, "Account test service is not configured") + return + } + + models, err := h.accountTestService.FetchUpstreamSupportedModels(c.Request.Context(), tempAccount) + if err != nil { + var syncErr *service.UpstreamModelSyncError + if errors.As(err, &syncErr) { + switch syncErr.Kind { + case service.UpstreamModelSyncErrorConfiguration, service.UpstreamModelSyncErrorUnsupported: + response.BadRequest(c, syncErr.SafeMessage()) + default: + slog.Warn("sync_upstream_models_preview_failed", "platform", req.Platform, "kind", syncErr.Kind) + response.Error(c, http.StatusBadGateway, syncErr.SafeMessage()) + } + return + } + + slog.Warn("sync_upstream_models_preview_failed", "platform", req.Platform) + response.Error(c, http.StatusBadGateway, "Failed to sync upstream models from upstream") + return + } + + response.Success(c, gin.H{"models": models}) +} + // SetPrivacy handles setting privacy for a single OpenAI/Antigravity OAuth account // POST /api/v1/admin/accounts/:id/set-privacy func (h *AccountHandler) SetPrivacy(c *gin.Context) { diff --git a/backend/internal/handler/admin/admin_service_stub_test.go b/backend/internal/handler/admin/admin_service_stub_test.go index fd0ec459ec4c..819f0cdc02a1 100644 --- a/backend/internal/handler/admin/admin_service_stub_test.go +++ b/backend/internal/handler/admin/admin_service_stub_test.go @@ -160,6 +160,10 @@ func (s *stubAdminService) GetUser(ctx context.Context, id int64) (*service.User return &user, nil } +func (s *stubAdminService) GetUserIncludeDeleted(ctx context.Context, id int64) (*service.User, error) { + return s.GetUser(ctx, id) +} + func (s *stubAdminService) CreateUser(ctx context.Context, input *service.CreateUserInput) (*service.User, error) { user := service.User{ID: 100, Email: input.Email, Status: service.StatusActive} return &user, nil diff --git a/backend/internal/handler/admin/usage_handler.go b/backend/internal/handler/admin/usage_handler.go index 0857a13880b3..11a4aeb8f601 100644 --- a/backend/internal/handler/admin/usage_handler.go +++ b/backend/internal/handler/admin/usage_handler.go @@ -325,10 +325,24 @@ func (h *UsageHandler) Stats(c *gin.Context) { EndTime: &endTime, } - stats, err := h.usageService.GetStatsWithFilters(c.Request.Context(), filters) - if err != nil { - response.ErrorFrom(c, err) - return + var stats *usagestats.UsageStats + // nocache: 绕过缓存直接回源,刷新者本人拿最新;不回写缓存(管理台"我刷新我自己拿最新"语义,非全局失效)。 + if parseBoolQueryWithDefault(c.Query("nocache"), false) { + s, err := h.usageService.GetStatsWithFilters(c.Request.Context(), filters) + if err != nil { + response.ErrorFrom(c, err) + return + } + stats = s + c.Header("X-Usage-Stats-Cache", "bypass") + } else { + s, hit, err := h.getStatsCached(c.Request.Context(), filters) + if err != nil { + response.ErrorFrom(c, err) + return + } + stats = s + c.Header("X-Usage-Stats-Cache", cacheStatusValue(hit)) } response.Success(c, stats) @@ -344,23 +358,25 @@ func (h *UsageHandler) SearchUsers(c *gin.Context) { } // Limit to 30 results - users, _, err := h.adminService.ListUsers(c.Request.Context(), 1, 30, service.UserListFilters{Search: keyword}, "email", "asc") + users, _, err := h.adminService.ListUsers(c.Request.Context(), 1, 30, service.UserListFilters{Search: keyword, IncludeDeleted: true}, "email", "asc") if err != nil { response.ErrorFrom(c, err) return } - // Return simplified user list (only id and email) + // Return simplified user list (only id, email and deleted flag) type SimpleUser struct { - ID int64 `json:"id"` - Email string `json:"email"` + ID int64 `json:"id"` + Email string `json:"email"` + Deleted bool `json:"deleted"` } result := make([]SimpleUser, len(users)) for i, u := range users { result[i] = SimpleUser{ - ID: u.ID, - Email: u.Email, + ID: u.ID, + Email: u.Email, + Deleted: u.DeletedAt != nil, } } diff --git a/backend/internal/handler/admin/usage_handler_search_users_test.go b/backend/internal/handler/admin/usage_handler_search_users_test.go new file mode 100644 index 000000000000..ca4350122abc --- /dev/null +++ b/backend/internal/handler/admin/usage_handler_search_users_test.go @@ -0,0 +1,56 @@ +package admin + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// 捕获 ListUsers 入参、返回一个已删用户的 admin service 桩。 +type searchUsersAdminStub struct { + service.AdminService + gotFilters service.UserListFilters +} + +func (s *searchUsersAdminStub) ListUsers(ctx context.Context, page, pageSize int, filters service.UserListFilters, sortBy, sortOrder string) ([]service.User, int64, error) { + s.gotFilters = filters + ts := time.Date(2026, 5, 28, 0, 0, 0, 0, time.UTC) + return []service.User{ + {ID: 1, Email: "active@test.com"}, + {ID: 2, Email: "deleted@test.com", DeletedAt: &ts}, + }, 2, nil +} + +func TestAdminUsageSearchUsers_IncludesDeletedAndFlags(t *testing.T) { + gin.SetMode(gin.TestMode) + stub := &searchUsersAdminStub{} + handler := NewUsageHandler(nil, nil, stub, nil) + router := gin.New() + router.GET("/admin/usage/search-users", handler.SearchUsers) + + req := httptest.NewRequest(http.MethodGet, "/admin/usage/search-users?q=test", nil) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.True(t, stub.gotFilters.IncludeDeleted, "SearchUsers 必须请求 IncludeDeleted") + + var resp struct { + Data []struct { + ID int64 `json:"id"` + Email string `json:"email"` + Deleted bool `json:"deleted"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + require.Len(t, resp.Data, 2) + require.False(t, resp.Data[0].Deleted) + require.True(t, resp.Data[1].Deleted, "已删用户必须标记 deleted=true") +} diff --git a/backend/internal/handler/admin/usage_query_cache.go b/backend/internal/handler/admin/usage_query_cache.go new file mode 100644 index 000000000000..b288a95ba445 --- /dev/null +++ b/backend/internal/handler/admin/usage_query_cache.go @@ -0,0 +1,62 @@ +package admin + +import ( + "context" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" +) + +// 与 dashboard 查询缓存同款:30s TTL 进程内缓存,仅服务 /admin/usage/stats 读路径。 +var usageStatsCache = newSnapshotCache(30 * time.Second) + +type usageStatsCacheKeyData struct { + StartTime string `json:"start_time"` + EndTime string `json:"end_time"` + UserID int64 `json:"user_id"` + APIKeyID int64 `json:"api_key_id"` + AccountID int64 `json:"account_id"` + GroupID int64 `json:"group_id"` + Model string `json:"model"` + BillingMode string `json:"billing_mode"` + RequestType *int16 `json:"request_type"` + Stream *bool `json:"stream"` + BillingType *int8 `json:"billing_type"` +} + +func usageStatsCacheKey(filters usagestats.UsageLogFilters) string { + start := "" + if filters.StartTime != nil { + start = filters.StartTime.UTC().Format(time.RFC3339) + } + end := "" + if filters.EndTime != nil { + end = filters.EndTime.UTC().Format(time.RFC3339) + } + return mustMarshalDashboardCacheKey(usageStatsCacheKeyData{ + StartTime: start, + EndTime: end, + UserID: filters.UserID, + APIKeyID: filters.APIKeyID, + AccountID: filters.AccountID, + GroupID: filters.GroupID, + Model: filters.Model, + BillingMode: filters.BillingMode, + RequestType: filters.RequestType, + Stream: filters.Stream, + BillingType: filters.BillingType, + }) +} + +// getStatsCached 命中则返回缓存,未命中则回源 usageService 并写缓存。 +func (h *UsageHandler) getStatsCached(ctx context.Context, filters usagestats.UsageLogFilters) (*usagestats.UsageStats, bool, error) { + key := usageStatsCacheKey(filters) + entry, hit, err := usageStatsCache.GetOrLoad(key, func() (any, error) { + return h.usageService.GetStatsWithFilters(ctx, filters) + }) + if err != nil { + return nil, hit, err + } + stats, err := snapshotPayloadAs[*usagestats.UsageStats](entry.Payload) + return stats, hit, err +} diff --git a/backend/internal/handler/admin/usage_query_cache_test.go b/backend/internal/handler/admin/usage_query_cache_test.go new file mode 100644 index 000000000000..857e507a58dc --- /dev/null +++ b/backend/internal/handler/admin/usage_query_cache_test.go @@ -0,0 +1,28 @@ +package admin + +import ( + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/stretchr/testify/require" +) + +func TestUsageStatsCacheKey_StableAndDistinct(t *testing.T) { + start := time.Date(2026, 5, 29, 0, 0, 0, 0, time.UTC) + end := time.Date(2026, 5, 31, 0, 0, 0, 0, time.UTC) + base := usagestats.UsageLogFilters{StartTime: &start, EndTime: &end, Model: "claude-3"} + + k1 := usageStatsCacheKey(base) + k2 := usageStatsCacheKey(base) + require.NotEmpty(t, k1) + require.Equal(t, k1, k2, "same filters must produce same key") + + other := base + other.Model = "gpt-4o" + require.NotEqual(t, k1, usageStatsCacheKey(other), "different model must change key") + + withUser := base + withUser.UserID = 7 + require.NotEqual(t, k1, usageStatsCacheKey(withUser), "different user must change key") +} diff --git a/backend/internal/handler/admin/user_handler.go b/backend/internal/handler/admin/user_handler.go index 32a216929e87..a21fe55a9381 100644 --- a/backend/internal/handler/admin/user_handler.go +++ b/backend/internal/handler/admin/user_handler.go @@ -48,14 +48,14 @@ func NewUserHandler( // CreateUserRequest represents admin create user request type CreateUserRequest struct { - Email string `json:"email" binding:"required,email"` - Password string `json:"password" binding:"required,min=6"` - Username string `json:"username"` - Notes string `json:"notes"` - Balance float64 `json:"balance"` - Concurrency int `json:"concurrency"` - RPMLimit int `json:"rpm_limit"` - AllowedGroups []int64 `json:"allowed_groups"` + Email string `json:"email" binding:"required,email"` + Password string `json:"password" binding:"required,min=6"` + Username string `json:"username"` + Notes string `json:"notes"` + Balance *float64 `json:"balance"` + Concurrency int `json:"concurrency"` + RPMLimit int `json:"rpm_limit"` + AllowedGroups []int64 `json:"allowed_groups"` } // UpdateUserRequest represents admin update user request @@ -195,7 +195,12 @@ func (h *UserHandler) GetByID(c *gin.Context) { return } - user, err := h.adminService.GetUser(c.Request.Context(), userID) + var user *service.User + if c.Query("include_deleted") == "true" { + user, err = h.adminService.GetUserIncludeDeleted(c.Request.Context(), userID) + } else { + user, err = h.adminService.GetUser(c.Request.Context(), userID) + } if err != nil { response.ErrorFrom(c, err) return @@ -743,7 +748,7 @@ func (h *UserHandler) UpdateUserPlatformQuotas(c *gin.Context) { if h.billingCache != nil { for _, p := range service.AllowedQuotaPlatforms { if err := h.billingCache.DeleteUserPlatformQuotaCache(ctx, userID, p); err != nil { - slog.Warn("quota cache invalidation failed", "user_id", userID, "platform", p, "err", err) + slog.Error("ALERT: quota cache invalidation failed after UpsertForUser; limit 生效可能延迟至 sentinel TTL(最长 1h),需人工确认或重试失效", "user_id", userID, "platform", p, "err", err) } } } @@ -827,7 +832,7 @@ func (h *UserHandler) ResetUserPlatformQuotaWindow(c *gin.Context) { if h.billingCache != nil { if err := h.billingCache.DeleteUserPlatformQuotaCache(ctx, userID, req.Platform); err != nil { - slog.Warn("quota cache invalidation failed", "user_id", userID, "platform", req.Platform, "err", err) + slog.Error("ALERT: quota cache invalidation failed after ResetExpiredWindow; 窗口重置可能延迟至 sentinel TTL(最长 1h)", "user_id", userID, "platform", req.Platform, "err", err) } } diff --git a/backend/internal/handler/admin/user_handler_get_deleted_test.go b/backend/internal/handler/admin/user_handler_get_deleted_test.go new file mode 100644 index 000000000000..1b3070cde9ba --- /dev/null +++ b/backend/internal/handler/admin/user_handler_get_deleted_test.go @@ -0,0 +1,51 @@ +package admin + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type getByIDAdminStub struct { + service.AdminService +} + +func (s *getByIDAdminStub) GetUser(_ context.Context, _ int64) (*service.User, error) { + return nil, service.ErrUserNotFound +} + +func (s *getByIDAdminStub) GetUserIncludeDeleted(_ context.Context, id int64) (*service.User, error) { + return &service.User{ID: id, Email: "del@test.com"}, nil +} + +func setupGetByIDRouter(svc service.AdminService) *gin.Engine { + gin.SetMode(gin.TestMode) + r := gin.New() + h := NewUserHandler(svc, nil, nil, nil) + r.GET("/admin/users/:id", h.GetByID) + return r +} + +func TestAdminUserGetByID_IncludeDeleted(t *testing.T) { + svc := &getByIDAdminStub{AdminService: newStubAdminService()} + router := setupGetByIDRouter(svc) + + t.Run("normal path returns 404 for deleted user", func(t *testing.T) { + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/admin/users/7", nil) + router.ServeHTTP(w, req) + require.Equal(t, http.StatusNotFound, w.Code) + }) + + t.Run("include_deleted=true returns 200", func(t *testing.T) { + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/admin/users/7?include_deleted=true", nil) + router.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + }) +} diff --git a/backend/internal/handler/auth_oauth_pending_flow_test.go b/backend/internal/handler/auth_oauth_pending_flow_test.go index 70fb160a30a1..2f8f4e58023e 100644 --- a/backend/internal/handler/auth_oauth_pending_flow_test.go +++ b/backend/internal/handler/auth_oauth_pending_flow_test.go @@ -2914,6 +2914,10 @@ func (r *oauthPendingFlowUserRepo) DisableTotp(ctx context.Context, userID int64 Exec(ctx) } +func (r *oauthPendingFlowUserRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) { + return r.GetByID(ctx, id) +} + func oauthPendingFlowServiceUser(entity *dbent.User) *service.User { if entity == nil { return nil diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index 65052217022d..19e1d2283c2b 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -30,6 +30,7 @@ func UserFromServiceShallow(u *service.User) *User { BalanceNotifyExtraEmails: NotifyEmailEntriesFromService(u.BalanceNotifyExtraEmails), TotalRecharged: u.TotalRecharged, RPMLimit: u.RPMLimit, + DeletedAt: u.DeletedAt, } } diff --git a/backend/internal/handler/dto/mappers_deleted_user_test.go b/backend/internal/handler/dto/mappers_deleted_user_test.go new file mode 100644 index 000000000000..8ce5388e040f --- /dev/null +++ b/backend/internal/handler/dto/mappers_deleted_user_test.go @@ -0,0 +1,20 @@ +package dto + +import ( + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestUserFromServiceShallow_MapsDeletedAt(t *testing.T) { + ts := time.Date(2026, 5, 28, 10, 0, 0, 0, time.UTC) + + deleted := UserFromServiceShallow(&service.User{ID: 1, Email: "d@test.com", DeletedAt: &ts}) + require.NotNil(t, deleted.DeletedAt) + require.Equal(t, ts, *deleted.DeletedAt) + + active := UserFromServiceShallow(&service.User{ID: 2, Email: "a@test.com"}) + require.Nil(t, active.DeletedAt, "active user must have nil DeletedAt") +} diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index 6b0b6bcd5b9d..d41fd7577f58 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -20,6 +20,7 @@ type User struct { LastActiveAt *time.Time `json:"last_active_at,omitempty"` CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` + DeletedAt *time.Time `json:"deleted_at,omitempty"` // 余额不足通知 BalanceNotifyEnabled bool `json:"balance_notify_enabled"` diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index a24611f91c72..8853bffbe6a1 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -154,7 +154,8 @@ func (h *GatewayHandler) Messages(c *gin.Context) { setOpsRequestContext(c, "", false) - parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic) + bodyRef := service.NewRequestBodyRef(body) + parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic) if err != nil { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return @@ -440,7 +441,17 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 记录 Forward 前已写入字节数，Forward 后若增加则说明 SSE 内容已发，禁止 failover writerSizeBeforeForward := c.Writer.Size() if account.Platform == service.PlatformAntigravity { - result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, reqModel, "generateContent", reqStream, body, hasBoundSession) + result, err = h.antigravityGatewayService.ForwardGemini( + requestCtx, + c, + account, + reqModel, + "generateContent", + reqStream, + body, + hasBoundSession, + service.WithForwardGeminiSession(derefGroupID(apiKey.GroupID), sessionKey), + ) } else { result, err = h.geminiCompatService.Forward(requestCtx, c, account, body) } @@ -509,11 +520,12 @@ func (h *GatewayHandler) Messages(c *gin.Context) { } // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 + // ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。 + forceCacheBilling := fs.ForceCacheBilling quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, - ParsedRequest: parsedReq, QuotaPlatform: quotaPlatform, APIKey: apiKey, User: apiKey.User, @@ -524,7 +536,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { UserAgent: userAgent, IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, - ForceCacheBilling: fs.ForceCacheBilling, + ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), }); err != nil { @@ -562,6 +574,12 @@ func (h *GatewayHandler) Messages(c *gin.Context) { retryWithFallback := false for { + attemptParsedReq, err := parsedReq.CloneForBody(body) + if err != nil { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return + } + // 选择支持该模型的账号 reqLog.Info("sticky.selecting_account", zap.String("session_key", sessionKey), @@ -693,7 +711,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // ===== 用户消息串行队列 START ===== var queueRelease func() - umqMode := h.getUserMsgQueueMode(account, parsedReq) + umqMode := h.getUserMsgQueueMode(account, attemptParsedReq) switch umqMode { case config.UMQModeSerialize: @@ -740,20 +758,26 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 用 wrapReleaseOnDone 确保 context 取消时自动释放（仅 serialize 模式有 queueRelease） queueRelease = wrapReleaseOnDone(c.Request.Context(), queueRelease) // 注入回调到 ParsedRequest：使用外层 wrapper 以便提前清理 AfterFunc - parsedReq.OnUpstreamAccepted = queueRelease + attemptParsedReq.OnUpstreamAccepted = queueRelease // ===== 用户消息串行队列 END ===== - // 应用渠道模型映射到请求 + // 渠道模型映射只作用于本次账号尝试，避免 failover 后污染原始 ParsedRequest。 if channelMapping.Mapped { - parsedReq.Model = channelMapping.MappedModel - parsedReq.Body = h.gatewayService.ReplaceModelInBody(parsedReq.Body, channelMapping.MappedModel) + attemptParsedReq.Model = channelMapping.MappedModel + if err := attemptParsedReq.ReplaceBody(h.gatewayService.ReplaceModelInBody(attemptParsedReq.Body.Bytes(), channelMapping.MappedModel)); err != nil { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return + } } // Bedrock CC 兼容：渠道模型映射后，清理 Anthropic API 专有字段、注入 Bedrock 必需字段 - parsedReq.Body = h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), parsedReq.Body, parsedReq.Model, account, apiKey.GroupID) - body = parsedReq.Body + if err := attemptParsedReq.ReplaceBody(h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), attemptParsedReq.Body.Bytes(), attemptParsedReq.Model, account, apiKey.GroupID)); err != nil { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return + } + attemptBody := attemptParsedReq.Body.Bytes() // 转发请求 - 根据账号平台分流 - c.Set("parsed_request", parsedReq) + c.Set("parsed_request", attemptParsedReq) var result *service.ForwardResult requestCtx := c.Request.Context() if fs.SwitchCount > 0 { @@ -762,9 +786,9 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 记录 Forward 前已写入字节数，Forward 后若增加则说明 SSE 内容已发，禁止 failover writerSizeBeforeForward := c.Writer.Size() if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey { - result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, body, hasBoundSession) + result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, attemptBody, hasBoundSession) } else { - result, err = h.gatewayService.Forward(requestCtx, c, account, parsedReq) + result, err = h.gatewayService.Forward(requestCtx, c, account, attemptParsedReq) } // 兜底释放串行锁（正常情况已通过回调提前释放） @@ -772,7 +796,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { queueRelease() } // 清理回调引用，防止 failover 重试时旧回调被错误调用 - parsedReq.OnUpstreamAccepted = nil + attemptParsedReq.OnUpstreamAccepted = nil if accountReleaseFunc != nil { accountReleaseFunc() @@ -895,20 +919,22 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context） userAgent := c.GetHeader("User-Agent") clientIP := ip.GetClientIP(c) - requestPayloadHash := service.HashUsageRequestPayload(body) + // Forward 内部可能继续改写 body，usage 去重指纹必须使用最终上游接受的当前 body。 + requestPayloadHash := service.HashUsageRequestPayload(attemptParsedReq.Body.Bytes()) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) if result.ReasoningEffort == nil { - result.ReasoningEffort = service.NormalizeClaudeOutputEffort(parsedReq.OutputEffort) + result.ReasoningEffort = service.NormalizeClaudeOutputEffort(attemptParsedReq.OutputEffort) } // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 + // ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。 + forceCacheBilling := fs.ForceCacheBilling quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, - ParsedRequest: parsedReq, QuotaPlatform: quotaPlatform, APIKey: currentAPIKey, User: currentAPIKey.User, @@ -919,7 +945,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { UserAgent: userAgent, IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, - ForceCacheBilling: fs.ForceCacheBilling, + ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), }); err != nil { @@ -1683,7 +1709,8 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) { setOpsRequestContext(c, "", false) - parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic) + bodyRef := service.NewRequestBodyRef(body) + parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic) if err != nil { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return @@ -2056,10 +2083,11 @@ func (h *GatewayHandler) maybeLogCompatibilityFallbackMetrics(reqLog *zap.Logger ) } -func (h *GatewayHandler) submitUsageRecordTask(task service.UsageRecordTask) { +func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task service.UsageRecordTask) { if task == nil { return } + task = wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { h.usageRecordWorkerPool.Submit(task) return diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index acbdc261d206..719700aa9f63 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -151,9 +151,10 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { } // Parse request for session hash - parsedReq, _ := service.ParseGatewayRequest(body, "chat_completions") + bodyRef := service.NewRequestBodyRef(body) + parsedReq, _ := service.ParseGatewayRequest(bodyRef, "chat_completions") if parsedReq == nil { - parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: body} + parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: bodyRef} } parsedReq.SessionContext = &service.SessionContext{ ClientIP: ip.GetClientIP(c), @@ -292,7 +293,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, QuotaPlatform: quotaPlatform, diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index 6a083f317b91..49f80d19c149 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -156,9 +156,10 @@ func (h *GatewayHandler) Responses(c *gin.Context) { } // Parse request for session hash - parsedReq, _ := service.ParseGatewayRequest(body, "responses") + bodyRef := service.NewRequestBodyRef(body) + parsedReq, _ := service.ParseGatewayRequest(bodyRef, "responses") if parsedReq == nil { - parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: body} + parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: bodyRef} } parsedReq.SessionContext = &service.SessionContext{ ClientIP: ip.GetClientIP(c), @@ -267,7 +268,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, QuotaPlatform: quotaPlatform, diff --git a/backend/internal/handler/gateway_helper.go b/backend/internal/handler/gateway_helper.go index e48975027830..4b6a47eb2fa7 100644 --- a/backend/internal/handler/gateway_helper.go +++ b/backend/internal/handler/gateway_helper.go @@ -18,18 +18,12 @@ import ( // claudeCodeValidator is a singleton validator for Claude Code client detection var claudeCodeValidator = service.NewClaudeCodeValidator() -const claudeCodeParsedRequestContextKey = "claude_code_parsed_request" - // SetClaudeCodeClientContext 检查请求是否来自 Claude Code 客户端，并设置到 context 中 // 返回更新后的 context func SetClaudeCodeClientContext(c *gin.Context, body []byte, parsedReq *service.ParsedRequest) { if c == nil || c.Request == nil { return } - if parsedReq != nil { - c.Set(claudeCodeParsedRequestContextKey, parsedReq) - } - ua := c.GetHeader("User-Agent") // Fast path：非 Claude CLI UA 直接判定 false，避免热路径二次 JSON 反序列化。 if !claudeCodeValidator.ValidateUserAgent(ua) { @@ -45,9 +39,6 @@ func SetClaudeCodeClientContext(c *gin.Context, body []byte, parsedReq *service. } else { // 仅在确认为 Claude CLI 且 messages 路径时再做 body 解析。 bodyMap := claudeCodeBodyMapFromParsedRequest(parsedReq) - if bodyMap == nil { - bodyMap = claudeCodeBodyMapFromContextCache(c) - } if bodyMap == nil && len(body) > 0 { _ = json.Unmarshal(body, &bodyMap) } @@ -74,8 +65,12 @@ func claudeCodeBodyMapFromParsedRequest(parsedReq *service.ParsedRequest) map[st bodyMap := map[string]any{ "model": parsedReq.Model, } - if parsedReq.System != nil || parsedReq.HasSystem { - bodyMap["system"] = parsedReq.System + if parsedReq.HasSystem { + if system, ok := parsedReq.SystemValue(); ok { + bodyMap["system"] = system + } else { + bodyMap["system"] = nil + } } if parsedReq.MetadataUserID != "" { bodyMap["metadata"] = map[string]any{"user_id": parsedReq.MetadataUserID} @@ -83,26 +78,6 @@ func claudeCodeBodyMapFromParsedRequest(parsedReq *service.ParsedRequest) map[st return bodyMap } -func claudeCodeBodyMapFromContextCache(c *gin.Context) map[string]any { - if c == nil { - return nil - } - if cached, ok := c.Get(service.OpenAIParsedRequestBodyKey); ok { - if bodyMap, ok := cached.(map[string]any); ok { - return bodyMap - } - } - if cached, ok := c.Get(claudeCodeParsedRequestContextKey); ok { - switch v := cached.(type) { - case *service.ParsedRequest: - return claudeCodeBodyMapFromParsedRequest(v) - case service.ParsedRequest: - return claudeCodeBodyMapFromParsedRequest(&v) - } - } - return nil -} - // 并发槽位等待相关常量 // // 性能优化说明： diff --git a/backend/internal/handler/gateway_helper_hotpath_test.go b/backend/internal/handler/gateway_helper_hotpath_test.go index d57c396cc1e4..a6b6a429e9b7 100644 --- a/backend/internal/handler/gateway_helper_hotpath_test.go +++ b/backend/internal/handler/gateway_helper_hotpath_test.go @@ -177,7 +177,7 @@ func TestSetClaudeCodeClientContext_FastPathAndStrictPath(t *testing.T) { }) } -func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing.T) { +func TestSetClaudeCodeClientContext_ReuseParsedRequest(t *testing.T) { t.Run("reuse parsed request without body unmarshal", func(t *testing.T) { c, _ := newHelperTestContext(http.MethodPost, "/v1/messages") c.Request.Header.Set("User-Agent", "claude-cli/1.0.1") @@ -185,36 +185,13 @@ func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing c.Request.Header.Set("anthropic-beta", "message-batches-2024-09-24") c.Request.Header.Set("anthropic-version", "2023-06-01") - parsedReq := &service.ParsedRequest{ - Model: "claude-3-5-sonnet-20241022", - System: []any{ - map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."}, - }, - MetadataUserID: "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", - } + parsedReq, err := service.ParseGatewayRequest(service.NewRequestBodyRef(validClaudeCodeBodyJSON()), "") + require.NoError(t, err) // body 非法 JSON，如果函数复用 parsedReq 成功则仍应判定为 Claude Code。 SetClaudeCodeClientContext(c, []byte(`{invalid`), parsedReq) require.True(t, service.IsClaudeCodeClient(c.Request.Context())) }) - - t.Run("reuse context cache without body unmarshal", func(t *testing.T) { - c, _ := newHelperTestContext(http.MethodPost, "/v1/messages") - c.Request.Header.Set("User-Agent", "claude-cli/1.0.1") - c.Request.Header.Set("X-App", "claude-code") - c.Request.Header.Set("anthropic-beta", "message-batches-2024-09-24") - c.Request.Header.Set("anthropic-version", "2023-06-01") - c.Set(service.OpenAIParsedRequestBodyKey, map[string]any{ - "model": "claude-3-5-sonnet-20241022", - "system": []any{ - map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."}, - }, - "metadata": map[string]any{"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"}, - }) - - SetClaudeCodeClientContext(c, []byte(`{invalid`), nil) - require.True(t, service.IsClaudeCodeClient(c.Request.Context())) - }) } func TestWaitForSlotWithPingTimeout_AccountAndUserAcquire(t *testing.T) { diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 27ea44044f26..b7781eec5ddc 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -262,7 +262,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { sessionHash := extractGeminiCLISessionHash(c, body) if sessionHash == "" { // Fallback: 使用通用的会话哈希生成逻辑（适用于其他客户端） - parsedReq, _ := service.ParseGatewayRequest(body, domain.PlatformGemini) + parsedReq, _ := service.ParseGatewayRequest(service.NewRequestBodyRef(body), domain.PlatformGemini) if parsedReq != nil { parsedReq.SessionContext = &service.SessionContext{ ClientIP: ip.GetClientIP(c), @@ -477,8 +477,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { if fs.SwitchCount > 0 { requestCtx = service.WithAccountSwitchCount(requestCtx, fs.SwitchCount, h.metadataBridgeEnabled()) } + sessionGroupID := derefGroupID(apiKey.GroupID) if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey { - result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, modelName, action, stream, body, hasBoundSession) + result, err = h.antigravityGatewayService.ForwardGemini( + requestCtx, + c, + account, + modelName, + action, + stream, + body, + hasBoundSession, + service.WithForwardGeminiSession(sessionGroupID, sessionKey), + ) } else { result, err = h.geminiCompatService.ForwardNative(requestCtx, c, account, modelName, action, stream, body) } @@ -527,8 +538,10 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { requestPayloadHash := service.HashUsageRequestPayload(body) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) + // ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。 + forceCacheBilling := fs.ForceCacheBilling quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsageWithLongContext(ctx, &service.RecordUsageLongContextInput{ Result: result, QuotaPlatform: quotaPlatform, @@ -543,7 +556,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { RequestPayloadHash: requestPayloadHash, LongContextThreshold: 200000, // Gemini 200K 阈值 LongContextMultiplier: 2.0, // 超出部分双倍计费 - ForceCacheBilling: fs.ForceCacheBilling, + ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), }); err != nil { diff --git a/backend/internal/handler/openai_audio_transcriptions.go b/backend/internal/handler/openai_audio_transcriptions.go index 2b56a6fadad7..efd6c093b250 100644 --- a/backend/internal/handler/openai_audio_transcriptions.go +++ b/backend/internal/handler/openai_audio_transcriptions.go @@ -233,7 +233,7 @@ func (h *OpenAIGatewayHandler) AudioTranscriptions(c *gin.Context) { if result != nil { upstreamModel = result.UpstreamModel } - h.submitMandatoryUsageRecordTask(func(ctx context.Context) { + h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index 9805bf8a703b..d58656205d3d 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -274,7 +274,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveRawCCUpstreamEndpoint(c, account) - h.submitOpenAIUsageRecordTask(result, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, diff --git a/backend/internal/handler/openai_embeddings.go b/backend/internal/handler/openai_embeddings.go index 81713f7f9f48..20f9073510e7 100644 --- a/backend/internal/handler/openai_embeddings.go +++ b/backend/internal/handler/openai_embeddings.go @@ -214,7 +214,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) - h.submitOpenAIUsageRecordTask(result, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index de34afef0aee..e6d92e6c8c56 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -12,6 +12,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" "github.com/Wei-Shaw/sub2api/internal/pkg/ip" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" @@ -46,6 +47,31 @@ func resolveOpenAIMessagesDispatchMappedModel(apiKey *service.APIKey, requestedM return strings.TrimSpace(apiKey.Group.ResolveMessagesDispatchModel(requestedModel)) } +func usageRecordContext(parent context.Context, base context.Context) context.Context { + if base == nil { + base = context.Background() + } + if parent == nil { + return base + } + if clientRequestID, _ := parent.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" { + base = context.WithValue(base, ctxkey.ClientRequestID, strings.TrimSpace(clientRequestID)) + } + if requestID, _ := parent.Value(ctxkey.RequestID).(string); strings.TrimSpace(requestID) != "" { + base = context.WithValue(base, ctxkey.RequestID, strings.TrimSpace(requestID)) + } + return base +} + +func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) service.UsageRecordTask { + if task == nil { + return nil + } + return func(ctx context.Context) { + task(usageRecordContext(parent, ctx)) + } +} + // NewOpenAIGatewayHandler creates a new OpenAIGatewayHandler func NewOpenAIGatewayHandler( gatewayService *service.OpenAIGatewayService, @@ -438,7 +464,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 - h.submitOpenAIUsageRecordTask(result, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, @@ -823,7 +849,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) - h.submitOpenAIUsageRecordTask(result, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, @@ -922,19 +948,12 @@ func (h *OpenAIGatewayHandler) validateFunctionCallOutputRequest(c *gin.Context, return true } - var reqBody map[string]any - if err := json.Unmarshal(body, &reqBody); err != nil { - // 保持原有容错语义：解析失败时跳过预校验，沿用后续上游校验结果。 - return true - } - - c.Set(service.OpenAIParsedRequestBodyKey, reqBody) - validation := service.ValidateFunctionCallOutputContext(reqBody) + validation := service.ValidateFunctionCallOutputContextBytes(body) if !validation.HasFunctionCallOutput { return true } - previousResponseID, _ := reqBody["previous_response_id"].(string) + previousResponseID := gjson.GetBytes(body, "previous_response_id").String() if strings.TrimSpace(previousResponseID) != "" || validation.HasToolCallContext { return true } @@ -1170,7 +1189,7 @@ func (h *OpenAIGatewayHandler) openAIWebSocket(c *gin.Context, opts openAIWebSoc defer func() { _ = wsConn.CloseNow() }() - wsConn.SetReadLimit(16 * 1024 * 1024) + wsConn.SetReadLimit(service.ResolveOpenAIWSClientReadLimitBytes(h.cfg)) ctx := c.Request.Context() var ( @@ -1387,6 +1406,7 @@ func (h *OpenAIGatewayHandler) openAIWebSocket(c *gin.Context, opts openAIWebSoc zap.Int("candidate_count", scheduleDecision.CandidateCount), ) + var requestPayloadHash string hooks := &service.OpenAIWSIngressHooks{ InitialRequestModel: reqModel, BeforeRequest: func(turn int, payload []byte, originalModel string) error { @@ -1461,7 +1481,7 @@ func (h *OpenAIGatewayHandler) openAIWebSocket(c *gin.Context, opts openAIWebSoc h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) - h.submitOpenAIUsageRecordTask(result, func(taskCtx context.Context) { + h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) { if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, @@ -1472,7 +1492,7 @@ func (h *OpenAIGatewayHandler) openAIWebSocket(c *gin.Context, opts openAIWebSoc UpstreamEndpoint: upstreamEndpoint, UserAgent: userAgent, IPAddress: clientIP, - RequestPayloadHash: service.HashUsageRequestPayload(firstMessage), + RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, ChannelUsageFields: channelMappingWS.ToUsageFields(reqModel, result.UpstreamModel), }); err != nil { @@ -1492,6 +1512,9 @@ func (h *OpenAIGatewayHandler) openAIWebSocket(c *gin.Context, opts openAIWebSoc wsFirstMessage = h.gatewayService.ReplaceModelInBody(firstMessage, channelMappingWS.MappedModel) } + // WebSocket 首包可能很大，hash 必须在 hooks 外算成字符串，避免 AfterTurn 闭包保活请求体。 + requestPayloadHash = service.HashUsageRequestPayload(wsFirstMessage) + var proxyErr error if opts.Realtime { realtimeModel := reqModel @@ -1712,10 +1735,11 @@ func getContextInt64(c *gin.Context, key string) (int64, bool) { } } -func (h *OpenAIGatewayHandler) submitUsageRecordTask(task service.UsageRecordTask) { +func (h *OpenAIGatewayHandler) submitUsageRecordTask(parent context.Context, task service.UsageRecordTask) { if task == nil { return } + task = wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { h.usageRecordWorkerPool.Submit(task) return @@ -1734,18 +1758,19 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(task service.UsageRecordTas task(ctx) } -func (h *OpenAIGatewayHandler) submitOpenAIUsageRecordTask(result *service.OpenAIForwardResult, task service.UsageRecordTask) { +func (h *OpenAIGatewayHandler) submitOpenAIUsageRecordTask(parent context.Context, result *service.OpenAIForwardResult, task service.UsageRecordTask) { if result != nil && result.ImageCount > 0 { - h.submitMandatoryUsageRecordTask(task) + h.submitMandatoryUsageRecordTask(parent, task) return } - h.submitUsageRecordTask(task) + h.submitUsageRecordTask(parent, task) } -func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(task service.UsageRecordTask) { +func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(parent context.Context, task service.UsageRecordTask) { if task == nil { return } + task = wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDropped { return diff --git a/backend/internal/handler/openai_gateway_usage_context_test.go b/backend/internal/handler/openai_gateway_usage_context_test.go new file mode 100644 index 000000000000..7091c9c09944 --- /dev/null +++ b/backend/internal/handler/openai_gateway_usage_context_test.go @@ -0,0 +1,41 @@ +package handler + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/stretchr/testify/require" +) + +func TestSubmitUsageRecordTaskCopiesRequestContext(t *testing.T) { + parent := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-request-123") + parent = context.WithValue(parent, ctxkey.RequestID, "request-456") + + var gotClientRequestID string + var gotRequestID string + h := &GatewayHandler{} + h.submitUsageRecordTask(parent, func(ctx context.Context) { + gotClientRequestID, _ = ctx.Value(ctxkey.ClientRequestID).(string) + gotRequestID, _ = ctx.Value(ctxkey.RequestID).(string) + }) + + require.Equal(t, "client-request-123", gotClientRequestID) + require.Equal(t, "request-456", gotRequestID) +} + +func TestOpenAISubmitUsageRecordTaskCopiesRequestContext(t *testing.T) { + parent := context.WithValue(context.Background(), ctxkey.ClientRequestID, "openai-client-request-123") + parent = context.WithValue(parent, ctxkey.RequestID, "openai-request-456") + + var gotClientRequestID string + var gotRequestID string + h := &OpenAIGatewayHandler{} + h.submitUsageRecordTask(parent, func(ctx context.Context) { + gotClientRequestID, _ = ctx.Value(ctxkey.ClientRequestID).(string) + gotRequestID, _ = ctx.Value(ctxkey.RequestID).(string) + }) + + require.Equal(t, "openai-client-request-123", gotClientRequestID) + require.Equal(t, "openai-request-456", gotRequestID) +} diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index bbb080149c2b..1e3b530601e3 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -73,9 +73,10 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", err.Error()) return } + requestModel := parsed.Model reqLog = reqLog.With( - zap.String("model", parsed.Model), + zap.String("model", requestModel), zap.Bool("stream", parsed.Stream), zap.Bool("multipart", parsed.Multipart), zap.String("capability", string(parsed.RequiredCapability)), @@ -85,7 +86,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { h.errorResponse(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage()) return } - if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, parsed.Model, parsed.ModerationBody()); decision != nil && decision.Blocked { + if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, requestModel, parsed.ModerationBody()); decision != nil && decision.Blocked { h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message) return } @@ -98,13 +99,13 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { } if parsed.Multipart { - setOpsRequestContext(c, parsed.Model, parsed.Stream) + setOpsRequestContext(c, requestModel, parsed.Stream) } else { - setOpsRequestContext(c, parsed.Model, parsed.Stream) + setOpsRequestContext(c, requestModel, parsed.Stream) } setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(parsed.Stream, false))) - channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, parsed.Model) + channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, requestModel) if h.errorPassthroughService != nil { service.BindErrorPassthroughService(c, h.errorPassthroughService) @@ -147,7 +148,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { c.Request.Context(), apiKey.GroupID, sessionHash, - parsed.Model, + requestModel, failedAccountIDs, parsed.RequiredCapability, ) @@ -311,7 +312,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { if result != nil { upstreamModel = result.UpstreamModel } - h.submitMandatoryUsageRecordTask(func(ctx context.Context) { + h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: apiKey, @@ -324,14 +325,14 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, - ChannelUsageFields: channelMapping.ToUsageFields(parsed.Model, upstreamModel), + ChannelUsageFields: channelMapping.ToUsageFields(requestModel, upstreamModel), }); err != nil { logger.L().With( zap.String("component", "handler.openai_gateway.images"), zap.Int64("user_id", subject.UserID), zap.Int64("api_key_id", apiKey.ID), zap.Any("group_id", apiKey.GroupID), - zap.String("model", parsed.Model), + zap.String("model", requestModel), zap.Int64("account_id", account.ID), ).Error("openai.images.record_usage_failed", zap.Error(err)) } diff --git a/backend/internal/handler/usage_record_submit_task_test.go b/backend/internal/handler/usage_record_submit_task_test.go index e4c2837a592c..ebe5c3df8266 100644 --- a/backend/internal/handler/usage_record_submit_task_test.go +++ b/backend/internal/handler/usage_record_submit_task_test.go @@ -29,7 +29,7 @@ func TestGatewayHandlerSubmitUsageRecordTask_WithPool(t *testing.T) { h := &GatewayHandler{usageRecordWorkerPool: pool} done := make(chan struct{}) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { close(done) }) @@ -44,7 +44,7 @@ func TestGatewayHandlerSubmitUsageRecordTask_WithoutPoolSyncFallback(t *testing. h := &GatewayHandler{} var called atomic.Bool - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { if _, ok := ctx.Deadline(); !ok { t.Fatal("expected deadline in fallback context") } @@ -57,7 +57,7 @@ func TestGatewayHandlerSubmitUsageRecordTask_WithoutPoolSyncFallback(t *testing. func TestGatewayHandlerSubmitUsageRecordTask_NilTask(t *testing.T) { h := &GatewayHandler{} require.NotPanics(t, func() { - h.submitUsageRecordTask(nil) + h.submitUsageRecordTask(context.Background(), nil) }) } @@ -66,12 +66,12 @@ func TestGatewayHandlerSubmitUsageRecordTask_WithoutPool_TaskPanicRecovered(t *t var called atomic.Bool require.NotPanics(t, func() { - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { panic("usage task panic") }) }) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { called.Store(true) }) require.True(t, called.Load(), "panic 后后续任务应仍可执行") @@ -82,7 +82,7 @@ func TestOpenAIGatewayHandlerSubmitUsageRecordTask_WithPool(t *testing.T) { h := &OpenAIGatewayHandler{usageRecordWorkerPool: pool} done := make(chan struct{}) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { close(done) }) @@ -97,7 +97,7 @@ func TestOpenAIGatewayHandlerSubmitUsageRecordTask_WithoutPoolSyncFallback(t *te h := &OpenAIGatewayHandler{} var called atomic.Bool - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { if _, ok := ctx.Deadline(); !ok { t.Fatal("expected deadline in fallback context") } @@ -110,7 +110,7 @@ func TestOpenAIGatewayHandlerSubmitUsageRecordTask_WithoutPoolSyncFallback(t *te func TestOpenAIGatewayHandlerSubmitUsageRecordTask_NilTask(t *testing.T) { h := &OpenAIGatewayHandler{} require.NotPanics(t, func() { - h.submitUsageRecordTask(nil) + h.submitUsageRecordTask(context.Background(), nil) }) } @@ -119,12 +119,12 @@ func TestOpenAIGatewayHandlerSubmitUsageRecordTask_WithoutPool_TaskPanicRecovere var called atomic.Bool require.NotPanics(t, func() { - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { panic("usage task panic") }) }) - h.submitUsageRecordTask(func(ctx context.Context) { + h.submitUsageRecordTask(context.Background(), func(ctx context.Context) { called.Store(true) }) require.True(t, called.Load(), "panic 后后续任务应仍可执行") @@ -152,7 +152,7 @@ func TestOpenAIGatewayHandlerSubmitMandatoryUsageRecordTask_DroppedTaskSyncFallb pool.Submit(func(ctx context.Context) {}) var called atomic.Bool - h.submitMandatoryUsageRecordTask(func(ctx context.Context) { + h.submitMandatoryUsageRecordTask(context.Background(), func(ctx context.Context) { called.Store(true) }) close(release) @@ -182,7 +182,7 @@ func TestOpenAIGatewayHandlerSubmitOpenAIUsageRecordTask_ImageResultUsesMandator pool.Submit(func(ctx context.Context) {}) var called atomic.Bool - h.submitOpenAIUsageRecordTask(&service.OpenAIForwardResult{ImageCount: 1}, func(ctx context.Context) { + h.submitOpenAIUsageRecordTask(context.Background(), &service.OpenAIForwardResult{ImageCount: 1}, func(ctx context.Context) { called.Store(true) }) close(release) diff --git a/backend/internal/handler/user_handler_test.go b/backend/internal/handler/user_handler_test.go index 416478022a35..2e366c230892 100644 --- a/backend/internal/handler/user_handler_test.go +++ b/backend/internal/handler/user_handler_test.go @@ -118,6 +118,9 @@ func (s *userHandlerRepoStub) RemoveGroupFromUserAllowedGroups(context.Context, func (s *userHandlerRepoStub) UpdateTotpSecret(context.Context, int64, *string) error { return nil } func (s *userHandlerRepoStub) EnableTotp(context.Context, int64) error { return nil } func (s *userHandlerRepoStub) DisableTotp(context.Context, int64) error { return nil } +func (s *userHandlerRepoStub) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) { + return s.GetByID(ctx, id) +} func (s *userHandlerRepoStub) ListUserAuthIdentities(context.Context, int64) ([]service.UserAuthIdentityRecord, error) { out := make([]service.UserAuthIdentityRecord, len(s.identities)) copy(out, s.identities) diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go index 09b680c7c73c..cc51cba2d211 100644 --- a/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_bridge.go @@ -42,14 +42,24 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR return out, nil } +// responsesInputToChatMessages converts a Responses request's instructions + +// input[] into Chat Completions messages. It is a three-stage pipeline: +// +// parse — instructions become a system message; input[] is split into items +// build — buildChatMessagesFromItems walks items, attaching reasoning to the +// assistant message that produced a tool call, merging parallel tool +// calls into one assistant message, and skipping item types that have +// no Chat equivalent +// normalize — normalizeChatMessages enforces the invariants DeepSeek requires +// +// The build + normalize split keeps every protocol rule in one place rather than +// scattered across per-item cases, and makes unknown future codex item types +// fail safe instead of leaking into the upstream request. func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) ([]ChatMessage, error) { var messages []ChatMessage if strings.TrimSpace(instructions) != "" { content, _ := json.Marshal(instructions) - messages = append(messages, ChatMessage{ - Role: "system", - Content: content, - }) + messages = append(messages, ChatMessage{Role: "system", Content: content}) } inputRaw = bytesTrimSpace(inputRaw) @@ -57,13 +67,11 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) return messages, nil } + // Bare string input is a single user turn. var inputText string if err := json.Unmarshal(inputRaw, &inputText); err == nil { content, _ := json.Marshal(inputText) - messages = append(messages, ChatMessage{ - Role: "user", - Content: content, - }) + messages = append(messages, ChatMessage{Role: "user", Content: content}) return messages, nil } @@ -72,6 +80,24 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) return nil, fmt.Errorf("parse responses input: %w", err) } + built, err := buildChatMessagesFromItems(messages, rawItems) + if err != nil { + return nil, err + } + return normalizeChatMessages(built), nil +} + +// buildChatMessagesFromItems walks the Responses input items and appends the +// corresponding Chat messages. +func buildChatMessagesFromItems(messages []ChatMessage, rawItems []json.RawMessage) ([]ChatMessage, error) { + // pendingReasoning holds the reasoning text from a reasoning item until the + // assistant message it belongs to is emitted. DeepSeek's thinking mode + // requires the reasoning_content that produced a tool call to be passed back + // on that assistant message; dropping it yields a 400. It only survives + // across an assistant message (so a following tool call in the same turn + // still receives it); any other role ends the thinking span. + var pendingReasoning string + for _, raw := range rawItems { raw = bytesTrimSpace(raw) if len(raw) == 0 || string(raw) == "null" { @@ -84,6 +110,7 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) if textErr := json.Unmarshal(raw, &text); textErr == nil { content, _ := json.Marshal(text) messages = append(messages, ChatMessage{Role: "user", Content: content}) + pendingReasoning = "" continue } return nil, fmt.Errorf("parse responses input item: %w", err) @@ -92,22 +119,40 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) role := chatCompletionsBridgeRole(rawString(item["role"])) itemType := rawString(item["type"]) switch itemType { + case "reasoning": + if txt := extractResponsesReasoningText(item); txt != "" { + pendingReasoning = txt + } + continue case "function_call": arguments := rawString(item["arguments"]) if strings.TrimSpace(arguments) == "" { arguments = "{}" } - messages = append(messages, ChatMessage{ - Role: "assistant", - ToolCalls: []ChatToolCall{{ - ID: rawString(item["call_id"]), - Type: "function", - Function: ChatFunctionCall{ - Name: rawString(item["name"]), - Arguments: arguments, - }, - }}, - }) + toolCall := ChatToolCall{ + ID: rawString(item["call_id"]), + Type: "function", + Function: ChatFunctionCall{ + Name: rawString(item["name"]), + Arguments: arguments, + }, + } + // Parallel tool calls arrive as consecutive function_call items and + // must share one assistant message; the matching tool replies then + // follow it. Merge into the immediately preceding assistant message. + if n := len(messages); n > 0 && messages[n-1].Role == "assistant" { + messages[n-1].ToolCalls = append(messages[n-1].ToolCalls, toolCall) + if messages[n-1].ReasoningContent == "" { + messages[n-1].ReasoningContent = pendingReasoning + } + } else { + messages = append(messages, ChatMessage{ + Role: "assistant", + ToolCalls: []ChatToolCall{toolCall}, + ReasoningContent: pendingReasoning, + }) + } + pendingReasoning = "" continue case "function_call_output": content, _ := json.Marshal(rawString(item["output"])) @@ -116,10 +161,12 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) ToolCallID: rawString(item["call_id"]), Content: content, }) + pendingReasoning = "" continue case "input_text", "text": content, _ := json.Marshal(rawString(item["text"])) messages = append(messages, ChatMessage{Role: "user", Content: content}) + pendingReasoning = "" continue case "input_image": content, err := chatContentFromSingleResponsesPart(itemType, item) @@ -127,6 +174,18 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) return nil, err } messages = append(messages, ChatMessage{Role: "user", Content: content}) + pendingReasoning = "" + continue + } + + // Only genuine message items become chat messages. Codex emits other + // Responses item types with no Chat equivalent (web_search_call, + // local_shell_call, custom tool calls, file_search_call, ...). Converting + // them via the generic path would insert a spurious message between an + // assistant tool_calls message and its tool reply, which DeepSeek rejects + // ("insufficient tool messages following tool_calls message"). Skip them. + if itemType != "" && itemType != "message" { + pendingReasoning = "" continue } @@ -140,15 +199,128 @@ func responsesInputToChatMessages(instructions string, inputRaw json.RawMessage) if err != nil { return nil, err } - messages = append(messages, ChatMessage{ - Role: role, - Content: chatContent, - }) + messages = append(messages, ChatMessage{Role: role, Content: chatContent}) + // Reasoning only survives across an assistant text message. + if role != "assistant" { + pendingReasoning = "" + } } return messages, nil } +// normalizeChatMessages is the single place that enforces the tool-call +// invariant the DeepSeek / OpenAI Chat Completions schema requires: an assistant +// message with tool_calls must be immediately followed by one tool message per +// tool_call_id, in order, with nothing in between. +// +// Codex histories violate this in several ways that the builder alone can't fix: +// - a non-tool message lands between an assistant tool_calls message and its +// tool replies (e.g. an "Approved command prefix saved" system notice codex +// injects mid tool-execution); +// - a parallel tool_call's sibling output never arrives, or a call is left +// dangling by a mid-execution reconnect (unanswered tool_call); +// - a tool reply has no announcing assistant tool_call (orphan). +// +// It rebuilds the sequence so each assistant's answered tool_calls are followed +// directly by their replies (in call order); unanswered tool_calls are dropped +// (and an assistant left with neither tool_calls nor content is dropped); orphan +// tool replies and intervening messages are emitted in their natural position +// but never between an assistant tool_calls message and its replies. +func normalizeChatMessages(messages []ChatMessage) []ChatMessage { + // Index every tool reply by its tool_call_id (last wins on duplicates). + replies := make(map[string]ChatMessage) + for _, m := range messages { + if m.Role == "tool" && m.ToolCallID != "" { + replies[m.ToolCallID] = m + } + } + + out := make([]ChatMessage, 0, len(messages)) + for _, m := range messages { + switch { + case m.Role == "tool": + // A bare tool message with no tool_call_id is a direct Chat + // Completions passthrough; keep it in place. A tool reply whose id is + // announced by an assistant is emitted right after that assistant + // (skip the standalone occurrence). Any other tool reply is an orphan + // and is dropped. + if m.ToolCallID == "" { + out = append(out, m) + } + continue + case len(m.ToolCalls) > 0: + kept := make([]ChatToolCall, 0, len(m.ToolCalls)) + for _, tc := range m.ToolCalls { + if tc.ID == "" { + continue + } + if _, ok := replies[tc.ID]; ok { + kept = append(kept, tc) + } + } + if len(kept) == 0 { + // No answered tool_calls left: keep as a plain message if it has + // content, otherwise drop it entirely. + if isBlankChatContent(m.Content) { + continue + } + m.ToolCalls = nil + out = append(out, m) + continue + } + m.ToolCalls = kept + out = append(out, m) + for _, tc := range kept { + out = append(out, replies[tc.ID]) + } + default: + out = append(out, m) + } + } + return out +} + +// isBlankChatContent reports whether a chat message content holds no usable text. +func isBlankChatContent(raw json.RawMessage) bool { + raw = bytesTrimSpace(raw) + if len(raw) == 0 || string(raw) == "null" || string(raw) == `""` { + return true + } + return chatMessageContentText(raw) == "" +} + +// extractResponsesReasoningText pulls the reasoning text out of a Responses +// reasoning item. The Chat→Responses bridge writes the upstream reasoning_content +// verbatim into the summary_text parts (see closeChatReasoningItem), so codex +// round-trips it there; prefer summary[].text and fall back to content. +func extractResponsesReasoningText(item map[string]json.RawMessage) string { + var parts []string + collect := func(raw json.RawMessage) { + raw = bytesTrimSpace(raw) + if len(raw) == 0 || string(raw) == "null" { + return + } + var arr []map[string]json.RawMessage + if err := json.Unmarshal(raw, &arr); err == nil { + for _, p := range arr { + if t := rawString(p["text"]); t != "" { + parts = append(parts, t) + } + } + return + } + if t := rawString(raw); t != "" { + parts = append(parts, t) + } + } + collect(item["summary"]) + if len(parts) == 0 { + collect(item["content"]) + } + return strings.Join(parts, "\n") +} + func chatCompletionsBridgeRole(role string) string { trimmed := strings.TrimSpace(role) if trimmed == "" { @@ -448,10 +620,32 @@ type ChatCompletionsToResponsesStreamState struct { CreatedSent bool CompletedSent bool + // nextOutputIndex assigns sequential output_index values to items as they + // are opened (reasoning, message, tool calls), so the streamed indices match + // the order of items in the final response.output array. + nextOutputIndex int + + // Reasoning item lifecycle. DeepSeek-style upstreams stream all + // reasoning_content before any content, so reasoning is modeled as its own + // "reasoning" output item that must be opened (output_item.added) before any + // reasoning delta and closed before the message/tool items open. + ReasoningItemID string + ReasoningIndex int + ReasoningOpen bool + ReasoningDone bool + + // Message item + output_text content-part lifecycle. MessageItemID string - Text strings.Builder - Reasoning strings.Builder - ToolCalls map[int]*ChatToolCall + MessageIndex int + TextPartOpen bool + + Text strings.Builder + Reasoning strings.Builder + + // Tool-call lifecycle, keyed by the upstream tool_call index. + ToolCalls map[int]*ChatToolCall + ToolItemIDs map[int]string + ToolOutputIndex map[int]int FinishReason string Usage *ResponsesUsage @@ -460,13 +654,21 @@ type ChatCompletionsToResponsesStreamState struct { // NewChatCompletionsToResponsesStreamState returns an initialized stream state. func NewChatCompletionsToResponsesStreamState(model string) *ChatCompletionsToResponsesStreamState { return &ChatCompletionsToResponsesStreamState{ - ResponseID: generateResponsesID(), - Model: model, - Created: time.Now().Unix(), - ToolCalls: make(map[int]*ChatToolCall), + ResponseID: generateResponsesID(), + Model: model, + Created: time.Now().Unix(), + ToolCalls: make(map[int]*ChatToolCall), + ToolItemIDs: make(map[int]string), + ToolOutputIndex: make(map[int]int), } } +func (state *ChatCompletionsToResponsesStreamState) allocOutputIndex() int { + idx := state.nextOutputIndex + state.nextOutputIndex++ + return idx +} + // ChatCompletionsChunkToResponsesEvents converts one Chat Completions stream // chunk into zero or more Responses stream events. func ChatCompletionsChunkToResponsesEvents( @@ -490,24 +692,34 @@ func ChatCompletionsChunkToResponsesEvents( events = append(events, ensureChatToResponsesCreated(state)...) for _, choice := range chunk.Choices { - if choice.Delta.Content != nil { + // Reasoning is emitted as its own output item and must be opened + // (output_item.added + reasoning_summary_part.added) before the first + // delta, otherwise a strict client discards the delta. The leading + // empty-string reasoning delta upstreams send is filtered out. + if choice.Delta.ReasoningContent != nil && *choice.Delta.ReasoningContent != "" { + events = append(events, ensureChatReasoningItem(state)...) + _, _ = state.Reasoning.WriteString(*choice.Delta.ReasoningContent) + events = append(events, chatToResponsesEvent(state, "response.reasoning_summary_text.delta", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + SummaryIndex: 0, + Delta: *choice.Delta.ReasoningContent, + ItemID: state.ReasoningItemID, + })) + } + if choice.Delta.Content != nil && *choice.Delta.Content != "" { + // First real content closes the reasoning item, then opens the + // message item and its output_text content part. + events = append(events, closeChatReasoningItem(state)...) events = append(events, ensureChatToResponsesMessageItem(state)...) + events = append(events, ensureChatToResponsesTextPart(state)...) _, _ = state.Text.WriteString(*choice.Delta.Content) events = append(events, chatToResponsesEvent(state, "response.output_text.delta", &ResponsesStreamEvent{ - OutputIndex: 0, + OutputIndex: state.MessageIndex, ContentIndex: 0, Delta: *choice.Delta.Content, ItemID: state.MessageItemID, })) } - if choice.Delta.ReasoningContent != nil { - _, _ = state.Reasoning.WriteString(*choice.Delta.ReasoningContent) - events = append(events, chatToResponsesEvent(state, "response.reasoning_summary_text.delta", &ResponsesStreamEvent{ - OutputIndex: 0, - SummaryIndex: 0, - Delta: *choice.Delta.ReasoningContent, - })) - } for _, toolCall := range choice.Delta.ToolCalls { idx := 0 if toolCall.Index != nil { @@ -515,6 +727,8 @@ func ChatCompletionsChunkToResponsesEvents( } stored, ok := state.ToolCalls[idx] if !ok { + // A tool call closes any open reasoning item first. + events = append(events, closeChatReasoningItem(state)...) copyCall := toolCall if copyCall.ID == "" { copyCall.ID = generateItemID() @@ -522,11 +736,14 @@ func ChatCompletionsChunkToResponsesEvents( copyCall.Type = "function" state.ToolCalls[idx] = ©Call stored = ©Call + itemID := generateItemID() + state.ToolItemIDs[idx] = itemID + state.ToolOutputIndex[idx] = state.allocOutputIndex() events = append(events, chatToResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{ - OutputIndex: idx + 1, + OutputIndex: state.ToolOutputIndex[idx], Item: &ResponsesOutput{ Type: "function_call", - ID: generateItemID(), + ID: itemID, CallID: stored.ID, Name: stored.Function.Name, Status: "in_progress", @@ -543,7 +760,8 @@ func ChatCompletionsChunkToResponsesEvents( if toolCall.Function.Arguments != "" { stored.Function.Arguments += toolCall.Function.Arguments events = append(events, chatToResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{ - OutputIndex: idx + 1, + OutputIndex: state.ToolOutputIndex[idx], + ItemID: state.ToolItemIDs[idx], Delta: toolCall.Function.Arguments, CallID: stored.ID, Name: stored.Function.Name, @@ -565,24 +783,44 @@ func FinalizeChatCompletionsResponsesStream(state *ChatCompletionsToResponsesStr } var events []ResponsesStreamEvent events = append(events, ensureChatToResponsesCreated(state)...) + + // Close a reasoning item that never transitioned to content (reasoning-only + // or empty completion). + events = append(events, closeChatReasoningItem(state)...) + if state.MessageItemID != "" { - events = append(events, chatToResponsesEvent(state, "response.output_text.done", &ResponsesStreamEvent{ - OutputIndex: 0, - ContentIndex: 0, - Text: state.Text.String(), - ItemID: state.MessageItemID, - })) + if state.TextPartOpen { + events = append(events, chatToResponsesEvent(state, "response.output_text.done", &ResponsesStreamEvent{ + OutputIndex: state.MessageIndex, + ContentIndex: 0, + Text: state.Text.String(), + ItemID: state.MessageItemID, + })) + events = append(events, chatToResponsesEvent(state, "response.content_part.done", &ResponsesStreamEvent{ + OutputIndex: state.MessageIndex, + ContentIndex: 0, + ItemID: state.MessageItemID, + Part: &ResponsesContentPart{Type: "output_text", Text: state.Text.String()}, + })) + } events = append(events, chatToResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{ - OutputIndex: 0, + OutputIndex: state.MessageIndex, Item: &ResponsesOutput{ - Type: "message", - ID: state.MessageItemID, - Role: "assistant", - Status: "completed", + Type: "message", + ID: state.MessageItemID, + Role: "assistant", + Content: []ResponsesContentPart{{Type: "output_text", Text: state.Text.String()}}, + Status: "completed", }, })) } + // Close every function_call item opened during the stream. Codex finalizes a + // tool call only after function_call_arguments.done + output_item.done for + // that item; without them the call never completes and the session wedges. + // Mirrors cc-switch's finalize_tools. + events = append(events, closeChatToolItems(state)...) + status := "completed" var incompleteDetails *ResponsesIncompleteDetails if state.FinishReason == "length" { @@ -621,22 +859,142 @@ func ensureChatToResponsesCreated(state *ChatCompletionsToResponsesStreamState) })} } +// ensureChatReasoningItem opens the reasoning output item (output_item.added + +// reasoning_summary_part.added) before the first reasoning delta. Codex renders +// streaming reasoning only when this summary-part lifecycle is present. +func ensureChatReasoningItem(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { + if state.ReasoningOpen || state.ReasoningDone { + return nil + } + state.ReasoningOpen = true + state.ReasoningItemID = generateItemID() + state.ReasoningIndex = state.allocOutputIndex() + return []ResponsesStreamEvent{ + chatToResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + Item: &ResponsesOutput{Type: "reasoning", ID: state.ReasoningItemID, Status: "in_progress"}, + }), + chatToResponsesEvent(state, "response.reasoning_summary_part.added", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + SummaryIndex: 0, + ItemID: state.ReasoningItemID, + Part: &ResponsesContentPart{Type: "summary_text"}, + }), + } +} + +// closeChatReasoningItem emits the reasoning item's terminal events +// (reasoning_summary_text.done + reasoning_summary_part.done + output_item.done). +func closeChatReasoningItem(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { + if !state.ReasoningOpen { + return nil + } + state.ReasoningOpen = false + state.ReasoningDone = true + reasoning := state.Reasoning.String() + return []ResponsesStreamEvent{ + chatToResponsesEvent(state, "response.reasoning_summary_text.done", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + SummaryIndex: 0, + Text: reasoning, + ItemID: state.ReasoningItemID, + }), + chatToResponsesEvent(state, "response.reasoning_summary_part.done", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + SummaryIndex: 0, + ItemID: state.ReasoningItemID, + Part: &ResponsesContentPart{Type: "summary_text", Text: reasoning}, + }), + chatToResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{ + OutputIndex: state.ReasoningIndex, + Item: &ResponsesOutput{ + Type: "reasoning", + ID: state.ReasoningItemID, + Status: "completed", + Summary: []ResponsesSummary{{Type: "summary_text", Text: reasoning}}, + }, + }), + } +} + func ensureChatToResponsesMessageItem(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { if state.MessageItemID != "" { return nil } state.MessageItemID = generateItemID() + state.MessageIndex = state.allocOutputIndex() return []ResponsesStreamEvent{chatToResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{ - OutputIndex: 0, + OutputIndex: state.MessageIndex, Item: &ResponsesOutput{ - Type: "message", - ID: state.MessageItemID, - Role: "assistant", - Status: "in_progress", + Type: "message", + ID: state.MessageItemID, + Role: "assistant", + Status: "in_progress", + Content: []ResponsesContentPart{{Type: "output_text"}}, }, })} } +func ensureChatToResponsesTextPart(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { + if state.TextPartOpen { + return nil + } + state.TextPartOpen = true + return []ResponsesStreamEvent{chatToResponsesEvent(state, "response.content_part.added", &ResponsesStreamEvent{ + OutputIndex: state.MessageIndex, + ContentIndex: 0, + ItemID: state.MessageItemID, + Part: &ResponsesContentPart{Type: "output_text", Text: ""}, + })} +} + +// closeChatToolItems emits function_call_arguments.done + output_item.done for +// every tool call opened during the stream, carrying the full call_id/name/ +// arguments so codex can deserialize and execute the call. Mirrors cc-switch's +// finalize_tools. +func closeChatToolItems(state *ChatCompletionsToResponsesStreamState) []ResponsesStreamEvent { + if len(state.ToolCalls) == 0 { + return nil + } + var events []ResponsesStreamEvent + for i := 0; i < len(state.ToolCalls); i++ { + toolCall, ok := state.ToolCalls[i] + if !ok || toolCall == nil { + continue + } + itemID, opened := state.ToolItemIDs[i] + if !opened { + continue + } + arguments := toolCall.Function.Arguments + if strings.TrimSpace(arguments) == "" { + arguments = "{}" + } + outputIndex := state.ToolOutputIndex[i] + events = append(events, + chatToResponsesEvent(state, "response.function_call_arguments.done", &ResponsesStreamEvent{ + OutputIndex: outputIndex, + ItemID: itemID, + CallID: toolCall.ID, + Name: toolCall.Function.Name, + Arguments: arguments, + }), + chatToResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{ + OutputIndex: outputIndex, + Item: &ResponsesOutput{ + Type: "function_call", + ID: itemID, + CallID: toolCall.ID, + Name: toolCall.Function.Name, + Arguments: arguments, + Status: "completed", + }, + }), + ) + } + return events +} + func (state *ChatCompletionsToResponsesStreamState) chatOutput() []ResponsesOutput { var outputs []ResponsesOutput if state.Reasoning.Len() > 0 { diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_request_invariants_test.go b/backend/internal/pkg/apicompat/chatcompletions_responses_request_invariants_test.go new file mode 100644 index 000000000000..e54a453279d0 --- /dev/null +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_request_invariants_test.go @@ -0,0 +1,187 @@ +package apicompat + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +// assertChatInvariants enforces the DeepSeek / OpenAI Chat Completions message +// invariants that, when violated, surface as upstream 400s. Used to validate the +// request-direction converter against golden codex request shapes. +func assertChatInvariants(t *testing.T, messages []ChatMessage) { + t.Helper() + for i, m := range messages { + // Every assistant tool_calls message must be immediately followed by one + // tool message per tool_call_id, in order. + if len(m.ToolCalls) > 0 { + for j, tc := range m.ToolCalls { + k := i + 1 + j + require.Lessf(t, k, len(messages), "tool_call %s has no following tool message", tc.ID) + require.Equalf(t, "tool", messages[k].Role, "tool_call %s not followed by a tool message", tc.ID) + require.Equalf(t, tc.ID, messages[k].ToolCallID, "tool reply order mismatch for %s", tc.ID) + } + } + // No two consecutive assistant messages. + if i > 0 && m.Role == "assistant" && messages[i-1].Role == "assistant" { + t.Fatalf("consecutive assistant messages at %d", i) + } + // No orphan tool replies. + if m.Role == "tool" { + require.NotEmptyf(t, m.ToolCallID, "tool message without tool_call_id at %d", i) + } + } +} + +func convertGolden(t *testing.T, input string) []ChatMessage { + t.Helper() + msgs, err := responsesInputToChatMessages("You are a helpful assistant.", json.RawMessage(input)) + require.NoError(t, err) + return msgs +} + +// Golden sample: a single tool-call turn (codex runs one shell/curl command), +// the shape that produced the original "no response" / 400. +func TestGolden_SingleToolCall(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"latest sha?"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"need to run curl"}]}, + {"type":"function_call","call_id":"call_a","name":"exec_command","arguments":"{\"cmd\":\"curl x\"}"}, + {"type":"function_call_output","call_id":"call_a","output":"deadbeef"} + ]`) + assertChatInvariants(t, msgs) + // reasoning_content must ride on the assistant tool-call message. + var asst *ChatMessage + for i := range msgs { + if len(msgs[i].ToolCalls) > 0 { + asst = &msgs[i] + } + } + require.NotNil(t, asst) + require.Equal(t, "need to run curl", asst.ReasoningContent) +} + +// Golden sample: parallel tool calls (codex runs git log + git tag at once). +func TestGolden_ParallelToolCalls(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"features?"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"inspect repo"}]}, + {"type":"function_call","call_id":"c0","name":"exec_command","arguments":"{\"cmd\":\"git log\"}"}, + {"type":"function_call","call_id":"c1","name":"exec_command","arguments":"{\"cmd\":\"git tag\"}"}, + {"type":"function_call_output","call_id":"c0","output":"log"}, + {"type":"function_call_output","call_id":"c1","output":"tags"} + ]`) + assertChatInvariants(t, msgs) + // Both parallel calls share ONE assistant message. + var toolMsgs int + for _, m := range msgs { + if len(m.ToolCalls) == 2 { + require.Equal(t, "c0", m.ToolCalls[0].ID) + require.Equal(t, "c1", m.ToolCalls[1].ID) + } + if m.Role == "tool" { + toolMsgs++ + } + } + require.Equal(t, 2, toolMsgs) +} + +// Golden sample: an unknown item type (web_search_call from a 联网查询) sitting +// between a function_call and its output must not break tool↔reply adjacency. +func TestGolden_UnknownItemBetweenToolCallAndOutput(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"search"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"let me search"}]}, + {"type":"function_call","call_id":"c0","name":"exec_command","arguments":"{}"}, + {"type":"web_search_call","id":"ws_1","status":"completed","action":{"type":"search","query":"x"}}, + {"type":"function_call_output","call_id":"c0","output":"result"} + ]`) + assertChatInvariants(t, msgs) +} + +// Sequential tool calls (a tool reply between two calls) must stay in distinct +// assistant messages. +func TestRequest_SequentialToolCallsStaySeparate(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"function_call","call_id":"c1","name":"exec","arguments":"{}"}, + {"type":"function_call_output","call_id":"c1","output":"r1"}, + {"type":"function_call","call_id":"c2","name":"exec","arguments":"{}"}, + {"type":"function_call_output","call_id":"c2","output":"r2"} + ]`) + assertChatInvariants(t, msgs) + assistants := 0 + for _, m := range msgs { + if len(m.ToolCalls) == 1 { + assistants++ + } + } + require.Equal(t, 2, assistants) +} + +// Golden sample: codex injects a message (e.g. an "Approved command prefix +// saved" notice) between a function_call and its output. The intervening message +// must be moved after the tool reply so the assistant tool_calls is immediately +// followed by its reply. +func TestGolden_MessageBetweenToolCallAndOutput(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"do it"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"run cmd"}]}, + {"type":"function_call","call_id":"A","name":"exec","arguments":"{}"}, + {"type":"message","role":"developer","content":[{"type":"input_text","text":"Approved command prefix saved"}]}, + {"type":"function_call_output","call_id":"A","output":"ok"} + ]`) + assertChatInvariants(t, msgs) + // The assistant tool_calls message is immediately followed by its tool reply. + for i, m := range msgs { + if len(m.ToolCalls) > 0 { + require.Equal(t, "tool", msgs[i+1].Role) + require.Equal(t, "A", msgs[i+1].ToolCallID) + } + } +} + +// Golden sample: a parallel tool call where one sibling's output is missing +// (codex interrupted/reconnected mid-execution). The unanswered tool_call must +// be dropped so the remaining assistant tool_calls are all answered. +func TestGolden_PartialParallelDropsUnansweredCall(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"q"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"r"}]}, + {"type":"function_call","call_id":"A","name":"exec","arguments":"{}"}, + {"type":"function_call","call_id":"B","name":"exec","arguments":"{}"}, + {"type":"function_call_output","call_id":"A","output":"oa"} + ]`) + assertChatInvariants(t, msgs) + for _, m := range msgs { + for _, tc := range m.ToolCalls { + require.NotEqual(t, "B", tc.ID, "unanswered tool_call B should have been dropped") + } + } +} + +// Golden sample: a dangling tool_call at the end of the history (no output yet). +// The assistant message holding only that call must be dropped entirely. +func TestGolden_DanglingToolCallDropped(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"q"}]}, + {"type":"reasoning","summary":[{"type":"summary_text","text":"r"}]}, + {"type":"function_call","call_id":"A","name":"exec","arguments":"{}"} + ]`) + assertChatInvariants(t, msgs) + for _, m := range msgs { + require.Empty(t, m.ToolCalls, "dangling unanswered tool_call should have been dropped") + } +} + +// normalizeChatMessages drops an orphan tool reply whose tool_call was never +// announced. +func TestNormalize_DropsOrphanToolReply(t *testing.T) { + msgs := convertGolden(t, `[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"q"}]}, + {"type":"function_call_output","call_id":"ghost","output":"orphan"} + ]`) + for _, m := range msgs { + require.NotEqualf(t, "tool", m.Role, "orphan tool reply should have been dropped") + } +} diff --git a/backend/internal/pkg/apicompat/chatcompletions_responses_stream_lifecycle_test.go b/backend/internal/pkg/apicompat/chatcompletions_responses_stream_lifecycle_test.go new file mode 100644 index 000000000000..beb473038bc2 --- /dev/null +++ b/backend/internal/pkg/apicompat/chatcompletions_responses_stream_lifecycle_test.go @@ -0,0 +1,103 @@ +package apicompat + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func collectStreamEvents(t *testing.T, chunks []string) []ResponsesStreamEvent { + t.Helper() + state := NewChatCompletionsToResponsesStreamState("deepseek-v4-pro") + var events []ResponsesStreamEvent + for _, payload := range chunks { + var chunk ChatCompletionsChunk + require.NoError(t, json.Unmarshal([]byte(payload), &chunk)) + events = append(events, ChatCompletionsChunkToResponsesEvents(&chunk, state)...) + } + events = append(events, FinalizeChatCompletionsResponsesStream(state)...) + return events +} + +// TestStream_ReasoningOpensItemBeforeDelta guards the bug where a strict client +// (Codex) drops reasoning deltas that reference an item not yet opened. +func TestStream_ReasoningOpensItemBeforeDelta(t *testing.T) { + events := collectStreamEvents(t, []string{ + `{"choices":[{"index":0,"delta":{"role":"assistant","content":null,"reasoning_content":""}}]}`, + `{"choices":[{"index":0,"delta":{"reasoning_content":"think"}}]}`, + `{"choices":[{"index":0,"delta":{"content":"hello"}}]}`, + `{"choices":[{"index":0,"delta":{"content":""},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}`, + }) + + open := map[int]string{} // output_index -> item type + for _, e := range events { + switch e.Type { + case "response.output_item.added": + require.NotNil(t, e.Item) + open[e.OutputIndex] = e.Item.Type + case "response.reasoning_summary_text.delta": + require.Equalf(t, "reasoning", open[e.OutputIndex], "reasoning delta before its item was opened") + case "response.output_text.delta": + require.Equalf(t, "message", open[e.OutputIndex], "text delta before its item was opened") + } + } +} + +// TestStream_ToolCallLifecycleComplete guards that a tool call is fully closed +// (function_call_arguments.done + output_item.done with full arguments), which +// codex needs to execute the call. +func TestStream_ToolCallLifecycleComplete(t *testing.T) { + events := collectStreamEvents(t, []string{ + `{"choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"plan"}}]}`, + `{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_a","type":"function","function":{"name":"exec","arguments":""}}]}}]}`, + `{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"cmd\":\"ls\"}"}}]}}]}`, + `{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}`, + }) + + var sawAdded, sawArgsDone, sawItemDone bool + for _, e := range events { + switch e.Type { + case "response.output_item.added": + if e.Item != nil && e.Item.Type == "function_call" { + sawAdded = true + } + case "response.function_call_arguments.done": + sawArgsDone = true + require.Equal(t, `{"cmd":"ls"}`, e.Arguments) + case "response.output_item.done": + if e.Item != nil && e.Item.Type == "function_call" { + sawItemDone = true + require.Equal(t, `{"cmd":"ls"}`, e.Item.Arguments) + require.Equal(t, "call_a", e.Item.CallID) + } + } + } + require.True(t, sawAdded, "function_call output_item.added missing") + require.True(t, sawArgsDone, "function_call_arguments.done missing") + require.True(t, sawItemDone, "function_call output_item.done missing") +} + +// TestStream_SSEWireComplete drives the full stream through SSE encoding and +// asserts the function_call events carry complete fields on the wire. +func TestStream_SSEWireComplete(t *testing.T) { + events := collectStreamEvents(t, []string{ + `{"choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"plan"}}]}`, + `{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_a","type":"function","function":{"name":"exec","arguments":"{}"}}]}}]}`, + `{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`, + }) + + var addedLine string + for _, e := range events { + sse, err := ResponsesEventToSSE(e) + require.NoError(t, err) + if e.Type == "response.output_item.added" && e.Item != nil && e.Item.Type == "function_call" { + addedLine = sse + } + } + require.NotEmpty(t, addedLine) + // The function_call added event must carry arguments:"" on the wire. + require.True(t, strings.Contains(addedLine, `"arguments":""`), "added line missing arguments: %s", addedLine) + require.Contains(t, addedLine, `"call_id":"call_a"`) +} diff --git a/backend/internal/pkg/apicompat/responses_stream_event_wire.go b/backend/internal/pkg/apicompat/responses_stream_event_wire.go new file mode 100644 index 000000000000..df7a82e39374 --- /dev/null +++ b/backend/internal/pkg/apicompat/responses_stream_event_wire.go @@ -0,0 +1,199 @@ +package apicompat + +import "encoding/json" + +// MarshalJSON renders a ResponsesStreamEvent into its wire form. +// +// The OpenAI Responses streaming protocol requires several fields to be present +// even when they hold a zero value: output_index/content_index/summary_index are +// meaningful at 0, a function_call item must always carry call_id/name/arguments +// (arguments may be ""), a message item must carry content:[] and an output_text +// part must carry text/annotations/logprobs. Go's `omitempty` drops exactly those +// zero values, and strict clients (Codex CLI) reject items/deltas whose required +// fields are missing. +// +// Rather than marshalling with omitempty and patching the JSON afterwards, every +// streamed event type is constructed explicitly here — the Go analogue of the +// reference gateways' (cc-switch, CCX) per-event object construction. This is the +// single source of truth for Responses SSE field presence and applies uniformly +// to every emitter (Chat→Responses bridge and Anthropic→Responses converter). +// +// Event types not listed fall back to the default struct marshalling, which +// bounds the blast radius of this method to the streamed item/part/text/tool +// events. +func (e ResponsesStreamEvent) MarshalJSON() ([]byte, error) { + switch e.Type { + case "response.output_text.delta", "response.output_text.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + m["content_index"] = e.ContentIndex + if e.Type == "response.output_text.done" { + m["text"] = e.Text + } else { + m["delta"] = e.Delta + } + return json.Marshal(m) + + case "response.content_part.added", "response.content_part.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + m["content_index"] = e.ContentIndex + m["part"] = outputTextPartWire(e.Part) + return json.Marshal(m) + + case "response.reasoning_summary_text.delta", "response.reasoning_summary_text.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + m["summary_index"] = e.SummaryIndex + if e.Type == "response.reasoning_summary_text.done" { + m["text"] = e.Text + } else { + m["delta"] = e.Delta + } + return json.Marshal(m) + + case "response.reasoning_summary_part.added", "response.reasoning_summary_part.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + m["summary_index"] = e.SummaryIndex + m["part"] = summaryTextPartWire(e.Part) + return json.Marshal(m) + + case "response.output_item.added", "response.output_item.done": + m := e.wireBase() + m["output_index"] = e.OutputIndex + m["item"] = responsesItemWire(e.Item) + return json.Marshal(m) + + case "response.function_call_arguments.delta", "response.function_call_arguments.done": + m := e.wireBase() + e.putItemID(m) + m["output_index"] = e.OutputIndex + if e.CallID != "" { + m["call_id"] = e.CallID + } + if e.Name != "" { + m["name"] = e.Name + } + if e.Type == "response.function_call_arguments.done" { + m["arguments"] = e.Arguments + } else { + m["delta"] = e.Delta + } + return json.Marshal(m) + + default: + // response.created / completed / done / failed / incomplete and any + // event type not shaped above keep the default struct marshalling. + type alias ResponsesStreamEvent + return json.Marshal(alias(e)) + } +} + +func (e ResponsesStreamEvent) wireBase() map[string]any { + m := map[string]any{ + "type": e.Type, + "sequence_number": e.SequenceNumber, + } + return m +} + +func (e ResponsesStreamEvent) putItemID(m map[string]any) { + if e.ItemID != "" { + m["item_id"] = e.ItemID + } +} + +// outputTextPartWire renders a content part for a message's output_text, always +// carrying text/annotations/logprobs (matching cc-switch's push_text_delta). +func outputTextPartWire(part *ResponsesContentPart) map[string]any { + text := "" + if part != nil { + text = part.Text + } + return map[string]any{ + "type": "output_text", + "text": text, + "annotations": []any{}, + "logprobs": []any{}, + } +} + +// summaryTextPartWire renders a reasoning summary part. +func summaryTextPartWire(part *ResponsesContentPart) map[string]any { + text := "" + if part != nil { + text = part.Text + } + return map[string]any{ + "type": "summary_text", + "text": text, + } +} + +// responsesItemWire renders an output_item with every field the item's type +// requires to be present, including the empty arrays/strings that omitempty +// would otherwise drop. Mirrors cc-switch's response_function_call_item and the +// message/reasoning item shapes codex expects. +func responsesItemWire(item *ResponsesOutput) map[string]any { + if item == nil { + return map[string]any{} + } + m := map[string]any{ + "type": item.Type, + "id": item.ID, + } + if item.Status != "" { + m["status"] = item.Status + } + switch item.Type { + case "message": + role := item.Role + if role == "" { + role = "assistant" + } + m["role"] = role + m["content"] = messageContentWire(item.Content) + case "reasoning": + m["summary"] = reasoningSummaryWire(item.Summary) + if item.EncryptedContent != "" { + m["encrypted_content"] = item.EncryptedContent + } + case "function_call": + m["call_id"] = item.CallID + m["name"] = item.Name + m["arguments"] = item.Arguments + } + return m +} + +// messageContentWire renders a message item's content array; always an array +// (never null), with each output_text part carrying its text. +func messageContentWire(parts []ResponsesContentPart) []map[string]any { + out := make([]map[string]any, 0, len(parts)) + for _, p := range parts { + typ := p.Type + if typ == "" { + typ = "output_text" + } + out = append(out, map[string]any{"type": typ, "text": p.Text}) + } + return out +} + +// reasoningSummaryWire renders a reasoning item's summary array; always an array. +func reasoningSummaryWire(summary []ResponsesSummary) []map[string]any { + out := make([]map[string]any, 0, len(summary)) + for _, s := range summary { + typ := s.Type + if typ == "" { + typ = "summary_text" + } + out = append(out, map[string]any{"type": typ, "text": s.Text}) + } + return out +} diff --git a/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go b/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go new file mode 100644 index 000000000000..b4f6871d5f4f --- /dev/null +++ b/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go @@ -0,0 +1,113 @@ +package apicompat + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +// marshalEvent marshals through the custom MarshalJSON and returns the decoded +// object plus the set of top-level keys. +func marshalEvent(t *testing.T, e ResponsesStreamEvent) map[string]any { + t.Helper() + b, err := json.Marshal(e) + require.NoError(t, err) + var m map[string]any + require.NoError(t, json.Unmarshal(b, &m)) + return m +} + +// TestWire_IndexFieldsPresentAtZero guards the omitempty trap: output_index/ +// content_index/summary_index must serialize even when 0. +func TestWire_IndexFieldsPresentAtZero(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.output_text.delta", OutputIndex: 0, ContentIndex: 0, ItemID: "msg_1", Delta: "hi", + }) + require.Contains(t, m, "output_index") + require.Contains(t, m, "content_index") + require.EqualValues(t, 0, m["output_index"]) + + r := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.reasoning_summary_text.delta", OutputIndex: 0, SummaryIndex: 0, ItemID: "rs_1", Delta: "think", + }) + require.Contains(t, r, "output_index") + require.Contains(t, r, "summary_index") +} + +// TestWire_FunctionCallItemAlwaysComplete guards that a function_call item +// always carries call_id/name/arguments, including arguments:"" on .added. +func TestWire_FunctionCallItemAlwaysComplete(t *testing.T) { + added := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.output_item.added", + OutputIndex: 1, + Item: &ResponsesOutput{Type: "function_call", ID: "fc_1", CallID: "call_a", Name: "exec", Status: "in_progress"}, + }) + item, ok := added["item"].(map[string]any) + require.True(t, ok, "item must be an object") + for _, k := range []string{"call_id", "name", "arguments"} { + require.Containsf(t, item, k, "function_call item missing %q", k) + } + require.Equal(t, "", item["arguments"]) +} + +// TestWire_MessageItemContentAlwaysArray guards content:[] presence. +func TestWire_MessageItemContentAlwaysArray(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.output_item.added", + OutputIndex: 0, + Item: &ResponsesOutput{Type: "message", ID: "msg_1", Role: "assistant", Status: "in_progress"}, + }) + item, ok := m["item"].(map[string]any) + require.True(t, ok, "item must be an object") + require.Contains(t, item, "content") + _, ok = item["content"].([]any) + require.True(t, ok, "content must be an array") +} + +// TestWire_ReasoningItemSummaryAlwaysArray guards summary:[] presence. +func TestWire_ReasoningItemSummaryAlwaysArray(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.output_item.added", + OutputIndex: 0, + Item: &ResponsesOutput{Type: "reasoning", ID: "rs_1", Status: "in_progress"}, + }) + item, ok := m["item"].(map[string]any) + require.True(t, ok, "item must be an object") + require.Contains(t, item, "summary") + _, ok = item["summary"].([]any) + require.True(t, ok, "summary must be an array") +} + +// TestWire_ContentPartCarriesAnnotationsLogprobs guards the output_text part shape. +func TestWire_ContentPartCarriesAnnotationsLogprobs(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.content_part.added", OutputIndex: 0, ContentIndex: 0, ItemID: "msg_1", + Part: &ResponsesContentPart{Type: "output_text", Text: ""}, + }) + part, ok := m["part"].(map[string]any) + require.True(t, ok, "part must be an object") + require.Equal(t, "output_text", part["type"]) + require.Contains(t, part, "text") + require.Contains(t, part, "annotations") + require.Contains(t, part, "logprobs") +} + +// TestWire_ArgumentsDonePresentEvenEmpty guards arguments presence on done. +func TestWire_ArgumentsDonePresentEvenEmpty(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.function_call_arguments.done", OutputIndex: 1, ItemID: "fc_1", CallID: "call_a", Name: "exec", Arguments: "", + }) + require.Contains(t, m, "arguments") + require.Equal(t, "", m["arguments"]) +} + +// TestWire_UnknownEventFallsBackToDefault ensures non-streamed event types keep +// default marshalling (the response object is preserved). +func TestWire_UnknownEventFallsBackToDefault(t *testing.T) { + m := marshalEvent(t, ResponsesStreamEvent{ + Type: "response.completed", + Response: &ResponsesResponse{ID: "resp_1", Object: "response", Status: "completed"}, + }) + require.Contains(t, m, "response") +} diff --git a/backend/internal/pkg/apicompat/types.go b/backend/internal/pkg/apicompat/types.go index b4451f235bb3..d29378027896 100644 --- a/backend/internal/pkg/apicompat/types.go +++ b/backend/internal/pkg/apicompat/types.go @@ -406,6 +406,10 @@ type ResponsesStreamEvent struct { // Reuses Text/Delta fields above, SummaryIndex identifies which summary part SummaryIndex int `json:"summary_index,omitempty"` + // response.content_part.added / done and + // response.reasoning_summary_part.added / done + Part *ResponsesContentPart `json:"part,omitempty"` + // error event fields Code string `json:"code,omitempty"` Param string `json:"param,omitempty"` diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index bfe092830022..7db35ecc7f90 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -679,6 +679,7 @@ func userEntityToService(u *dbent.User) *service.User { RPMLimit: u.RpmLimit, CreatedAt: u.CreatedAt, UpdatedAt: u.UpdatedAt, + DeletedAt: u.DeletedAt, } // Parse extra emails JSON (supports both old []string and new []NotifyEmailEntry format) if u.BalanceNotifyExtraEmails != "" && u.BalanceNotifyExtraEmails != "[]" { diff --git a/backend/internal/repository/billing_cache.go b/backend/internal/repository/billing_cache.go index 60dae9540d34..de229da9c351 100644 --- a/backend/internal/repository/billing_cache.go +++ b/backend/internal/repository/billing_cache.go @@ -7,6 +7,7 @@ import ( "log" "math/rand/v2" "strconv" + "strings" "time" "github.com/Wei-Shaw/sub2api/internal/service" @@ -338,38 +339,26 @@ func userPlatformQuotaCacheKey(userID int64, platform string) string { return fmt.Sprintf("billing:user_platform_quota:%d:%s", userID, platform) } -func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*service.UserPlatformQuotaCacheEntry, bool, error) { - key := userPlatformQuotaCacheKey(userID, platform) - fields := []string{ - "daily_usage", "weekly_usage", "monthly_usage", "version", "schema_version", - "daily_limit", "weekly_limit", "monthly_limit", - "daily_window_start", "weekly_window_start", "monthly_window_start", - } - vals, err := c.rdb.HMGet(ctx, key, fields...).Result() - if err != nil { - return nil, false, err - } - // 前4个全为nil → key 不存在 - if vals[0] == nil && vals[1] == nil && vals[2] == nil && vals[3] == nil { - return nil, false, nil +// parseUserPlatformQuotaHash 将 Redis HGETALL 返回的 map[string]string 反序列化为 +// *service.UserPlatformQuotaCacheEntry。空 map（key 不存在）返回 nil。 +// GetUserPlatformQuotaCache 和 BatchGetUserPlatformQuotaCache 共用此函数，确保解析逻辑一致。 +func parseUserPlatformQuotaHash(m map[string]string) *service.UserPlatformQuotaCacheEntry { + if len(m) == 0 { + return nil } - parseFloat := func(v any) float64 { - if v == nil { + parseFloat := func(s string) float64 { + if s == "" { return 0 } - s, ok := v.(string) - if !ok { + f, err := strconv.ParseFloat(s, 64) + if err != nil { + log.Printf("billing_cache: corrupt quota usage field %q (using 0): %v", s, err) return 0 } - f, _ := strconv.ParseFloat(s, 64) return f } - parseFloatPtr := func(v any) *float64 { - if v == nil { - return nil - } - s, ok := v.(string) - if !ok || s == "" { + parseFloatPtr := func(s string) *float64 { + if s == "" { return nil } f, err := strconv.ParseFloat(s, 64) @@ -378,12 +367,8 @@ func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int } return &f } - parseTimePtr := func(v any) *time.Time { - if v == nil { - return nil - } - s, ok := v.(string) - if !ok || s == "" { + parseTimePtr := func(s string) *time.Time { + if s == "" { return nil } n, err := strconv.ParseInt(s, 10, 64) @@ -393,30 +378,37 @@ func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int t := time.Unix(n, 0).UTC() return &t } - parseInt64 := func(v any) int64 { - if v == nil { - return 0 - } - s, ok := v.(string) - if !ok { - return 0 - } + parseInt64 := func(s string) int64 { n, _ := strconv.ParseInt(s, 10, 64) return n } return &service.UserPlatformQuotaCacheEntry{ - DailyUsageUSD: parseFloat(vals[0]), - WeeklyUsageUSD: parseFloat(vals[1]), - MonthlyUsageUSD: parseFloat(vals[2]), - Version: parseInt64(vals[3]), - SchemaVersion: parseInt64(vals[4]), - DailyLimitUSD: parseFloatPtr(vals[5]), - WeeklyLimitUSD: parseFloatPtr(vals[6]), - MonthlyLimitUSD: parseFloatPtr(vals[7]), - DailyWindowStart: parseTimePtr(vals[8]), - WeeklyWindowStart: parseTimePtr(vals[9]), - MonthlyWindowStart: parseTimePtr(vals[10]), - }, true, nil + DailyUsageUSD: parseFloat(m["daily_usage"]), + WeeklyUsageUSD: parseFloat(m["weekly_usage"]), + MonthlyUsageUSD: parseFloat(m["monthly_usage"]), + Version: parseInt64(m["version"]), + SchemaVersion: parseInt64(m["schema_version"]), + DailyLimitUSD: parseFloatPtr(m["daily_limit"]), + WeeklyLimitUSD: parseFloatPtr(m["weekly_limit"]), + MonthlyLimitUSD: parseFloatPtr(m["monthly_limit"]), + DailyWindowStart: parseTimePtr(m["daily_window_start"]), + WeeklyWindowStart: parseTimePtr(m["weekly_window_start"]), + MonthlyWindowStart: parseTimePtr(m["monthly_window_start"]), + } +} + +func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*service.UserPlatformQuotaCacheEntry, bool, error) { + key := userPlatformQuotaCacheKey(userID, platform) + m, err := c.rdb.HGetAll(ctx, key).Result() + if err != nil { + return nil, false, err + } + entry := parseUserPlatformQuotaHash(m) + if entry == nil { + // 空 map → key 不存在 → MISS + return nil, false, nil + } + return entry, true, nil } func (c *billingCache) SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *service.UserPlatformQuotaCacheEntry, ttl time.Duration) error { @@ -468,9 +460,12 @@ func (c *billingCache) DeleteUserPlatformQuotaCache(ctx context.Context, userID // SetCache 重建为新版 entry —— 若此处仍累加，上层覆盖时会丢失这部分增量，导致 Redis usage 比真实偏小。 // key 不存在同样跳过（由下次 SetCache 重建）。 // KEYS[1] = hash key +// KEYS[2] = 脏集 key（dirty set） // ARGV[1] = cost (string float) // ARGV[2] = ttl seconds // ARGV[3] = expected schema_version (Go 侧 UserPlatformQuotaCacheSchemaV1) +// ARGV[4] = dirty set member（空串则不 SADD） +// ARGV[5] = 脏集兜底 TTL 秒 const updateUserPlatformQuotaUsageScript = ` if redis.call("EXISTS", KEYS[1]) == 0 then return 0 @@ -484,18 +479,125 @@ redis.call("HINCRBYFLOAT", KEYS[1], "weekly_usage", ARGV[1]) redis.call("HINCRBYFLOAT", KEYS[1], "monthly_usage", ARGV[1]) redis.call("HINCRBY", KEYS[1], "version", 1) redis.call("EXPIRE", KEYS[1], ARGV[2]) +if ARGV[4] ~= "" then + redis.call("SADD", KEYS[2], ARGV[4]) + redis.call("EXPIRE", KEYS[2], ARGV[5]) +end return 1 ` -func (c *billingCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { - key := userPlatformQuotaCacheKey(userID, platform) - _, err := c.rdb.Eval(ctx, updateUserPlatformQuotaUsageScript, []string{key}, +// userPlatformQuotaDirtySetKey 返回脏集（dirty set）的 Redis key。 +// 使用与 userPlatformQuotaCacheKey 相同的前缀 "billing:"。 +func userPlatformQuotaDirtySetKey() string { return "billing:" + "upq:dirty" } + +// userPlatformQuotaDirtyTTLSeconds 脏集兜底 TTL（秒）：初始 SADD（Lua）与 Readd 共用， +// 确保 flusher 长期停摆时脏集最终过期；正常运行因持续 SADD 不断续期。 +const userPlatformQuotaDirtyTTLSeconds = 86400 + +// userPlatformQuotaDirtyMember 构造脏集成员字符串 "userID:platform"。 +func userPlatformQuotaDirtyMember(userID int64, platform string) string { + return strconv.FormatInt(userID, 10) + ":" + platform +} + +func (c *billingCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { + member := "" + if markDirty { + member = userPlatformQuotaDirtyMember(userID, platform) + } + _, err := c.rdb.Eval(ctx, updateUserPlatformQuotaUsageScript, + []string{userPlatformQuotaCacheKey(userID, platform), userPlatformQuotaDirtySetKey()}, strconv.FormatFloat(cost, 'f', -1, 64), int(ttl.Seconds()), service.UserPlatformQuotaCacheSchemaV1, + member, + userPlatformQuotaDirtyTTLSeconds, ).Result() if err != nil && !errors.Is(err, redis.Nil) { return err } return nil } + +// parseUserPlatformQuotaDirtyMember 将脏集成员字符串 "userID:platform" 解析为 +// service.UserPlatformQuotaKey。解析失败返回 ok=false。 +func parseUserPlatformQuotaDirtyMember(m string) (service.UserPlatformQuotaKey, bool) { + parts := strings.SplitN(m, ":", 2) + if len(parts) != 2 { + return service.UserPlatformQuotaKey{}, false + } + uid, err := strconv.ParseInt(parts[0], 10, 64) + if err != nil { + return service.UserPlatformQuotaKey{}, false + } + return service.UserPlatformQuotaKey{UserID: uid, Platform: parts[1]}, true +} + +// PopDirtyUserPlatformQuotaKeys 从脏集随机弹出最多 n 个 key。 +// 脏集为空时返回 (nil, nil)。 +func (c *billingCache) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]service.UserPlatformQuotaKey, error) { + members, err := c.rdb.SPopN(ctx, userPlatformQuotaDirtySetKey(), int64(n)).Result() + if err != nil { + if errors.Is(err, redis.Nil) { + return nil, nil + } + return nil, err + } + keys := make([]service.UserPlatformQuotaKey, 0, len(members)) + for _, m := range members { + k, ok := parseUserPlatformQuotaDirtyMember(m) + if !ok { + log.Printf("billing_cache: skipping invalid dirty member %q", m) + continue + } + keys = append(keys, k) + } + return keys, nil +} + +// ReaddDirtyUserPlatformQuotaKeys 将 keys 重新加入脏集（flush 失败时回填）。 +// 通过 pipeline 同时执行 SAdd + Expire，确保 Readd 后脏集具有兜底 TTL。 +// 空切片时直接返回 nil。 +func (c *billingCache) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []service.UserPlatformQuotaKey) error { + if len(keys) == 0 { + return nil + } + dirtyKey := userPlatformQuotaDirtySetKey() + members := make([]any, len(keys)) + for i, k := range keys { + members[i] = userPlatformQuotaDirtyMember(k.UserID, k.Platform) + } + pipe := c.rdb.Pipeline() + pipe.SAdd(ctx, dirtyKey, members...) + pipe.Expire(ctx, dirtyKey, userPlatformQuotaDirtyTTLSeconds*time.Second) + _, err := pipe.Exec(ctx) + return err +} + +// BatchGetUserPlatformQuotaCache 通过 Pipeline 批量 HGETALL 获取多个 user×platform 的 +// quota cache。返回切片与 keys 顺序、长度对齐；MISS 或解析失败位置返回 nil。 +func (c *billingCache) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []service.UserPlatformQuotaKey) ([]*service.UserPlatformQuotaCacheEntry, error) { + if len(keys) == 0 { + return nil, nil + } + pipe := c.rdb.Pipeline() + cmds := make([]*redis.MapStringStringCmd, len(keys)) + for i, k := range keys { + cmds[i] = pipe.HGetAll(ctx, userPlatformQuotaCacheKey(k.UserID, k.Platform)) + } + if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) { + return nil, err + } + results := make([]*service.UserPlatformQuotaCacheEntry, len(keys)) + for i, cmd := range cmds { + m, err := cmd.Result() + if err != nil { + if !errors.Is(err, redis.Nil) { + log.Printf("billing_cache: BatchGet HGETALL cmd[%d] failed: %v (skip, self-heal)", i, err) + } + // 单个命令失败 → 对应位置 nil，继续 + continue + } + results[i] = parseUserPlatformQuotaHash(m) + } + return results, nil +} diff --git a/backend/internal/repository/billing_cache_user_platform_quota_test.go b/backend/internal/repository/billing_cache_user_platform_quota_test.go index 8d49fd31cfde..15b185e77693 100644 --- a/backend/internal/repository/billing_cache_user_platform_quota_test.go +++ b/backend/internal/repository/billing_cache_user_platform_quota_test.go @@ -88,7 +88,7 @@ func TestUserPlatformQuotaCache_NilLimitSetThenGet(t *testing.T) { func TestUserPlatformQuotaCache_IncrMissIsNoop(t *testing.T) { c, _ := newMiniRedisCache(t) - if err := c.IncrUserPlatformQuotaUsageCache(context.Background(), 1, "openai", 0.5, time.Minute); err != nil { + if err := c.IncrUserPlatformQuotaUsageCache(context.Background(), 1, "openai", 0.5, time.Minute, false); err != nil { t.Fatal(err) } _, ok, _ := c.GetUserPlatformQuotaCache(context.Background(), 1, "openai") @@ -105,10 +105,10 @@ func TestUserPlatformQuotaCache_IncrHitAccumulates(t *testing.T) { Version: 1, SchemaVersion: service.UserPlatformQuotaCacheSchemaV1, }, time.Minute) - if err := c.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 0.5, time.Minute); err != nil { + if err := c.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 0.5, time.Minute, false); err != nil { t.Fatal(err) } - if err := c.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 0.25, time.Minute); err != nil { + if err := c.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 0.25, time.Minute, false); err != nil { t.Fatal(err) } got, _, _ := c.GetUserPlatformQuotaCache(ctx, 1, "openai") diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index ab01a8638ace..921aa081cae5 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -548,6 +548,18 @@ func filterSchedulerExtra(extra map[string]any) map[string]any { "openai_ws_force_http", "openai_responses_mode", "openai_responses_supported", + "codex_5h_used_percent", + "codex_7d_used_percent", + "codex_5h_reset_at", + "codex_7d_reset_at", + "codex_5h_reset_after_seconds", + "codex_7d_reset_after_seconds", + "codex_usage_updated_at", + "auto_pause_5h_threshold", + "auto_pause_7d_threshold", + "auto_pause_5h_disabled", + "auto_pause_7d_disabled", + "model_rate_limits", } filtered := make(map[string]any) for _, key := range keys { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index 86de87c70995..c14721cd7c18 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -75,3 +75,62 @@ func TestBuildSchedulerMetadataAccount_KeepsSlimGroupMembership(t *testing.T) { require.Equal(t, int64(11), got.AccountGroups[1].GroupID) require.Nil(t, got.Groups) } + +func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) { + account := service.Account{ + ID: 88, + Extra: map[string]any{ + "codex_5h_used_percent": 12.34, + "codex_7d_used_percent": 56.78, + "codex_5h_reset_at": "2026-05-29T10:00:00Z", + "codex_7d_reset_at": "2026-06-01T10:00:00Z", + "codex_5h_reset_after_seconds": 300, + "codex_7d_reset_after_seconds": 600, + "codex_usage_updated_at": "2026-05-29T09:00:00Z", + "auto_pause_5h_threshold": 0.95, + "auto_pause_7d_threshold": 0.96, + "auto_pause_5h_disabled": true, + "auto_pause_7d_disabled": false, + }, + } + + got := buildSchedulerMetadataAccount(account) + + require.Equal(t, 12.34, got.Extra["codex_5h_used_percent"]) + require.Equal(t, 56.78, got.Extra["codex_7d_used_percent"]) + require.Equal(t, "2026-05-29T10:00:00Z", got.Extra["codex_5h_reset_at"]) + require.Equal(t, "2026-06-01T10:00:00Z", got.Extra["codex_7d_reset_at"]) + require.Equal(t, 300, got.Extra["codex_5h_reset_after_seconds"]) + require.Equal(t, 600, got.Extra["codex_7d_reset_after_seconds"]) + require.Equal(t, "2026-05-29T09:00:00Z", got.Extra["codex_usage_updated_at"]) + require.Equal(t, 0.95, got.Extra["auto_pause_5h_threshold"]) + require.Equal(t, 0.96, got.Extra["auto_pause_7d_threshold"]) + require.Equal(t, true, got.Extra["auto_pause_5h_disabled"]) + require.Equal(t, false, got.Extra["auto_pause_7d_disabled"]) +} + +func TestBuildSchedulerMetadataAccount_KeepsModelRateLimits(t *testing.T) { + account := service.Account{ + ID: 90, + Platform: service.PlatformAntigravity, + Extra: map[string]any{ + "model_rate_limits": map[string]any{ + "gemini-3-flash": map[string]any{ + "rate_limit_reset_at": "2026-05-30T10:10:00Z", + }, + "antigravity:gemini": map[string]any{ + "rate_limit_reset_at": "2026-05-30T10:10:00Z", + }, + }, + "unused_large_field": "drop-me", + }, + } + + got := buildSchedulerMetadataAccount(account) + + limits, ok := got.Extra["model_rate_limits"].(map[string]any) + require.True(t, ok) + require.Contains(t, limits, "gemini-3-flash") + require.Contains(t, limits, "antigravity:gemini") + require.Nil(t, got.Extra["unused_large_field"]) +} diff --git a/backend/internal/repository/usage_log_repo.go b/backend/internal/repository/usage_log_repo.go index dd3596f7f5c0..9ee696246597 100644 --- a/backend/internal/repository/usage_log_repo.go +++ b/backend/internal/repository/usage_log_repo.go @@ -17,6 +17,7 @@ import ( dbaccount "github.com/Wei-Shaw/sub2api/ent/account" dbapikey "github.com/Wei-Shaw/sub2api/ent/apikey" dbgroup "github.com/Wei-Shaw/sub2api/ent/group" + "github.com/Wei-Shaw/sub2api/ent/schema/mixins" dbuser "github.com/Wei-Shaw/sub2api/ent/user" dbusersub "github.com/Wei-Shaw/sub2api/ent/usersubscription" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" @@ -26,6 +27,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/service" "github.com/lib/pq" gocache "github.com/patrickmn/go-cache" + "golang.org/x/sync/errgroup" ) 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, audio_input_tokens, audio_output_tokens, audio_cache_creation_tokens, audio_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, 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" @@ -3763,28 +3765,6 @@ func (r *usageLogRepository) GetStatsWithFilters(ctx context.Context, filters Us stats := &UsageStats{} var totalAccountCost float64 - if err := scanSingleRow( - ctx, - r.sql, - query, - args, - &stats.TotalRequests, - &stats.TotalInputTokens, - &stats.TotalOutputTokens, - &stats.TotalCacheTokens, - &stats.TotalAudioInputTokens, - &stats.TotalAudioOutputTokens, - &stats.TotalAudioCacheCreationTokens, - &stats.TotalAudioCacheReadTokens, - &stats.TotalCost, - &stats.TotalActualCost, - &totalAccountCost, - &stats.AverageDurationMs, - ); err != nil { - return nil, err - } - stats.TotalAccountCost = &totalAccountCost - stats.TotalTokens = usageStatsTotalTokens(stats) start := time.Unix(0, 0).UTC() if filters.StartTime != nil { @@ -3795,21 +3775,80 @@ func (r *usageLogRepository) GetStatsWithFilters(ctx context.Context, filters Us end = *filters.EndTime } - endpoints, endpointErr := r.GetEndpointStatsWithFilters(ctx, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) - if endpointErr != nil { - logger.LegacyPrintf("repository.usage_log", "GetEndpointStatsWithFilters failed in GetStatsWithFilters: %v", endpointErr) - endpoints = []EndpointStat{} + var endpoints, upstreamEndpoints, endpointPaths []EndpointStat + + // 汇总查询:失败即致命。 + runSummary := func(c context.Context) error { + return scanSingleRow( + c, r.sql, query, args, + &stats.TotalRequests, + &stats.TotalInputTokens, + &stats.TotalOutputTokens, + &stats.TotalCacheTokens, + &stats.TotalAudioInputTokens, + &stats.TotalAudioOutputTokens, + &stats.TotalAudioCacheCreationTokens, + &stats.TotalAudioCacheReadTokens, + &stats.TotalCost, + &stats.TotalActualCost, + &totalAccountCost, + &stats.AverageDurationMs, + ) } - upstreamEndpoints, upstreamEndpointErr := r.GetUpstreamEndpointStatsWithFilters(ctx, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) - if upstreamEndpointErr != nil { - logger.LegacyPrintf("repository.usage_log", "GetUpstreamEndpointStatsWithFilters failed in GetStatsWithFilters: %v", upstreamEndpointErr) - upstreamEndpoints = []EndpointStat{} + // endpoint 明细:best-effort(失败 log + 返空),不致命。 + runEndpoints := func(c context.Context) { + res, err := r.GetEndpointStatsWithFilters(c, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) + if err != nil { + if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("repository.usage_log", "GetEndpointStatsWithFilters failed in GetStatsWithFilters: %v", err) + } + res = []EndpointStat{} + } + endpoints = res + } + runUpstream := func(c context.Context) { + res, err := r.GetUpstreamEndpointStatsWithFilters(c, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) + if err != nil { + if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("repository.usage_log", "GetUpstreamEndpointStatsWithFilters failed in GetStatsWithFilters: %v", err) + } + res = []EndpointStat{} + } + upstreamEndpoints = res + } + runPaths := func(c context.Context) { + res, err := r.getEndpointPathStatsWithFilters(c, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) + if err != nil { + if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("repository.usage_log", "getEndpointPathStatsWithFilters failed in GetStatsWithFilters: %v", err) + } + res = []EndpointStat{} + } + endpointPaths = res } - endpointPaths, endpointPathErr := r.getEndpointPathStatsWithFilters(ctx, start, end, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType) - if endpointPathErr != nil { - logger.LegacyPrintf("repository.usage_log", "getEndpointPathStatsWithFilters failed in GetStatsWithFilters: %v", endpointPathErr) - endpointPaths = []EndpointStat{} + + if r.db != nil { + // 生产路径:r.sql 是 *sql.DB 连接池,可并发。4 条查询并行,延迟取最大值。 + g, gctx := errgroup.WithContext(ctx) + g.Go(func() error { return runSummary(gctx) }) + g.Go(func() error { runEndpoints(gctx); return nil }) + g.Go(func() error { runUpstream(gctx); return nil }) + g.Go(func() error { runPaths(gctx); return nil }) + if err := g.Wait(); err != nil { + return nil, err + } + } else { + // 事务路径(ent.Tx 不能并发查询):顺序执行,行为与重构前一致。 + if err := runSummary(ctx); err != nil { + return nil, err + } + runEndpoints(ctx) + runUpstream(ctx) + runPaths(ctx) } + + stats.TotalAccountCost = &totalAccountCost + stats.TotalTokens = usageStatsTotalTokens(stats) stats.Endpoints = endpoints stats.UpstreamEndpoints = upstreamEndpoints stats.EndpointPaths = endpointPaths @@ -4351,7 +4390,8 @@ func (r *usageLogRepository) loadUsers(ctx context.Context, ids []int64) (map[in if len(ids) == 0 { return out, nil } - models, err := r.client.User.Query().Where(dbuser.IDIn(ids...)).All(ctx) + // 无条件穿透软删除：ids 来自调用方已按 user_id 筛选的日志行；普通用户路径强制 UserID=本人（本人必为活跃用户），不会借此解析他人已删身份；仅 admin 路径可借此显示已删用户。 + models, err := r.client.User.Query().Where(dbuser.IDIn(ids...)).All(mixins.SkipSoftDelete(ctx)) if err != nil { return nil, err } diff --git a/backend/internal/repository/usage_log_repo_deleted_user_integration_test.go b/backend/internal/repository/usage_log_repo_deleted_user_integration_test.go new file mode 100644 index 000000000000..70835b03d56b --- /dev/null +++ b/backend/internal/repository/usage_log_repo_deleted_user_integration_test.go @@ -0,0 +1,65 @@ +//go:build integration + +package repository + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestUsageLog_ListWithFilters_ResolvesSoftDeletedUser(t *testing.T) { + ctx := context.Background() + tx := testEntTx(t) + client := tx.Client() + repo := newUsageLogRepositoryWithSQL(client, tx) + + // 一个活跃用户、一个将被软删的用户，各一条日志。 + active := mustCreateUser(t, client, &service.User{Email: "active-listfilter@test.com"}) + deleted := mustCreateUser(t, client, &service.User{Email: "deleted-listfilter@test.com"}) + apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: deleted.ID, Key: "sk-del-1", Name: "k"}) + apiKey2 := mustCreateApiKey(t, client, &service.APIKey{UserID: active.ID, Key: "sk-act-1", Name: "k"}) + account := mustCreateAccount(t, client, &service.Account{Name: "acc-listfilter"}) + + now := time.Now().UTC() + for _, u := range []struct { + uid int64 + kid int64 + }{{deleted.ID, apiKey.ID}, {active.ID, apiKey2.ID}} { + _, err := repo.Create(ctx, &service.UsageLog{ + UserID: u.uid, APIKeyID: u.kid, AccountID: account.ID, + Model: "claude-3", InputTokens: 1, OutputTokens: 1, + TotalCost: 0.1, ActualCost: 0.1, CreatedAt: now, + }) + require.NoError(t, err) + } + + // 软删除该用户（触发 SoftDeleteMixin Hook → UPDATE deleted_at）。 + require.NoError(t, client.User.DeleteOneID(deleted.ID).Exec(ctx)) + + logs, _, err := repo.ListWithFilters(ctx, pagination.PaginationParams{Page: 1, PageSize: 50}, + usagestats.UsageLogFilters{ExactTotal: true}) + require.NoError(t, err) + + byUser := map[int64]service.UsageLog{} + for _, l := range logs { + byUser[l.UserID] = l + } + + // 已删用户的日志行：富化后 User 非 nil、邮箱正确、DeletedAt 非 nil。 + delLog, ok := byUser[deleted.ID] + require.True(t, ok, "deleted user's usage log must still be listed") + require.NotNil(t, delLog.User, "deleted user identity must resolve") + require.Equal(t, "deleted-listfilter@test.com", delLog.User.Email) + require.NotNil(t, delLog.User.DeletedAt, "DeletedAt must be set for soft-deleted user") + + // 活跃用户：DeletedAt 为 nil。 + actLog := byUser[active.ID] + require.NotNil(t, actLog.User) + require.Nil(t, actLog.User.DeletedAt) +} diff --git a/backend/internal/repository/usage_log_repo_stats_integration_test.go b/backend/internal/repository/usage_log_repo_stats_integration_test.go new file mode 100644 index 000000000000..09ac2aee1417 --- /dev/null +++ b/backend/internal/repository/usage_log_repo_stats_integration_test.go @@ -0,0 +1,51 @@ +//go:build integration + +package repository + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestUsageLog_GetStatsWithFilters_AggregatesAndEndpoints(t *testing.T) { + ctx := context.Background() + tx := testEntTx(t) + client := tx.Client() + repo := newUsageLogRepositoryWithSQL(client, tx) + + user := mustCreateUser(t, client, &service.User{Email: "stats@test.com"}) + apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-stats-1", Name: "k"}) + account := mustCreateAccount(t, client, &service.Account{Name: "acc-stats"}) + + now := time.Now().UTC() + inboundEndpoint := "/v1/messages" + upstreamEndpoint := "/v1/responses" + for i := 0; i < 3; i++ { + _, err := repo.Create(ctx, &service.UsageLog{ + UserID: user.ID, APIKeyID: apiKey.ID, AccountID: account.ID, + Model: "claude-3", InputTokens: 2, OutputTokens: 3, + TotalCost: 0.5, ActualCost: 0.4, CreatedAt: now, + InboundEndpoint: &inboundEndpoint, UpstreamEndpoint: &upstreamEndpoint, + }) + require.NoError(t, err) + } + + start := now.Add(-1 * time.Hour) + end := now.Add(1 * time.Hour) + // 按本测试创建的 user 维度过滤:集成库为共享实例,其它用 testEntClient 的兄弟测试会留下 + // 已提交的 usage_log 行(含零 token 的失败请求),不限定 user 会把它们计入 TotalRequests。 + stats, err := repo.GetStatsWithFilters(ctx, usagestats.UsageLogFilters{UserID: user.ID, StartTime: &start, EndTime: &end}) + require.NoError(t, err) + require.Equal(t, int64(3), stats.TotalRequests) + require.Equal(t, int64(6), stats.TotalInputTokens) + require.Equal(t, int64(9), stats.TotalOutputTokens) + require.InDelta(t, 1.2, stats.TotalActualCost, 1e-9) + require.NotEmpty(t, stats.Endpoints) + require.NotEmpty(t, stats.UpstreamEndpoints) + require.NotEmpty(t, stats.EndpointPaths) +} diff --git a/backend/internal/repository/user_platform_quota_adapter_test.go b/backend/internal/repository/user_platform_quota_adapter_test.go index a55d2e9c0330..f31defe5d55a 100644 --- a/backend/internal/repository/user_platform_quota_adapter_test.go +++ b/backend/internal/repository/user_platform_quota_adapter_test.go @@ -38,6 +38,9 @@ func (f *fakeRepoForAdapter) UpsertForUser(_ context.Context, userID int64, reco f.upsertCalledWith = records return f.upsertErr } +func (f *fakeRepoForAdapter) BatchSnapshotUsage(_ context.Context, _ []UserPlatformQuotaSnapshot, _ time.Time) error { + return nil +} func TestGenericAdapter_UpsertForUser_ForwardsRecords(t *testing.T) { fake := &fakeRepoForAdapter{} diff --git a/backend/internal/repository/user_platform_quota_repo.go b/backend/internal/repository/user_platform_quota_repo.go index 1e2e7f51bd48..ccba2330ced2 100644 --- a/backend/internal/repository/user_platform_quota_repo.go +++ b/backend/internal/repository/user_platform_quota_repo.go @@ -2,6 +2,7 @@ package repository import ( "context" + "errors" "fmt" "strings" "time" @@ -9,6 +10,7 @@ import ( dbent "github.com/Wei-Shaw/sub2api/ent" "github.com/Wei-Shaw/sub2api/ent/userplatformquota" "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" + "github.com/lib/pq" ) // UserPlatformQuotaRecord 是 repository 层的传输结构体， @@ -30,6 +32,22 @@ type UserPlatformQuotaRecord struct { // ErrUserPlatformQuotaNotFound 用于 ResetExpiredWindow 等需要"必须命中已有记录"的方法。 var ErrUserPlatformQuotaNotFound = fmt.Errorf("user platform quota record not found") +// ErrUserPlatformQuotaFKViolation 当批量 UPSERT 中存在 user_id 不在 users 表的记录时返回。 +var ErrUserPlatformQuotaFKViolation = errors.New("user platform quota snapshot FK violation") + +// UserPlatformQuotaSnapshot 是 BatchSnapshotUsage 的输入结构体， +// 表示 Redis 当前窗口快照（用于绝对值覆盖写入 DB）。 +type UserPlatformQuotaSnapshot struct { + UserID int64 + Platform string + DailyUsageUSD float64 + WeeklyUsageUSD float64 + MonthlyUsageUSD float64 + DailyWindowStart time.Time + WeeklyWindowStart time.Time + MonthlyWindowStart time.Time +} + // UserPlatformQuotaRepository 定义用户平台配额的数据访问接口。 type UserPlatformQuotaRepository interface { // BulkInsertInitial 幂等批量插入初始配额记录（ON CONFLICT DO NOTHING）。 @@ -44,6 +62,10 @@ type UserPlatformQuotaRepository interface { ResetExpiredWindow(ctx context.Context, userID int64, platform string, window string, newStart time.Time) error // UpsertForUser 全量替换该用户所有平台限额配置（详见 service.UserPlatformQuotaRepository.UpsertForUser）。 UpsertForUser(ctx context.Context, userID int64, records []UserPlatformQuotaRecord) error + // BatchSnapshotUsage 用一条多行 UPSERT 把整批 usage 以绝对值覆盖写入(非累加)。 + // usage/window_start 直接取 EXCLUDED(Redis 当前窗口快照),无 CASE。整批共用 now 作 created/updated_at。 + // 要求 snapshots 内 (user,platform) 不重复。FK 违反返回 ErrUserPlatformQuotaFKViolation。 + BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error } type userPlatformQuotaRepository struct { @@ -414,3 +436,73 @@ func insertLimitsRow(ctx context.Context, client *dbent.Client, userID int64, re } return nil } + +// batchRows 是 BatchSnapshotUsage 每批最大行数（9 参/行 × 6000 ≈ 54000 参,低于 Postgres 65535 上限）。 +const batchRows = 6000 + +// BatchSnapshotUsage 用一条多行 UPSERT 把整批 usage 以绝对值覆盖写入（非累加）。 +// 每批最多 batchRows 行；$1=now 共用；每行 8 个 per-row 参（user_id, platform, 3×usage, 3×window_start）。 +// FK 违反（user_id 不存在）返回 ErrUserPlatformQuotaFKViolation。 +// +// 注意:snapshots 超过 batchRows 会分多条 SQL 执行且【非单事务】——若某子批 FK 失败, +// 先前子批已写入无法回滚。调用方(flusher)应保证单次 batchSize ≤ batchRows +// (默认 flush_batch_size=1000 < 6000,安全)。 +// 另注:启用 flusher 后,本绝对值覆盖与 admin 直写 DB(ResetExpiredWindow/UpsertForUser)存在覆盖竞态, +// 详见 service/user_platform_quota_flusher.go 中 flushOneBatch 的"已知竞态"注释。 +func (r *userPlatformQuotaRepository) BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error { + if len(snapshots) == 0 { + return nil + } + + client := clientFromContext(ctx, r.client) + + for start := 0; start < len(snapshots); start += batchRows { + end := start + batchRows + if end > len(snapshots) { + end = len(snapshots) + } + batch := snapshots[start:end] + + var sb strings.Builder + _, _ = sb.WriteString( + "INSERT INTO user_platform_quotas" + + " (user_id, platform, daily_usage_usd, weekly_usage_usd, monthly_usage_usd," + + " daily_window_start, weekly_window_start, monthly_window_start, created_at, updated_at)" + + " VALUES ") + + // $1 = now（共用）；每行 8 个 per-row 参，从 $2 起连续编号。 + args := []any{now} + for i, s := range batch { + if i > 0 { + _, _ = sb.WriteString(",") + } + b := len(args) // 当前 per-row 第一个参数的 0-based 索引，实际占位符 = b+1 + fmt.Fprintf(&sb, "($%d,$%d,$%d,$%d,$%d,$%d,$%d,$%d,$1,$1)", + b+1, b+2, b+3, b+4, b+5, b+6, b+7, b+8) + args = append(args, + s.UserID, s.Platform, + s.DailyUsageUSD, s.WeeklyUsageUSD, s.MonthlyUsageUSD, + s.DailyWindowStart, s.WeeklyWindowStart, s.MonthlyWindowStart, + ) + } + + _, _ = sb.WriteString( + " ON CONFLICT (user_id, platform) WHERE deleted_at IS NULL DO UPDATE SET" + + " daily_usage_usd = EXCLUDED.daily_usage_usd," + + " weekly_usage_usd = EXCLUDED.weekly_usage_usd," + + " monthly_usage_usd = EXCLUDED.monthly_usage_usd," + + " daily_window_start = EXCLUDED.daily_window_start," + + " weekly_window_start = EXCLUDED.weekly_window_start," + + " monthly_window_start = EXCLUDED.monthly_window_start," + + " updated_at = EXCLUDED.updated_at") + + if _, err := client.ExecContext(ctx, sb.String(), args...); err != nil { + var pqErr *pq.Error + if errors.As(err, &pqErr) && pqErr.Code == "23503" { + return ErrUserPlatformQuotaFKViolation + } + return err + } + } + return nil +} diff --git a/backend/internal/repository/user_platform_quota_repo_integration_test.go b/backend/internal/repository/user_platform_quota_repo_integration_test.go index f02eeaa900cc..39e2f6e0ba54 100644 --- a/backend/internal/repository/user_platform_quota_repo_integration_test.go +++ b/backend/internal/repository/user_platform_quota_repo_integration_test.go @@ -267,3 +267,101 @@ func TestUserPlatformQuotaRepository_ResetExpiredWindow_NotFoundReturnsSentinel( require.True(t, errors.Is(err, ErrUserPlatformQuotaNotFound), "expected ErrUserPlatformQuotaNotFound, got %v", err) } + +// TestBatchSnapshotUsage_InsertOverwriteMultiKey 验证 BatchSnapshotUsage 的绝对值覆盖语义： +// 1. 首批插入 2 条（不同 user），验证 daily 等于首批值； +// 2. 对同一 key 传不同值，验证 daily 等于新值（绝对覆盖，非累加）。 +func TestBatchSnapshotUsage_InsertOverwriteMultiKey(t *testing.T) { + ctx := context.Background() + // BatchSnapshotUsage 不开事务（直接写），使用独立 client 保证跨调用可见性。 + client := testEntClient(t) + + userID1 := mustCreateUserForQuota(t, client) + userID2 := mustCreateUserForQuota(t, client) + + repo := NewUserPlatformQuotaRepository(client) + + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + dailyStart := time.Date(2026, 5, 29, 0, 0, 0, 0, time.UTC) + weeklyStart := time.Date(2026, 5, 25, 0, 0, 0, 0, time.UTC) // 当周一 + monthlyStart := time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC) + + // ── 第一批：插入 2 行 ────────────────────────────────────────────────────── + firstBatch := []UserPlatformQuotaSnapshot{ + { + UserID: userID1, + Platform: "anthropic", + DailyUsageUSD: 1.0, + WeeklyUsageUSD: 3.0, + MonthlyUsageUSD: 5.0, + DailyWindowStart: dailyStart, + WeeklyWindowStart: weeklyStart, + MonthlyWindowStart: monthlyStart, + }, + { + UserID: userID2, + Platform: "openai", + DailyUsageUSD: 2.0, + WeeklyUsageUSD: 4.0, + MonthlyUsageUSD: 6.0, + DailyWindowStart: dailyStart, + WeeklyWindowStart: weeklyStart, + MonthlyWindowStart: monthlyStart, + }, + } + require.NoError(t, repo.BatchSnapshotUsage(ctx, firstBatch, now), "first batch upsert") + + // 验证首批 daily 值 + rec1, err := repo.GetByUserPlatform(ctx, userID1, "anthropic") + require.NoError(t, err) + require.NotNil(t, rec1, "user1/anthropic should exist after first batch") + require.InDelta(t, 1.0, rec1.DailyUsageUSD, 1e-9, "user1 daily after first batch") + require.InDelta(t, 3.0, rec1.WeeklyUsageUSD, 1e-9, "user1 weekly after first batch") + require.InDelta(t, 5.0, rec1.MonthlyUsageUSD, 1e-9, "user1 monthly after first batch") + + rec2, err := repo.GetByUserPlatform(ctx, userID2, "openai") + require.NoError(t, err) + require.NotNil(t, rec2, "user2/openai should exist after first batch") + require.InDelta(t, 2.0, rec2.DailyUsageUSD, 1e-9, "user2 daily after first batch") + + // ── 第二批：对同一 key 传不同值，验证绝对覆盖（非累加）────────────────── + now2 := now.Add(5 * time.Minute) + secondBatch := []UserPlatformQuotaSnapshot{ + { + UserID: userID1, + Platform: "anthropic", + DailyUsageUSD: 9.9, // 新值，不是 1.0+9.9=10.9 + WeeklyUsageUSD: 19.9, // 新值，不是 3.0+19.9=22.9 + MonthlyUsageUSD: 29.9, // 新值 + DailyWindowStart: dailyStart, + WeeklyWindowStart: weeklyStart, + MonthlyWindowStart: monthlyStart, + }, + { + UserID: userID2, + Platform: "openai", + DailyUsageUSD: 8.8, + WeeklyUsageUSD: 18.8, + MonthlyUsageUSD: 28.8, + DailyWindowStart: dailyStart, + WeeklyWindowStart: weeklyStart, + MonthlyWindowStart: monthlyStart, + }, + } + require.NoError(t, repo.BatchSnapshotUsage(ctx, secondBatch, now2), "second batch upsert") + + // 验证第二批覆盖：daily 应为新值，不是累加 + rec1After, err := repo.GetByUserPlatform(ctx, userID1, "anthropic") + require.NoError(t, err) + require.NotNil(t, rec1After) + require.InDelta(t, 9.9, rec1After.DailyUsageUSD, 1e-9, "user1 daily must be overwritten to 9.9 (not accumulated)") + require.InDelta(t, 19.9, rec1After.WeeklyUsageUSD, 1e-9, "user1 weekly must be overwritten to 19.9") + require.InDelta(t, 29.9, rec1After.MonthlyUsageUSD, 1e-9, "user1 monthly must be overwritten to 29.9") + + rec2After, err := repo.GetByUserPlatform(ctx, userID2, "openai") + require.NoError(t, err) + require.NotNil(t, rec2After) + require.InDelta(t, 8.8, rec2After.DailyUsageUSD, 1e-9, "user2 daily must be overwritten to 8.8 (not accumulated)") + require.InDelta(t, 18.8, rec2After.WeeklyUsageUSD, 1e-9, "user2 weekly must be overwritten to 18.8") + require.InDelta(t, 28.8, rec2After.MonthlyUsageUSD, 1e-9, "user2 monthly must be overwritten to 28.8") +} diff --git a/backend/internal/repository/user_platform_quota_service_adapter.go b/backend/internal/repository/user_platform_quota_service_adapter.go index 7495cd26ba2b..5240bb5412d1 100644 --- a/backend/internal/repository/user_platform_quota_service_adapter.go +++ b/backend/internal/repository/user_platform_quota_service_adapter.go @@ -94,6 +94,29 @@ func (a *userPlatformQuotaServiceAdapter) ResetExpiredWindow(ctx context.Context return err } +// BatchSnapshotUsage 转换 []service.UserPlatformQuotaSnapshot → []UserPlatformQuotaSnapshot， +// 调底层 repo，并将 repository FK sentinel 包装为 service sentinel。 +func (a *userPlatformQuotaServiceAdapter) BatchSnapshotUsage(ctx context.Context, snapshots []service.UserPlatformQuotaSnapshot, now time.Time) error { + repoSnaps := make([]UserPlatformQuotaSnapshot, len(snapshots)) + for i, s := range snapshots { + repoSnaps[i] = UserPlatformQuotaSnapshot{ + UserID: s.UserID, + Platform: s.Platform, + DailyUsageUSD: s.DailyUsageUSD, + WeeklyUsageUSD: s.WeeklyUsageUSD, + MonthlyUsageUSD: s.MonthlyUsageUSD, + DailyWindowStart: s.DailyWindowStart, + WeeklyWindowStart: s.WeeklyWindowStart, + MonthlyWindowStart: s.MonthlyWindowStart, + } + } + err := a.inner.BatchSnapshotUsage(ctx, repoSnaps, now) + if errors.Is(err, ErrUserPlatformQuotaFKViolation) { + return fmt.Errorf("%w: %v", service.ErrUserPlatformQuotaFKViolation, err) + } + return err +} + // genericUserPlatformQuotaAdapter 通过通用接口适配（用于测试 fake 或非标准实现）。 type genericUserPlatformQuotaAdapter struct { inner UserPlatformQuotaRepository @@ -167,6 +190,29 @@ func (a *genericUserPlatformQuotaAdapter) ResetExpiredWindow(ctx context.Context return err } +// BatchSnapshotUsage 转换 []service.UserPlatformQuotaSnapshot → []UserPlatformQuotaSnapshot（通用 adapter）， +// 并将 repository FK sentinel 包装为 service sentinel。 +func (a *genericUserPlatformQuotaAdapter) BatchSnapshotUsage(ctx context.Context, snapshots []service.UserPlatformQuotaSnapshot, now time.Time) error { + repoSnaps := make([]UserPlatformQuotaSnapshot, len(snapshots)) + for i, s := range snapshots { + repoSnaps[i] = UserPlatformQuotaSnapshot{ + UserID: s.UserID, + Platform: s.Platform, + DailyUsageUSD: s.DailyUsageUSD, + WeeklyUsageUSD: s.WeeklyUsageUSD, + MonthlyUsageUSD: s.MonthlyUsageUSD, + DailyWindowStart: s.DailyWindowStart, + WeeklyWindowStart: s.WeeklyWindowStart, + MonthlyWindowStart: s.MonthlyWindowStart, + } + } + err := a.inner.BatchSnapshotUsage(ctx, repoSnaps, now) + if errors.Is(err, ErrUserPlatformQuotaFKViolation) { + return fmt.Errorf("%w: %v", service.ErrUserPlatformQuotaFKViolation, err) + } + return err +} + // toServiceRecord 将 repository.UserPlatformQuotaRecord 转换为 service.UserPlatformQuotaRecord。 func toServiceRecord(rec *UserPlatformQuotaRecord) *service.UserPlatformQuotaRecord { return &service.UserPlatformQuotaRecord{ diff --git a/backend/internal/repository/user_repo.go b/backend/internal/repository/user_repo.go index 610d9a7b9908..fb05452d0d7f 100644 --- a/backend/internal/repository/user_repo.go +++ b/backend/internal/repository/user_repo.go @@ -16,6 +16,7 @@ import ( dbgroup "github.com/Wei-Shaw/sub2api/ent/group" "github.com/Wei-Shaw/sub2api/ent/identityadoptiondecision" "github.com/Wei-Shaw/sub2api/ent/predicate" + "github.com/Wei-Shaw/sub2api/ent/schema/mixins" dbuser "github.com/Wei-Shaw/sub2api/ent/user" "github.com/Wei-Shaw/sub2api/ent/userallowedgroup" "github.com/Wei-Shaw/sub2api/ent/usersubscription" @@ -133,6 +134,23 @@ func (r *userRepository) GetByID(ctx context.Context, id int64) (*service.User, return out, nil } +func (r *userRepository) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) { + ctx = mixins.SkipSoftDelete(ctx) + m, err := r.client.User.Query().Where(dbuser.IDEQ(id)).Only(ctx) + if err != nil { + return nil, translatePersistenceError(err, service.ErrUserNotFound, nil) + } + out := userEntityToService(m) + groups, err := r.loadAllowedGroups(ctx, []int64{id}) + if err != nil { + return nil, err + } + if v, ok := groups[id]; ok { + out.AllowedGroups = v + } + return out, nil +} + func (r *userRepository) GetByEmail(ctx context.Context, email string) (*service.User, error) { matches, err := r.client.User.Query(). Where(userEmailLookupPredicate(email)). @@ -405,6 +423,12 @@ func (r *userRepository) List(ctx context.Context, params pagination.PaginationP } func (r *userRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters service.UserListFilters) ([]service.User, *pagination.PaginationResult, error) { + // SkipSoftDelete 仅作用于 User 身份解析（下方 Count/All）；订阅、分组等关联实体沿用原始 ctx，避免穿透到这些同样带软删除的实体而带出已删除行。 + userCtx := ctx + if filters.IncludeDeleted { + userCtx = mixins.SkipSoftDelete(ctx) + } + q := r.client.User.Query() if filters.Status != "" { @@ -445,7 +469,7 @@ func (r *userRepository) ListWithFilters(ctx context.Context, params pagination. q = q.Where(dbuser.IDIn(allowedUserIDs...)) } - total, err := q.Clone().Count(ctx) + total, err := q.Clone().Count(userCtx) if err != nil { return nil, nil, err } @@ -457,7 +481,7 @@ func (r *userRepository) ListWithFilters(ctx context.Context, params pagination. usersQuery = usersQuery.Order(order) } - users, err := usersQuery.All(ctx) + users, err := usersQuery.All(userCtx) if err != nil { return nil, nil, err } diff --git a/backend/internal/repository/user_repo_include_deleted_integration_test.go b/backend/internal/repository/user_repo_include_deleted_integration_test.go new file mode 100644 index 000000000000..014b24f94983 --- /dev/null +++ b/backend/internal/repository/user_repo_include_deleted_integration_test.go @@ -0,0 +1,69 @@ +//go:build integration + +package repository + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestUserRepo_ListWithFilters_IncludeDeleted(t *testing.T) { + ctx := context.Background() + tx := testEntTx(t) + client := tx.Client() + repo := NewUserRepository(client, integrationDB) + + active := mustCreateUser(t, client, &service.User{Email: "shared-keyword-active@test.com"}) + deleted := mustCreateUser(t, client, &service.User{Email: "shared-keyword-deleted@test.com"}) + require.NoError(t, client.User.DeleteOneID(deleted.ID).Exec(ctx)) + + params := pagination.PaginationParams{Page: 1, PageSize: 50, SortBy: "email", SortOrder: "asc"} + + // 默认（不含已删）：只返回活跃用户。 + usersDefault, resDefault, err := repo.ListWithFilters(ctx, params, + service.UserListFilters{Search: "shared-keyword-"}) + require.NoError(t, err) + require.Len(t, usersDefault, 1) + require.Equal(t, active.ID, usersDefault[0].ID) + require.EqualValues(t, 1, resDefault.Total) + + // IncludeDeleted=true：两个都返回，且 Total 与结果集一致。 + usersAll, resAll, err := repo.ListWithFilters(ctx, params, + service.UserListFilters{Search: "shared-keyword-", IncludeDeleted: true}) + require.NoError(t, err) + require.Len(t, usersAll, 2) + require.EqualValues(t, 2, resAll.Total, "Count 必须与结果集行数一致") + + var delUser *service.User + for i := range usersAll { + if usersAll[i].ID == deleted.ID { + delUser = &usersAll[i] + } + } + require.NotNil(t, delUser) + require.NotNil(t, delUser.DeletedAt) +} + +func TestUserRepo_GetByIDIncludeDeleted(t *testing.T) { + ctx := context.Background() + tx := testEntTx(t) + client := tx.Client() + repo := NewUserRepository(client, integrationDB) + + u := mustCreateUser(t, client, &service.User{Email: "getbyid-deleted@test.com"}) + require.NoError(t, client.User.DeleteOneID(u.ID).Exec(ctx)) + + // 默认 GetByID：找不到（被软删过滤）。 + _, err := repo.GetByID(ctx, u.ID) + require.ErrorIs(t, err, service.ErrUserNotFound) + + // GetByIDIncludeDeleted：找得到，且 DeletedAt 非空。 + got, err := repo.GetByIDIncludeDeleted(ctx, u.ID) + require.NoError(t, err) + require.Equal(t, "getbyid-deleted@test.com", got.Email) + require.NotNil(t, got.DeletedAt) +} diff --git a/backend/internal/server/api_contract_test.go b/backend/internal/server/api_contract_test.go index f04f56cd3114..917eef643828 100644 --- a/backend/internal/server/api_contract_test.go +++ b/backend/internal/server/api_contract_test.go @@ -1500,6 +1500,10 @@ func (r *stubUserRepo) DisableTotp(ctx context.Context, userID int64) error { return errors.New("not implemented") } +func (r *stubUserRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) { + panic("unexpected GetByIDIncludeDeleted call") +} + type stubApiKeyCache struct{} func (stubApiKeyCache) GetCreateAttemptCount(ctx context.Context, userID int64) (int, error) { diff --git a/backend/internal/server/middleware/admin_auth_test.go b/backend/internal/server/middleware/admin_auth_test.go index 303d0db879df..3110c6c10443 100644 --- a/backend/internal/server/middleware/admin_auth_test.go +++ b/backend/internal/server/middleware/admin_auth_test.go @@ -236,3 +236,7 @@ func (s *stubUserRepo) EnableTotp(ctx context.Context, userID int64) error { func (s *stubUserRepo) DisableTotp(ctx context.Context, userID int64) error { panic("unexpected DisableTotp call") } + +func (s *stubUserRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) { + panic("unexpected GetByIDIncludeDeleted call") +} diff --git a/backend/internal/server/middleware/client_request_id.go b/backend/internal/server/middleware/client_request_id.go index 6838d6afd50c..5f886646eee9 100644 --- a/backend/internal/server/middleware/client_request_id.go +++ b/backend/internal/server/middleware/client_request_id.go @@ -11,6 +11,8 @@ import ( "go.uber.org/zap" ) +const clientRequestIDHeader = "X-Client-Request-ID" + // ClientRequestID ensures every request has a unique client_request_id in request.Context(). // // This is used by the Ops monitoring module for end-to-end request correlation. @@ -21,12 +23,14 @@ func ClientRequestID() gin.HandlerFunc { return } - if v := c.Request.Context().Value(ctxkey.ClientRequestID); v != nil { + if v, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(v) != "" { + c.Header(clientRequestIDHeader, strings.TrimSpace(v)) c.Next() return } id := uuid.New().String() + c.Header(clientRequestIDHeader, id) ctx := context.WithValue(c.Request.Context(), ctxkey.ClientRequestID, id) requestLogger := logger.FromContext(ctx).With(zap.String("client_request_id", strings.TrimSpace(id))) ctx = logger.IntoContext(ctx, requestLogger) diff --git a/backend/internal/server/middleware/client_request_id_test.go b/backend/internal/server/middleware/client_request_id_test.go new file mode 100644 index 000000000000..394c16125d0f --- /dev/null +++ b/backend/internal/server/middleware/client_request_id_test.go @@ -0,0 +1,50 @@ +package middleware + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestClientRequestIDGeneratesAndExposesID(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(ClientRequestID()) + router.GET("/", func(c *gin.Context) { + value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string) + c.String(http.StatusOK, value) + }) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.NotEmpty(t, w.Body.String()) + require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader)) +} + +func TestClientRequestIDPreservesExistingContextID(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(ClientRequestID()) + router.GET("/", func(c *gin.Context) { + value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string) + c.String(http.StatusOK, value) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/", nil) + req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "existing-client-request-id")) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Equal(t, "existing-client-request-id", w.Body.String()) + require.Equal(t, "existing-client-request-id", w.Header().Get(clientRequestIDHeader)) +} diff --git a/backend/internal/server/routes/admin.go b/backend/internal/server/routes/admin.go index 2301adda6bf9..9a3253b55dd6 100644 --- a/backend/internal/server/routes/admin.go +++ b/backend/internal/server/routes/admin.go @@ -302,6 +302,7 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers) { accounts.GET("/:id/temp-unschedulable", h.Admin.Account.GetTempUnschedulable) accounts.DELETE("/:id/temp-unschedulable", h.Admin.Account.ClearTempUnschedulable) accounts.POST("/:id/schedulable", h.Admin.Account.SetSchedulable) + accounts.POST("/models/sync-upstream-preview", h.Admin.Account.SyncUpstreamModelsPreview) accounts.GET("/:id/models", h.Admin.Account.GetAvailableModels) accounts.POST("/:id/models/sync-upstream", h.Admin.Account.SyncUpstreamModels) accounts.POST("/batch", h.Admin.Account.BatchCreate) diff --git a/backend/internal/service/account_test_service_openai_test.go b/backend/internal/service/account_test_service_openai_test.go index 910567fb25b6..970c723a9b63 100644 --- a/backend/internal/service/account_test_service_openai_test.go +++ b/backend/internal/service/account_test_service_openai_test.go @@ -132,7 +132,7 @@ func TestAccountTestService_OpenAISuccessPersistsSnapshotFromHeaders(t *testing. require.Len(t, upstream.requests, 1) require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.requests[0].Context())) require.NotEmpty(t, repo.updatedExtra) - require.Equal(t, 42.0, repo.updatedExtra["codex_5h_used_percent"]) + require.Equal(t, 58.0, repo.updatedExtra["codex_5h_used_percent"]) require.Equal(t, 88.0, repo.updatedExtra["codex_7d_used_percent"]) require.Contains(t, recorder.Body.String(), "test_complete") } @@ -170,7 +170,7 @@ func TestAccountTestService_OpenAI429PersistsSnapshotAndRateLimitState(t *testin resp.Header.Set("x-codex-primary-used-percent", "100") resp.Header.Set("x-codex-primary-reset-after-seconds", "604800") resp.Header.Set("x-codex-primary-window-minutes", "10080") - resp.Header.Set("x-codex-secondary-used-percent", "100") + resp.Header.Set("x-codex-secondary-used-percent", "0") resp.Header.Set("x-codex-secondary-reset-after-seconds", "18000") resp.Header.Set("x-codex-secondary-window-minutes", "300") diff --git a/backend/internal/service/account_usage_service_test.go b/backend/internal/service/account_usage_service_test.go index e0390c4c2a77..5f37aadbab7b 100644 --- a/backend/internal/service/account_usage_service_test.go +++ b/backend/internal/service/account_usage_service_test.go @@ -73,7 +73,7 @@ func TestExtractOpenAICodexProbeUpdatesAccepts429WithCodexHeaders(t *testing.T) headers.Set("x-codex-primary-used-percent", "100") headers.Set("x-codex-primary-reset-after-seconds", "604800") headers.Set("x-codex-primary-window-minutes", "10080") - headers.Set("x-codex-secondary-used-percent", "100") + headers.Set("x-codex-secondary-used-percent", "0") headers.Set("x-codex-secondary-reset-after-seconds", "18000") headers.Set("x-codex-secondary-window-minutes", "300") @@ -92,6 +92,33 @@ func TestExtractOpenAICodexProbeUpdatesAccepts429WithCodexHeaders(t *testing.T) } } +func TestBuildCodexUsageProgressFromExtra_UsesCanonicalUsedPercent(t *testing.T) { + t.Parallel() + now := time.Date(2026, 5, 30, 7, 4, 9, 0, time.UTC) + extra := map[string]any{ + "codex_5h_used_percent": 94.0, + "codex_5h_reset_at": now.Add(2 * time.Hour).Format(time.RFC3339), + "codex_7d_used_percent": 93.0, + "codex_7d_reset_at": now.Add(5 * 24 * time.Hour).Format(time.RFC3339), + } + + fiveHour := buildCodexUsageProgressFromExtra(extra, "5h", now) + if fiveHour == nil { + t.Fatal("expected non-nil 5h progress") + } + if fiveHour.Utilization != 94.0 { + t.Fatalf("5h Utilization = %v, want 94", fiveHour.Utilization) + } + + sevenDay := buildCodexUsageProgressFromExtra(extra, "7d", now) + if sevenDay == nil { + t.Fatal("expected non-nil 7d progress") + } + if sevenDay.Utilization != 93.0 { + t.Fatalf("7d Utilization = %v, want 93", sevenDay.Utilization) + } +} + func TestAccountUsageService_PersistOpenAICodexProbeSnapshotOnlyUpdatesExtra(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index d46b636f2cb1..00205d1f7926 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -33,6 +33,7 @@ type AdminService interface { // User management ListUsers(ctx context.Context, page, pageSize int, filters UserListFilters, sortBy, sortOrder string) ([]User, int64, error) GetUser(ctx context.Context, id int64) (*User, error) + GetUserIncludeDeleted(ctx context.Context, id int64) (*User, error) CreateUser(ctx context.Context, input *CreateUserInput) (*User, error) UpdateUser(ctx context.Context, id int64, input *UpdateUserInput) (*User, error) DeleteUser(ctx context.Context, id int64) error @@ -128,7 +129,7 @@ type CreateUserInput struct { Password string Username string Notes string - Balance float64 + Balance *float64 Concurrency int RPMLimit int AllowedGroups []int64 @@ -674,13 +675,24 @@ func (s *adminServiceImpl) GetUser(ctx context.Context, id int64) (*User, error) return user, nil } +func (s *adminServiceImpl) GetUserIncludeDeleted(ctx context.Context, id int64) (*User, error) { + return s.userRepo.GetByIDIncludeDeleted(ctx, id) +} + func (s *adminServiceImpl) CreateUser(ctx context.Context, input *CreateUserInput) (*User, error) { + balance := 0.0 + if input.Balance != nil { + balance = *input.Balance + } else if s.settingService != nil { + balance = s.settingService.GetDefaultBalance(ctx) + } + user := &User{ Email: input.Email, Username: input.Username, Notes: input.Notes, Role: RoleUser, // Always create as regular user, never admin - Balance: input.Balance, + Balance: balance, Concurrency: input.Concurrency, RPMLimit: input.RPMLimit, Status: StatusActive, diff --git a/backend/internal/service/admin_service_apikey_test.go b/backend/internal/service/admin_service_apikey_test.go index 3b3dbc21cfb4..f26fadb8c8e8 100644 --- a/backend/internal/service/admin_service_apikey_test.go +++ b/backend/internal/service/admin_service_apikey_test.go @@ -69,8 +69,12 @@ func (s *userRepoStubForGroupUpdate) UpdateConcurrency(context.Context, int64, i panic("unexpected") } -func (s *userRepoStubForGroupUpdate) BatchSetConcurrency(context.Context, []int64, int) (int, error) { return 0, nil } -func (s *userRepoStubForGroupUpdate) BatchAddConcurrency(context.Context, []int64, int) (int, error) { return 0, nil } +func (s *userRepoStubForGroupUpdate) BatchSetConcurrency(context.Context, []int64, int) (int, error) { + return 0, nil +} +func (s *userRepoStubForGroupUpdate) BatchAddConcurrency(context.Context, []int64, int) (int, error) { + return 0, nil +} func (s *userRepoStubForGroupUpdate) ExistsByEmail(context.Context, string) (bool, error) { panic("unexpected") } @@ -82,6 +86,9 @@ func (s *userRepoStubForGroupUpdate) UpdateTotpSecret(context.Context, int64, *s } func (s *userRepoStubForGroupUpdate) EnableTotp(context.Context, int64) error { panic("unexpected") } func (s *userRepoStubForGroupUpdate) DisableTotp(context.Context, int64) error { panic("unexpected") } +func (s *userRepoStubForGroupUpdate) GetByIDIncludeDeleted(ctx context.Context, id int64) (*User, error) { + panic("unexpected GetByIDIncludeDeleted call") +} func (s *userRepoStubForGroupUpdate) ListUserAuthIdentities(context.Context, int64) ([]UserAuthIdentityRecord, error) { panic("unexpected") } diff --git a/backend/internal/service/admin_service_create_user_test.go b/backend/internal/service/admin_service_create_user_test.go index c5b1e38d3903..5e9578abecca 100644 --- a/backend/internal/service/admin_service_create_user_test.go +++ b/backend/internal/service/admin_service_create_user_test.go @@ -14,13 +14,14 @@ import ( func TestAdminService_CreateUser_Success(t *testing.T) { repo := &userRepoStub{nextID: 10} svc := &adminServiceImpl{userRepo: repo} + balance := 12.5 input := &CreateUserInput{ Email: "user@test.com", Password: "strong-pass", Username: "tester", Notes: "note", - Balance: 12.5, + Balance: &balance, Concurrency: 7, AllowedGroups: []int64{3, 5}, } @@ -32,7 +33,7 @@ func TestAdminService_CreateUser_Success(t *testing.T) { require.Equal(t, input.Email, user.Email) require.Equal(t, input.Username, user.Username) require.Equal(t, input.Notes, user.Notes) - require.Equal(t, input.Balance, user.Balance) + require.Equal(t, balance, user.Balance) require.Equal(t, input.Concurrency, user.Concurrency) require.Equal(t, input.AllowedGroups, user.AllowedGroups) require.Equal(t, RoleUser, user.Role) @@ -42,6 +43,56 @@ func TestAdminService_CreateUser_Success(t *testing.T) { require.Equal(t, user, repo.created[0]) } +func TestAdminService_CreateUser_UsesDefaultBalanceWhenBalanceOmitted(t *testing.T) { + repo := &userRepoStub{nextID: 11} + cfg := &config.Config{ + Default: config.DefaultConfig{ + UserBalance: 0, + }, + } + settingService := NewSettingService(&settingRepoStub{values: map[string]string{ + SettingKeyDefaultBalance: "0.02", + }}, cfg) + svc := &adminServiceImpl{userRepo: repo, settingService: settingService} + + user, err := svc.CreateUser(context.Background(), &CreateUserInput{ + Email: "default-balance@test.com", + Password: "strong-pass", + }) + + require.NoError(t, err) + require.NotNil(t, user) + require.Equal(t, 0.02, user.Balance) + require.Len(t, repo.created, 1) + require.Equal(t, 0.02, repo.created[0].Balance) +} + +func TestAdminService_CreateUser_ExplicitZeroBalanceOverridesDefault(t *testing.T) { + repo := &userRepoStub{nextID: 12} + cfg := &config.Config{ + Default: config.DefaultConfig{ + UserBalance: 0, + }, + } + settingService := NewSettingService(&settingRepoStub{values: map[string]string{ + SettingKeyDefaultBalance: "0.02", + }}, cfg) + svc := &adminServiceImpl{userRepo: repo, settingService: settingService} + balance := 0.0 + + user, err := svc.CreateUser(context.Background(), &CreateUserInput{ + Email: "zero-balance@test.com", + Password: "strong-pass", + Balance: &balance, + }) + + require.NoError(t, err) + require.NotNil(t, user) + require.Equal(t, 0.0, user.Balance) + require.Len(t, repo.created, 1) + require.Equal(t, 0.0, repo.created[0].Balance) +} + func TestAdminService_CreateUser_EmailExists(t *testing.T) { repo := &userRepoStub{createErr: ErrEmailExists} svc := &adminServiceImpl{userRepo: repo} diff --git a/backend/internal/service/admin_service_delete_test.go b/backend/internal/service/admin_service_delete_test.go index d01b11e67241..150c4f53413f 100644 --- a/backend/internal/service/admin_service_delete_test.go +++ b/backend/internal/service/admin_service_delete_test.go @@ -173,6 +173,10 @@ func (s *userRepoStub) DisableTotp(ctx context.Context, userID int64) error { panic("unexpected DisableTotp call") } +func (s *userRepoStub) GetByIDIncludeDeleted(ctx context.Context, id int64) (*User, error) { + return s.GetByID(ctx, id) +} + type groupRepoStub struct { affectedUserIDs []int64 deleteErr error @@ -471,10 +475,22 @@ func (s *billingCacheStub) DeleteUserPlatformQuotaCache(ctx context.Context, use panic("unexpected DeleteUserPlatformQuotaCache call") } -func (s *billingCacheStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { +func (s *billingCacheStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { panic("unexpected IncrUserPlatformQuotaUsageCache call") } +func (s *billingCacheStub) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) { + panic("unexpected PopDirtyUserPlatformQuotaKeys call") +} + +func (s *billingCacheStub) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error { + panic("unexpected ReaddDirtyUserPlatformQuotaKeys call") +} + +func (s *billingCacheStub) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + panic("unexpected BatchGetUserPlatformQuotaCache call") +} + func waitForInvalidations(t *testing.T, ch <-chan subscriptionInvalidateCall, expected int) []subscriptionInvalidateCall { t.Helper() calls := make([]subscriptionInvalidateCall, 0, expected) diff --git a/backend/internal/service/admin_service_email_identity_sync_test.go b/backend/internal/service/admin_service_email_identity_sync_test.go index c791b747cf7b..c3737f5a1a6e 100644 --- a/backend/internal/service/admin_service_email_identity_sync_test.go +++ b/backend/internal/service/admin_service_email_identity_sync_test.go @@ -113,8 +113,12 @@ func (s *emailSyncRepoStub) RemoveGroupFromAllowedGroups(context.Context, int64) return 0, nil } -func (s *emailSyncRepoStub) BatchSetConcurrency(context.Context, []int64, int) (int, error) { return 0, nil } -func (s *emailSyncRepoStub) BatchAddConcurrency(context.Context, []int64, int) (int, error) { return 0, nil } +func (s *emailSyncRepoStub) BatchSetConcurrency(context.Context, []int64, int) (int, error) { + return 0, nil +} +func (s *emailSyncRepoStub) BatchAddConcurrency(context.Context, []int64, int) (int, error) { + return 0, nil +} func (s *emailSyncRepoStub) AddGroupToAllowedGroups(context.Context, int64, int64) error { return nil } @@ -133,6 +137,9 @@ func (s *emailSyncRepoStub) UpdateTotpSecret(context.Context, int64, *string) er func (s *emailSyncRepoStub) EnableTotp(context.Context, int64) error { return nil } func (s *emailSyncRepoStub) DisableTotp(context.Context, int64) error { return nil } +func (s *emailSyncRepoStub) GetByIDIncludeDeleted(ctx context.Context, id int64) (*User, error) { + return s.GetByID(ctx, id) +} func (s *emailSyncRepoStub) EnsureEmailAuthIdentity(_ context.Context, userID int64, email string) error { s.ensureCalls = append(s.ensureCalls, ensureEmailCall{userID: userID, email: email}) diff --git a/backend/internal/service/admin_service_get_deleted_test.go b/backend/internal/service/admin_service_get_deleted_test.go new file mode 100644 index 000000000000..6ad17f60671a --- /dev/null +++ b/backend/internal/service/admin_service_get_deleted_test.go @@ -0,0 +1,22 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestAdminService_GetUserIncludeDeleted(t *testing.T) { + ts := time.Date(2026, 5, 28, 0, 0, 0, 0, time.UTC) + repo := &userRepoStub{user: &User{ID: 7, Email: "del@test.com", DeletedAt: &ts}} + svc := &adminServiceImpl{userRepo: repo} + + got, err := svc.GetUserIncludeDeleted(context.Background(), 7) + require.NoError(t, err) + require.Equal(t, int64(7), got.ID) + require.NotNil(t, got.DeletedAt) +} diff --git a/backend/internal/service/anthropic_session.go b/backend/internal/service/anthropic_session.go index 26544c68cd48..bca8cc7ffba7 100644 --- a/backend/internal/service/anthropic_session.go +++ b/backend/internal/service/anthropic_session.go @@ -4,6 +4,8 @@ import ( "encoding/json" "strings" "time" + + "github.com/tidwall/gjson" ) // Anthropic 会话 Fallback 相关常量 @@ -30,30 +32,39 @@ func BuildAnthropicDigestChain(parsed *ParsedRequest) string { var parts []string - // 1. system prompt - if parsed.System != nil { - systemData, _ := json.Marshal(parsed.System) - if len(systemData) > 0 && string(systemData) != "null" { - parts = append(parts, "s:"+shortHash(systemData)) - } + if systemRaw := parsed.SystemRaw(); len(systemRaw) > 0 && string(systemRaw) != "null" { + parts = append(parts, "s:"+shortHash(canonicalAnthropicDigestJSON(systemRaw))) } - // 2. messages - for _, msg := range parsed.Messages { - msgMap, ok := msg.(map[string]any) - if !ok { - continue - } - role, _ := msgMap["role"].(string) - prefix := rolePrefix(role) - content := msgMap["content"] - contentData, _ := json.Marshal(content) - parts = append(parts, prefix+":"+shortHash(contentData)) + messages := parsed.MessagesRaw() + if len(messages) > 0 { + gjson.ParseBytes(messages).ForEach(func(_, msg gjson.Result) bool { + prefix := rolePrefix(msg.Get("role").String()) + content := msg.Get("content") + parts = append(parts, prefix+":"+shortHash(canonicalAnthropicDigestJSON([]byte(content.Raw)))) + return true + }) } return strings.Join(parts, "-") } +// canonicalAnthropicDigestJSON 保持 digest 对 JSON key 顺序和空白不敏感。 +func canonicalAnthropicDigestJSON(raw []byte) []byte { + if len(raw) == 0 { + return raw + } + var value any + if err := json.Unmarshal(raw, &value); err != nil { + return raw + } + canonical, err := json.Marshal(value) + if err != nil { + return raw + } + return canonical +} + // rolePrefix 将 Anthropic 的 role 映射为单字符前缀 func rolePrefix(role string) string { switch role { diff --git a/backend/internal/service/anthropic_session_test.go b/backend/internal/service/anthropic_session_test.go index 10406643bda8..4d88fc8b3419 100644 --- a/backend/internal/service/anthropic_session_test.go +++ b/backend/internal/service/anthropic_session_test.go @@ -1,3 +1,5 @@ +//go:build unit + package service import ( @@ -5,6 +7,15 @@ import ( "testing" ) +func mustParseAnthropicDigestRequest(t *testing.T, body string) *ParsedRequest { + t.Helper() + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(body)), "") + if err != nil { + t.Fatalf("ParseGatewayRequest failed: %v", err) + } + return parsed +} + func TestBuildAnthropicDigestChain_NilRequest(t *testing.T) { result := BuildAnthropicDigestChain(nil) if result != "" { @@ -13,9 +24,7 @@ func TestBuildAnthropicDigestChain_NilRequest(t *testing.T) { } func TestBuildAnthropicDigestChain_EmptyMessages(t *testing.T) { - parsed := &ParsedRequest{ - Messages: []any{}, - } + parsed := mustParseAnthropicDigestRequest(t, `{"messages":[]}`) result := BuildAnthropicDigestChain(parsed) if result != "" { t.Errorf("expected empty string for empty messages, got: %s", result) @@ -23,11 +32,7 @@ func TestBuildAnthropicDigestChain_EmptyMessages(t *testing.T) { } func TestBuildAnthropicDigestChain_SingleUserMessage(t *testing.T) { - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 1 { @@ -39,12 +44,7 @@ func TestBuildAnthropicDigestChain_SingleUserMessage(t *testing.T) { } func TestBuildAnthropicDigestChain_UserAndAssistant(t *testing.T) { - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi there"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi there"}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 2 { @@ -59,12 +59,7 @@ func TestBuildAnthropicDigestChain_UserAndAssistant(t *testing.T) { } func TestBuildAnthropicDigestChain_WithSystemString(t *testing.T) { - parsed := &ParsedRequest{ - System: "You are a helpful assistant", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 2 { @@ -79,14 +74,7 @@ func TestBuildAnthropicDigestChain_WithSystemString(t *testing.T) { } func TestBuildAnthropicDigestChain_WithSystemContentBlocks(t *testing.T) { - parsed := &ParsedRequest{ - System: []any{ - map[string]any{"type": "text", "text": "You are a helpful assistant"}, - }, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"system":[{"type":"text","text":"You are a helpful assistant"}],"messages":[{"role":"user","content":"hello"}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 2 { @@ -100,74 +88,33 @@ func TestBuildAnthropicDigestChain_WithSystemContentBlocks(t *testing.T) { func TestBuildAnthropicDigestChain_ConversationPrefixRelationship(t *testing.T) { // 核心测试：验证对话增长时链的前缀关系 // 上一轮的完整链一定是下一轮链的前缀 - system := "You are a helpful assistant" - - // 第 1 轮: system + user - round1 := &ParsedRequest{ - System: system, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + round1 := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"}]}`) chain1 := BuildAnthropicDigestChain(round1) - // 第 2 轮: system + user + assistant + user - round2 := &ParsedRequest{ - System: system, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi there"}, - map[string]any{"role": "user", "content": "how are you?"}, - }, - } + round2 := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi there"},{"role":"user","content":"how are you?"}]}`) chain2 := BuildAnthropicDigestChain(round2) - // 第 3 轮: system + user + assistant + user + assistant + user - round3 := &ParsedRequest{ - System: system, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi there"}, - map[string]any{"role": "user", "content": "how are you?"}, - map[string]any{"role": "assistant", "content": "I'm doing well"}, - map[string]any{"role": "user", "content": "great"}, - }, - } + round3 := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi there"},{"role":"user","content":"how are you?"},{"role":"assistant","content":"I'm doing well"},{"role":"user","content":"great"}]}`) chain3 := BuildAnthropicDigestChain(round3) t.Logf("Chain1: %s", chain1) t.Logf("Chain2: %s", chain2) t.Logf("Chain3: %s", chain3) - // chain1 是 chain2 的前缀 if !strings.HasPrefix(chain2, chain1) { t.Errorf("chain1 should be prefix of chain2:\n chain1: %s\n chain2: %s", chain1, chain2) } - - // chain2 是 chain3 的前缀 if !strings.HasPrefix(chain3, chain2) { t.Errorf("chain2 should be prefix of chain3:\n chain2: %s\n chain3: %s", chain2, chain3) } - - // chain1 也是 chain3 的前缀（传递性） if !strings.HasPrefix(chain3, chain1) { t.Errorf("chain1 should be prefix of chain3:\n chain1: %s\n chain3: %s", chain1, chain3) } } func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *testing.T) { - parsed1 := &ParsedRequest{ - System: "System A", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } - parsed2 := &ParsedRequest{ - System: "System B", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed1 := mustParseAnthropicDigestRequest(t, `{"system":"System A","messages":[{"role":"user","content":"hello"}]}`) + parsed2 := mustParseAnthropicDigestRequest(t, `{"system":"System B","messages":[{"role":"user","content":"hello"}]}`) chain1 := BuildAnthropicDigestChain(parsed1) chain2 := BuildAnthropicDigestChain(parsed2) @@ -176,7 +123,6 @@ func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *test t.Error("Different system prompts should produce different chains") } - // 但 user 部分的 hash 应该相同 parts1 := splitChain(chain1) parts2 := splitChain(chain2) if parts1[1] != parts2[1] { @@ -185,20 +131,8 @@ func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *test } func TestBuildAnthropicDigestChain_DifferentContentProducesDifferentChain(t *testing.T) { - parsed1 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "ORIGINAL reply"}, - map[string]any{"role": "user", "content": "next"}, - }, - } - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "TAMPERED reply"}, - map[string]any{"role": "user", "content": "next"}, - }, - } + parsed1 := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"ORIGINAL reply"},{"role":"user","content":"next"}]}`) + parsed2 := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"TAMPERED reply"},{"role":"user","content":"next"}]}`) chain1 := BuildAnthropicDigestChain(parsed1) chain2 := BuildAnthropicDigestChain(parsed2) @@ -209,24 +143,16 @@ func TestBuildAnthropicDigestChain_DifferentContentProducesDifferentChain(t *tes parts1 := splitChain(chain1) parts2 := splitChain(chain2) - // 第一个 user message hash 应该相同 if parts1[0] != parts2[0] { t.Error("First user message hash should be the same") } - // assistant reply hash 应该不同 if parts1[1] == parts2[1] { t.Error("Assistant reply hash should differ") } } func TestBuildAnthropicDigestChain_Deterministic(t *testing.T) { - parsed := &ParsedRequest{ - System: "test system", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi"}, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"system":"test system","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi"}]}`) chain1 := BuildAnthropicDigestChain(parsed) chain2 := BuildAnthropicDigestChain(parsed) @@ -236,6 +162,18 @@ func TestBuildAnthropicDigestChain_Deterministic(t *testing.T) { } } +func TestBuildAnthropicDigestChain_CanonicalJSON(t *testing.T) { + parsed1 := mustParseAnthropicDigestRequest(t, `{"system":[{"type":"text","text":"system"}],"messages":[{"role":"user","content":{"type":"text","text":"hello"}}]}`) + parsed2 := mustParseAnthropicDigestRequest(t, `{"system":[{"text":"system","type":"text"}],"messages":[{"role":"user","content":{"text":"hello","type":"text"}}]}`) + + chain1 := BuildAnthropicDigestChain(parsed1) + chain2 := BuildAnthropicDigestChain(parsed2) + + if chain1 != chain2 { + t.Errorf("semantically equivalent JSON should produce same chain: %s vs %s", chain1, chain2) + } +} + func TestGenerateAnthropicDigestSessionKey(t *testing.T) { tests := []struct { name string @@ -278,7 +216,6 @@ func TestGenerateAnthropicDigestSessionKey(t *testing.T) { }) } - // 验证不同 uuid 产生不同 sessionKey t.Run("different uuid different key", func(t *testing.T) { hash := "sameprefix123456" result1 := GenerateAnthropicDigestSessionKey(hash, "uuid0001-session-a") @@ -297,18 +234,7 @@ func TestAnthropicSessionTTL(t *testing.T) { } func TestBuildAnthropicDigestChain_ContentBlocks(t *testing.T) { - // 测试 content 为 content blocks 数组的情况 - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{"type": "text", "text": "describe this image"}, - map[string]any{"type": "image", "source": map[string]any{"type": "base64"}}, - }, - }, - }, - } + parsed := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":[{"type":"text","text":"describe this image"},{"type":"image","source":{"type":"base64"}}]}]}`) result := BuildAnthropicDigestChain(parsed) parts := splitChain(result) if len(parts) != 1 { diff --git a/backend/internal/service/antigravity_default_test_stubs_test.go b/backend/internal/service/antigravity_default_test_stubs_test.go new file mode 100644 index 000000000000..d3c2c57a272b --- /dev/null +++ b/backend/internal/service/antigravity_default_test_stubs_test.go @@ -0,0 +1,61 @@ +//go:build !unit + +package service + +import ( + "context" + "time" +) + +type defaultRateLimitCall struct { + accountID int64 + resetAt time.Time +} + +type defaultModelRateLimitCall struct { + accountID int64 + modelKey string + resetAt time.Time +} + +type defaultExtraUpdateCall struct { + accountID int64 + updates map[string]any +} + +type stubAntigravityAccountRepo struct { + AccountRepository + rateCalls []defaultRateLimitCall + modelRateLimitCalls []defaultModelRateLimitCall + extraUpdateCalls []defaultExtraUpdateCall +} + +func (s *stubAntigravityAccountRepo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error { + s.rateCalls = append(s.rateCalls, defaultRateLimitCall{accountID: id, resetAt: resetAt}) + return nil +} + +func (s *stubAntigravityAccountRepo) SetModelRateLimit(_ context.Context, id int64, modelKey string, resetAt time.Time, _ ...string) error { + s.modelRateLimitCalls = append(s.modelRateLimitCalls, defaultModelRateLimitCall{accountID: id, modelKey: modelKey, resetAt: resetAt}) + return nil +} + +func (s *stubAntigravityAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error { + s.extraUpdateCalls = append(s.extraUpdateCalls, defaultExtraUpdateCall{accountID: id, updates: updates}) + return nil +} + +type defaultDeleteSessionCall struct { + groupID int64 + sessionHash string +} + +type stubSmartRetryCache struct { + GatewayCache + deleteCalls []defaultDeleteSessionCall +} + +func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID int64, sessionHash string) error { + c.deleteCalls = append(c.deleteCalls, defaultDeleteSessionCall{groupID: groupID, sessionHash: sessionHash}) + return nil +} diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index 2b849bdd8de0..f79d20a2ca90 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -228,12 +228,11 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration, truncateForLog(respBody, 200)) resetAt := time.Now().Add(rateLimitDuration) - if !setModelRateLimitByModelName(p.ctx, p.accountRepo, p.account.ID, modelName, p.prefix, resp.StatusCode, resetAt, false) { + if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, false) { p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited account=%d (no model mapping)", p.prefix, resp.StatusCode, p.account.ID) - } else { - s.updateAccountModelRateLimitInCache(p.ctx, p.account, modelName, resetAt) } + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) // 返回账号切换信号，让上层切换账号重试 return &smartRetryResult{ @@ -392,20 +391,10 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, rateLimitDuration, truncateForLog(retryBody, 200)) resetAt := time.Now().Add(rateLimitDuration) - if p.accountRepo != nil && modelName != "" { - if err := p.accountRepo.SetModelRateLimit(p.ctx, p.account.ID, modelName, resetAt); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_failed model=%s error=%v", p.prefix, resp.StatusCode, modelName, err) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_after_smart_retry model=%s account=%d reset_in=%v", - p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration) - s.updateAccountModelRateLimitInCache(p.ctx, p.account, modelName, resetAt) - } - } + s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, true) // 清除粘性会话绑定，避免下次请求仍命中限流账号 - if s.cache != nil && p.sessionHash != "" { - _ = s.cache.DeleteSessionAccountID(p.ctx, p.groupID, p.sessionHash) - } + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) // 返回账号切换信号，让上层切换账号重试 return &smartRetryResult{ @@ -662,7 +651,7 @@ urlFallbackLoop: // 统一处理错误响应 if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) { @@ -875,6 +864,22 @@ type AntigravityGatewayService struct { internal500Cache Internal500CounterCache // INTERNAL 500 渐进惩罚计数器 } +func (s *AntigravityGatewayService) upstreamErrorBodyReadLimit() int64 { + limit := gatewayUpstreamErrorBodyReadLimit + if s != nil && s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.LogUpstreamErrorBody && s.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { + limit = int64(s.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return limit +} + +func (s *AntigravityGatewayService) readUpstreamErrorBody(resp *http.Response) []byte { + if resp == nil || resp.Body == nil { + return nil + } + body, _ := io.ReadAll(io.LimitReader(resp.Body, s.upstreamErrorBodyReadLimit())) + return body +} + func NewAntigravityGatewayService( accountRepo AccountRepository, cache GatewayCache, @@ -938,8 +943,14 @@ func (s *AntigravityGatewayService) checkErrorPolicy(ctx context.Context, accoun func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParams, statusCode int, headers http.Header, respBody []byte) (handled bool, outStatus int, retErr error) { switch s.checkErrorPolicy(p.ctx, p.account, statusCode, respBody) { case ErrorPolicySkipped: + if s.handleAntigravityModelRateLimitBeforePolicy(p, statusCode, headers, respBody) { + return true, statusCode, nil + } return true, http.StatusInternalServerError, nil case ErrorPolicyMatched: + if s.handleAntigravityModelRateLimitBeforePolicy(p, statusCode, headers, respBody) { + return true, statusCode, nil + } _ = p.handleError(p.ctx, p.prefix, p.account, statusCode, headers, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) return true, statusCode, nil @@ -951,6 +962,31 @@ func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParam return false, statusCode, nil } +func (s *AntigravityGatewayService) handleAntigravityModelRateLimitBeforePolicy(p antigravityRetryLoopParams, statusCode int, headers http.Header, respBody []byte) bool { + if statusCode != http.StatusTooManyRequests && statusCode != http.StatusServiceUnavailable { + return false + } + if p.account == nil || p.account.Platform != PlatformAntigravity { + return false + } + _, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody) + if isModelCapacityExhausted || !shouldRateLimitModel || strings.TrimSpace(modelName) == "" { + return false + } + rateLimitDuration := waitDuration + if rateLimitDuration <= 0 { + rateLimitDuration = antigravityDefaultRateLimitDuration + } + resetAt := time.Now().Add(rateLimitDuration) + if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, statusCode, resetAt, false) { + return false + } + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_before_error_policy model=%s account=%d reset_in=%v", + p.prefix, statusCode, modelName, p.account.ID, rateLimitDuration) + return true +} + // mapAntigravityModel 获取映射后的模型名 // 完全依赖映射配置：账户映射（通配符）→ 默认映射兜底（DefaultAntigravityModelMapping） // 注意：返回空字符串表示模型不被支持，调度时会过滤掉该账号 @@ -958,6 +994,7 @@ func mapAntigravityModel(account *Account, requestedModel string) string { if account == nil { return "" } + requestedModel = strings.TrimPrefix(requestedModel, "models/") // 获取映射表（未配置时自动使用 DefaultAntigravityModelMapping） mapping := account.GetModelMapping() @@ -1090,7 +1127,7 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account } defer func() { _ = result.resp.Body.Close() }() - respBody, err := io.ReadAll(io.LimitReader(result.resp.Body, 2<<20)) + respBody, err := io.ReadAll(io.LimitReader(result.resp.Body, s.upstreamErrorBodyReadLimit())) if err != nil { return nil, fmt.Errorf("读取响应失败: %w", err) } @@ -1427,7 +1464,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, defer func() { _ = resp.Body.Close() }() if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) // 优先检测 thinking block 的 signature 相关错误（400）并重试一次： // Antigravity /v1internal 链路在部分场景会对 thought/thinking signature 做严格校验， @@ -1622,7 +1659,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, resp = retryResp respBody = nil } else { - retryBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20)) + retryBody := s.readUpstreamErrorBody(retryResp) _ = retryResp.Body.Close() respBody = retryBody resp = &http.Response{ @@ -2057,8 +2094,28 @@ func stripSignatureSensitiveBlocksFromClaudeRequest(req *antigravity.ClaudeReque // └─ retryDelay < 7s → 等待后重试 1 次 // ├─ 成功 → 正常返回 // └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 -func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool) (*ForwardResult, error) { +type ForwardGeminiOption func(*forwardGeminiOptions) + +type forwardGeminiOptions struct { + groupID int64 + sessionHash string +} + +func WithForwardGeminiSession(groupID int64, sessionHash string) ForwardGeminiOption { + return func(opts *forwardGeminiOptions) { + opts.groupID = groupID + opts.sessionHash = sessionHash + } +} + +func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool, options ...ForwardGeminiOption) (*ForwardResult, error) { startTime := time.Now() + forwardOpts := forwardGeminiOptions{} + for _, apply := range options { + if apply != nil { + apply(&forwardOpts) + } + } sessionID := getSessionID(c) prefix := logPrefix(sessionID, account.Name) @@ -2163,8 +2220,8 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co handleError: s.handleUpstreamError, requestedModel: originalModel, isStickySession: isStickySession, // ForwardGemini 由上层判断粘性会话 - groupID: 0, // ForwardGemini 方法没有 groupID，由上层处理粘性会话清除 - sessionHash: "", // ForwardGemini 方法没有 sessionHash，由上层处理粘性会话清除 + groupID: forwardOpts.groupID, + sessionHash: forwardOpts.sessionHash, }) if err != nil { // 检查是否是账号切换信号，转换为 UpstreamFailoverError 让 Handler 切换账号 @@ -2189,7 +2246,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co // 处理错误响应 if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) contentType := resp.Header.Get("Content-Type") // 尽早关闭原始响应体，释放连接；后续逻辑仍可能需要读取 body，因此用内存副本重新包装。 _ = resp.Body.Close() @@ -2262,15 +2319,15 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co handleError: s.handleUpstreamError, requestedModel: originalModel, isStickySession: isStickySession, - groupID: 0, - sessionHash: "", + groupID: forwardOpts.groupID, + sessionHash: forwardOpts.sessionHash, }) if retryErr == nil { retryResp := retryResult.resp if retryResp.StatusCode < 400 { resp = retryResp } else { - retryRespBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20)) + retryRespBody := s.readUpstreamErrorBody(retryResp) _ = retryResp.Body.Close() retryOpsBody := retryRespBody if retryUnwrapped, unwrapErr := s.unwrapV1InternalResponse(retryRespBody); unwrapErr == nil && len(retryUnwrapped) > 0 { @@ -2339,7 +2396,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co if unwrapErr != nil || len(unwrappedForOps) == 0 { unwrappedForOps = respBody } - s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", isStickySession) + s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, forwardOpts.groupID, forwardOpts.sessionHash, isStickySession) upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(unwrappedForOps)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) upstreamDetail := s.getUpstreamErrorDetail(unwrappedForOps) @@ -2550,6 +2607,34 @@ func setModelRateLimitByModelName(ctx context.Context, repo AccountRepository, a return true } +func (s *AntigravityGatewayService) setAntigravityModelRateLimits(ctx context.Context, repo AccountRepository, account *Account, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { + if account == nil || repo == nil { + return false + } + keys := antigravityModelRateLimitKeys(modelName) + if len(keys) == 0 { + return false + } + + success := false + for _, key := range keys { + if setModelRateLimitByModelName(ctx, repo, account.ID, key, prefix, statusCode, resetAt, afterSmartRetry) { + s.updateAccountModelRateLimitInCache(ctx, account, key, resetAt) + success = true + } + } + return success +} + +func (s *AntigravityGatewayService) clearStickySession(ctx context.Context, groupID int64, sessionHash string) { + if s == nil || s.cache == nil || strings.TrimSpace(sessionHash) == "" { + return + } + if err := s.cache.DeleteSessionAccountID(ctx, groupID, sessionHash); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] sticky_session_clear_failed group_id=%d session=%s err=%v", groupID, shortSessionHash(sessionHash), err) + } +} + func antigravityFallbackCooldownSeconds() (time.Duration, bool) { raw := strings.TrimSpace(os.Getenv(antigravityFallbackSecondsEnv)) if raw == "" { @@ -2628,7 +2713,7 @@ func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo { if atType == googleRPCTypeErrorInfo { if meta, ok := dm["metadata"].(map[string]any); ok { if model, ok := meta["model"].(string); ok { - modelName = model + modelName = normalizeAntigravityModelName(model) } } // 检查 reason @@ -2802,13 +2887,7 @@ func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleMo logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", p.prefix, p.statusCode, info.ModelName, p.account.ID, info.RetryDelay) - // 设置模型限流状态（数据库） - if err := s.accountRepo.SetModelRateLimit(p.ctx, p.account.ID, info.ModelName, resetAt); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s model_rate_limit_failed model=%s error=%v", p.prefix, info.ModelName, err) - } - - // 立即更新 Redis 快照中账号的限流状态，避免并发请求重复选中 - s.updateAccountModelRateLimitInCache(p.ctx, p.account, info.ModelName, resetAt) + s.setAntigravityModelRateLimits(p.ctx, s.accountRepo, p.account, info.ModelName, p.prefix, p.statusCode, resetAt, false) // 清除粘性会话绑定 if p.cache != nil && p.sessionHash != "" { @@ -2898,12 +2977,11 @@ func (s *AntigravityGatewayService) handleUpstreamError( } if modelKey != "" { ra := s.resolveResetTime(resetAt, defaultDur) - if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, ra); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s error=%v", prefix, modelKey, err) + if !s.setAntigravityModelRateLimits(ctx, s.accountRepo, account, modelKey, prefix, statusCode, ra, false) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s", prefix, modelKey) } else { logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limited model=%s account=%d reset_at=%v reset_in=%v", prefix, modelKey, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) - s.updateAccountModelRateLimitInCache(ctx, account, modelKey, ra) } return nil } @@ -4252,7 +4330,7 @@ func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin. // 处理错误响应 if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) // 429 错误时标记账号限流 if resp.StatusCode == http.StatusTooManyRequests { diff --git a/backend/internal/service/antigravity_gateway_service_test.go b/backend/internal/service/antigravity_gateway_service_test.go index 2212437467cf..0fac7a1eb617 100644 --- a/backend/internal/service/antigravity_gateway_service_test.go +++ b/backend/internal/service/antigravity_gateway_service_test.go @@ -42,6 +42,15 @@ func newAntigravityTestService(cfg *config.Config) *AntigravityGatewayService { } } +func TestAntigravityUpstreamErrorBodyReadLimit_RespectsDiagnosticLimit(t *testing.T) { + svc := newAntigravityTestService(&config.Config{Gateway: config.GatewayConfig{ + LogUpstreamErrorBody: true, + LogUpstreamErrorBodyMaxBytes: int(gatewayUpstreamErrorBodyReadLimit) + 1024, + }}) + + require.Equal(t, int64(svc.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), svc.upstreamErrorBodyReadLimit()) +} + func TestStripSignatureSensitiveBlocksFromClaudeRequest(t *testing.T) { req := &antigravity.ClaudeRequest{ Model: "claude-sonnet-4-5", @@ -491,6 +500,86 @@ func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling( require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch") } +func TestAntigravityGatewayService_ForwardGemini_ClearsStickySessionOnGeminiRateLimit(t *testing.T) { + gin.SetMode(gin.TestMode) + writer := httptest.NewRecorder() + c, _ := gin.CreateTestContext(writer) + + body, err := json.Marshal(map[string]any{ + "contents": []map[string]any{ + {"role": "user", "parts": []map[string]any{{"text": "hi"}}}, + }, + }) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3-flash-preview:generateContent", bytes.NewReader(body)) + c.Request = req + + respBody := []byte(`{ + "error": { + "status": "RESOURCE_EXHAUSTED", + "details": [ + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"}, + {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"} + ] + } + }`) + upstream := &httpUpstreamStub{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{}, + Body: io.NopCloser(bytes.NewReader(respBody)), + }} + repo := &stubAntigravityAccountRepo{} + cache := &stubSmartRetryCache{} + svc := &AntigravityGatewayService{ + tokenProvider: &AntigravityTokenProvider{}, + httpUpstream: upstream, + accountRepo: repo, + cache: cache, + } + + account := &Account{ + ID: 44, + Name: "acc-gemini-runtime-rate-limited", + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "token", + "expires_at": time.Now().Add(time.Hour).Format(time.RFC3339), + "project_id": "proj", + }, + Extra: map[string]any{ + "mixed_scheduling": true, + }, + } + + result, err := svc.ForwardGemini( + context.Background(), + c, + account, + "gemini-3-flash-preview", + "generateContent", + false, + body, + true, + WithForwardGeminiSession(77, "gemini:sticky-runtime"), + ) + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusServiceUnavailable, failoverErr.StatusCode) + require.Len(t, repo.modelRateLimitCalls, 2) + require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) + require.Len(t, cache.deleteCalls, 1) + require.Equal(t, int64(77), cache.deleteCalls[0].groupID) + require.Equal(t, "gemini:sticky-runtime", cache.deleteCalls[0].sessionHash) +} + // TestAntigravityGatewayService_Forward_BillsWithMappedModel // 验证：Antigravity Claude 转发返回的计费模型使用映射后的模型 func TestAntigravityGatewayService_Forward_BillsWithMappedModel(t *testing.T) { diff --git a/backend/internal/service/antigravity_quota_scope.go b/backend/internal/service/antigravity_quota_scope.go index b536d16cf893..758626331b33 100644 --- a/backend/internal/service/antigravity_quota_scope.go +++ b/backend/internal/service/antigravity_quota_scope.go @@ -8,7 +8,17 @@ import ( func normalizeAntigravityModelName(model string) string { normalized := strings.ToLower(strings.TrimSpace(model)) - normalized = strings.TrimPrefix(normalized, "models/") + if idx := strings.LastIndex(normalized, "/publishers/google/models/"); idx != -1 { + normalized = normalized[idx+len("/publishers/google/models/"):] + } else if idx := strings.LastIndex(normalized, "/publishers/anthropic/models/"); idx != -1 { + normalized = normalized[idx+len("/publishers/anthropic/models/"):] + } else if idx := strings.LastIndex(normalized, "/models/"); idx != -1 { + normalized = normalized[idx+len("/models/"):] + } else { + normalized = strings.TrimPrefix(normalized, "publishers/google/models/") + normalized = strings.TrimPrefix(normalized, "publishers/anthropic/models/") + normalized = strings.TrimPrefix(normalized, "models/") + } return normalized } diff --git a/backend/internal/service/antigravity_rate_limit_test.go b/backend/internal/service/antigravity_rate_limit_test.go index c3e49458d876..374b29f610cc 100644 --- a/backend/internal/service/antigravity_rate_limit_test.go +++ b/backend/internal/service/antigravity_rate_limit_test.go @@ -821,6 +821,51 @@ func TestSetModelRateLimitByModelName_NotConvertToScope(t *testing.T) { require.NotEqual(t, "claude_sonnet", call.modelKey, "should NOT be scope") } +func TestSetAntigravityModelRateLimits_GeminiWritesFamilyScope(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + svc := &AntigravityGatewayService{} + account := &Account{ID: 789, Platform: PlatformAntigravity} + resetAt := time.Now().Add(30 * time.Second) + + success := svc.setAntigravityModelRateLimits( + context.Background(), + repo, + account, + "gemini-3-pro", + "[test]", + 429, + resetAt, + false, + ) + + require.True(t, success) + require.Len(t, repo.modelRateLimitCalls, 2) + require.Equal(t, "gemini-3-pro", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) +} + +func TestSetAntigravityModelRateLimits_ClaudeDoesNotWriteGeminiScope(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + svc := &AntigravityGatewayService{} + account := &Account{ID: 790, Platform: PlatformAntigravity} + resetAt := time.Now().Add(30 * time.Second) + + success := svc.setAntigravityModelRateLimits( + context.Background(), + repo, + account, + "claude-sonnet-4-5", + "[test]", + 429, + resetAt, + false, + ) + + require.True(t, success) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey) +} + func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRateLimited(t *testing.T) { upstream := &recordingOKUpstream{} account := &Account{ @@ -1124,3 +1169,53 @@ func TestSchedulerSnapshotService_UpdateAccountInCache(t *testing.T) { require.ErrorIs(t, err, expectedErr) }) } +func TestNormalizeAntigravityModelName(t *testing.T) { + tests := []struct { + name string + model string + expected string + }{ + { + name: "plain model name", + model: "gemini-1.5-pro", + expected: "gemini-1.5-pro", + }, + { + name: "models/ prefix", + model: "models/gemini-1.5-pro", + expected: "gemini-1.5-pro", + }, + { + name: "publishers/google/models/ prefix", + model: "publishers/google/models/gemini-1.5-pro", + expected: "gemini-1.5-pro", + }, + { + name: "projects/.../publishers/google/models/ path", + model: "projects/my-proj/locations/us-central1/publishers/google/models/gemini-2.5-flash", + expected: "gemini-2.5-flash", + }, + { + name: "publishers/anthropic/models/ prefix", + model: "publishers/anthropic/models/claude-sonnet-4-5", + expected: "claude-sonnet-4-5", + }, + { + name: "projects/.../publishers/anthropic/models/ path", + model: "projects/my-proj/locations/global/publishers/anthropic/models/claude-sonnet-4-5", + expected: "claude-sonnet-4-5", + }, + { + name: "mixed case and spaces", + model: " Models/Gemini-1.5-Pro ", + expected: "gemini-1.5-pro", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + actual := normalizeAntigravityModelName(tt.model) + require.Equal(t, tt.expected, actual) + }) + } +} diff --git a/backend/internal/service/antigravity_single_account_retry_test.go b/backend/internal/service/antigravity_single_account_retry_test.go index 675e9c0cbf4d..6e58ab75d2cb 100644 --- a/backend/internal/service/antigravity_single_account_retry_test.go +++ b/backend/internal/service/antigravity_single_account_retry_test.go @@ -196,8 +196,10 @@ func TestHandleSmartRetry_503_LongDelay_NoSingleAccountRetry_StillSwitches(t *te require.Nil(t, result.resp, "should not return resp when switchError is set") // 对照：多账号模式应设模型限流 - require.Len(t, repo.modelRateLimitCalls, 1, + require.Len(t, repo.modelRateLimitCalls, 2, "multi-account mode SHOULD set model rate limit") + require.Equal(t, "gemini-3-pro-high", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) } // TestHandleSmartRetry_429_LongDelay_SingleAccountRetry_StillSwitches @@ -412,8 +414,10 @@ func TestHandleSmartRetry_503_ShortDelay_NoSingleAccountRetry_SetsRateLimit(t *t // 对照：多账号模式应返回 switchError require.NotNil(t, result.switchError, "multi-account mode should return switchError for 503") // 对照：多账号模式应设模型限流 - require.Len(t, repo.modelRateLimitCalls, 1, + require.Len(t, repo.modelRateLimitCalls, 2, "multi-account mode should set model rate limit") + require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) } // --------------------------------------------------------------------------- diff --git a/backend/internal/service/antigravity_smart_retry_test.go b/backend/internal/service/antigravity_smart_retry_test.go index e3b60a279027..9f06b13b0632 100644 --- a/backend/internal/service/antigravity_smart_retry_test.go +++ b/backend/internal/service/antigravity_smart_retry_test.go @@ -328,9 +328,10 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test require.Equal(t, "gemini-3-flash", result.switchError.RateLimitedModel) require.False(t, result.switchError.IsStickySession) - // 验证模型限流已设置 - require.Len(t, repo.modelRateLimitCalls, 1) + // 验证模型限流已设置：Gemini 同时写入精确模型和家族级 scope + require.Len(t, repo.modelRateLimitCalls, 2) require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) require.Len(t, upstream.calls, 1, "should have made one retry call (max attempts)") } @@ -1104,10 +1105,9 @@ func TestHandleSmartRetry_ShortDelay_StickySession_SuccessRetry_NoDeleteSession( require.Len(t, cache.deleteCalls, 0, "should NOT call DeleteSessionAccountID on successful retry") } -// TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry -// 长延迟路径（情况1）在 handleSmartRetry 中不直接调用 DeleteSessionAccountID -// （清除由 handler 层的 shouldClearStickySession 在下次请求时处理） -func TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry(t *testing.T) { +// TestHandleSmartRetry_LongDelay_StickySession_ClearsSession +// 长延迟路径（情况1）应立即清除 sticky 绑定，避免下一次请求继续命中已限流账号。 +func TestHandleSmartRetry_LongDelay_StickySession_ClearsSession(t *testing.T) { repo := &stubAntigravityAccountRepo{} cache := &stubSmartRetryCache{} account := &Account{ @@ -1159,10 +1159,9 @@ func TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry(t * require.NotNil(t, result.switchError) require.True(t, result.switchError.IsStickySession) - // 长延迟路径不在 handleSmartRetry 中调用 DeleteSessionAccountID - // （由上游 handler 的 shouldClearStickySession 处理） - require.Len(t, cache.deleteCalls, 0, - "long delay path should NOT call DeleteSessionAccountID in handleSmartRetry (handled by handler layer)") + require.Len(t, cache.deleteCalls, 1, "long delay path should clear sticky session in handleSmartRetry") + require.Equal(t, int64(42), cache.deleteCalls[0].groupID) + require.Equal(t, "sticky-hash-long-delay", cache.deleteCalls[0].sessionHash) } // TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession @@ -1227,6 +1226,10 @@ func TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession(t require.Len(t, cache.deleteCalls, 1, "should call DeleteSessionAccountID after network error exhausts retry") require.Equal(t, int64(99), cache.deleteCalls[0].groupID) require.Equal(t, "sticky-net-error", cache.deleteCalls[0].sessionHash) + + require.Len(t, repo.modelRateLimitCalls, 2) + require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) } // TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession @@ -1308,9 +1311,10 @@ func TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession require.Equal(t, int64(77), cache.deleteCalls[0].groupID) require.Equal(t, "sticky-503-short", cache.deleteCalls[0].sessionHash) - // 验证模型限流已设置 - require.Len(t, repo.modelRateLimitCalls, 1) + // 验证模型限流已设置：Gemini 同时写入精确模型和家族级 scope + require.Len(t, repo.modelRateLimitCalls, 2) require.Equal(t, "gemini-3-pro", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) } // TestAntigravityRetryLoop_SmartRetryFailed_StickySession_SwitchErrorPropagates diff --git a/backend/internal/service/auth_service_email_bind_test.go b/backend/internal/service/auth_service_email_bind_test.go index 87867395a809..28bb0a3b8ebc 100644 --- a/backend/internal/service/auth_service_email_bind_test.go +++ b/backend/internal/service/auth_service_email_bind_test.go @@ -850,6 +850,9 @@ func (s *emailBindUserRepoStub) UnbindUserAuthProvider(context.Context, int64, s func (s *emailBindUserRepoStub) UpdateTotpSecret(context.Context, int64, *string) error { return nil } func (s *emailBindUserRepoStub) EnableTotp(context.Context, int64) error { return nil } func (s *emailBindUserRepoStub) DisableTotp(context.Context, int64) error { return nil } +func (s *emailBindUserRepoStub) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) { + return s.GetByID(ctx, id) +} func cloneEmailBindUser(user *service.User) *service.User { if user == nil { diff --git a/backend/internal/service/auth_service_platform_quota_test.go b/backend/internal/service/auth_service_platform_quota_test.go index f58dc48c35a2..46069814b04c 100644 --- a/backend/internal/service/auth_service_platform_quota_test.go +++ b/backend/internal/service/auth_service_platform_quota_test.go @@ -43,6 +43,10 @@ func (f *fakeInsertRecorder) ResetExpiredWindow(_ context.Context, _ int64, _ st return nil } +func (f *fakeInsertRecorder) BatchSnapshotUsage(_ context.Context, _ []UserPlatformQuotaSnapshot, _ time.Time) error { + return nil +} + func TestSnapshotPlatformQuotaDefaults_PassesToRepoBulkInsert(t *testing.T) { fakeRepo := &fakeInsertRecorder{} s := &AuthService{userPlatformQuotaRepo: fakeRepo} diff --git a/backend/internal/service/auth_service_register_test.go b/backend/internal/service/auth_service_register_test.go index a7c0d260c872..2ee9f21a562c 100644 --- a/backend/internal/service/auth_service_register_test.go +++ b/backend/internal/service/auth_service_register_test.go @@ -105,6 +105,10 @@ func (s *userPlatformQuotaRepoStub) ResetExpiredWindow(context.Context, int64, s panic("unexpected ResetExpiredWindow call") } +func (s *userPlatformQuotaRepoStub) BatchSnapshotUsage(_ context.Context, _ []UserPlatformQuotaSnapshot, _ time.Time) error { + return nil +} + func (s *defaultSubscriptionAssignerStub) AssignOrExtendSubscription(_ context.Context, input *AssignSubscriptionInput) (*UserSubscription, bool, error) { if input != nil { s.calls = append(s.calls, *input) diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index 2b7c06baa287..b734fab13ee8 100644 --- a/backend/internal/service/billing_cache_service.go +++ b/backend/internal/service/billing_cache_service.go @@ -689,7 +689,8 @@ func (s *BillingCacheService) IncrementUserPlatformQuotaUsage(userID int64, plat ctx, cancel := context.WithTimeout(context.Background(), cacheWriteTimeout) defer cancel() ttl := time.Duration(s.cfg.Billing.UserPlatformQuotaCacheTTLSeconds) * time.Second - if err := s.cache.IncrUserPlatformQuotaUsageCache(ctx, userID, platform, cost, ttl); err != nil { + markDirty := s.cfg.Database.UserPlatformQuotaFlusherEnabled + if err := s.cache.IncrUserPlatformQuotaUsageCache(ctx, userID, platform, cost, ttl, markDirty); err != nil { logger.LegacyPrintf("service.billing_cache", "ALERT: incr user platform quota cache failed user=%d platform=%s cost=%f: %v", userID, platform, cost, err) @@ -1096,7 +1097,12 @@ func (s *BillingCacheService) checkUserPlatformQuotaEligibility( // 超时 50ms:覆盖正常路径与可接受抖动;Redis 异常时 hot path 不阻塞超过此值。 // 用 context.Background()+短超时,避免请求 ctx 取消导致刷新丢失。 // 显式 setCancel()(而非 defer):缩短 context 生命周期,避免 defer 延迟到函数返回。 - if windowExpired && s.cache != nil { + // isSentinel 判定「该 entry 无任何 limit」,涵盖两类,跨窗口命中时都跳过 refresh: + // 1) A3 回填的 sentinel(DB 无行,短 TTL):refresh 会把短 TTL 误升级为 86400s,有害; + // 2) DB 有行但三 limit 全未配置的用户(TTL 86400s):refresh 纯属无意义(TTL 升级本身无害)。 + // 两类的 enforcement(下方 limit!=nil 比较)都因 limit 全 nil 永远放行,跳过 refresh 均正确。 + isSentinel := entry.DailyLimitUSD == nil && entry.WeeklyLimitUSD == nil && entry.MonthlyLimitUSD == nil + if windowExpired && s.cache != nil && !isSentinel { refreshed := &UserPlatformQuotaCacheEntry{ DailyUsageUSD: dailyUsage, WeeklyUsageUSD: weeklyUsage, @@ -1159,6 +1165,33 @@ func (s *BillingCacheService) checkUserPlatformQuotaEligibility( } rec, _ := v.(*UserPlatformQuotaRecord) if rec == nil { + // 仅在 cache 可用且本次 GET 未出错时回填 sentinel:Redis GET 故障(cacheErr!=nil) + // 时不回填,与下方 line ~1201 "Redis 故障时 fail-open:不回填" 保持一致, + // 避免在 Redis 异常期做一次注定失败的 SET。 + if s.cache != nil && cacheErr == nil { + now := time.Now() + startOfDay := timezone.StartOfDay(now) + startOfWeek := timezone.StartOfWeek(now) + sentinel := &UserPlatformQuotaCacheEntry{ + SchemaVersion: UserPlatformQuotaCacheSchemaV1, + DailyWindowStart: &startOfDay, + WeeklyWindowStart: &startOfWeek, + MonthlyWindowStart: &now, + // limits 全 nil, usage 全 0(零值) + } + sentinelTTL := time.Duration(s.cfg.Billing.UserPlatformQuotaSentinelTTLSeconds) * time.Second + if sentinelTTL <= 0 { + // 防御:TTL<=0 时 Redis EXPIRE 会立即删除整个 key(见 billing_cache.go 的 pipe.Expire), + // sentinel 不持久化 → 每请求击穿 DB。配置缺失/误配为 0 时 fallback 到 1h。 + sentinelTTL = time.Hour + } + setCtx, setCancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + if setErr := s.cache.SetUserPlatformQuotaCache(setCtx, userID, platform, sentinel, sentinelTTL); setErr != nil { + userPlatformQuotaSentinelSetCacheErrorTotal.Add(1) + logger.LegacyPrintf("service.billing_cache", "Warning: set sentinel quota cache failed user=%d platform=%s: %v", userID, platform, setErr) + } + setCancel() + } return nil } @@ -1278,3 +1311,20 @@ func monthlyQuotaWindowExpired(start *time.Time, now time.Time) bool { } return now.Sub(*start) >= 30*24*time.Hour } + +// HasUserPlatformQuotaLimit 判断该 user×platform 是否设了任一非 nil limit。 +// 写入点守卫:无 limit 直接跳过 Redis 写 + 脏集标记,消除无谓写入。 +// fail-safe:任何不确定(simple 模式除外)都返回 true 维持写入。 +func (s *BillingCacheService) HasUserPlatformQuotaLimit(ctx context.Context, userID int64, platform string) bool { + if s.cfg.RunMode == config.RunModeSimple { + return false + } + if s.cache == nil { + return true + } + entry, ok, err := s.cache.GetUserPlatformQuotaCache(ctx, userID, platform) + if err != nil || !ok || entry == nil { + return true + } + return entry.DailyLimitUSD != nil || entry.WeeklyLimitUSD != nil || entry.MonthlyLimitUSD != nil +} diff --git a/backend/internal/service/billing_cache_service_singleflight_test.go b/backend/internal/service/billing_cache_service_singleflight_test.go index b443d97eaa0b..235b13a6c86c 100644 --- a/backend/internal/service/billing_cache_service_singleflight_test.go +++ b/backend/internal/service/billing_cache_service_singleflight_test.go @@ -79,10 +79,22 @@ func (s *billingCacheMissStub) DeleteUserPlatformQuotaCache(ctx context.Context, return nil } -func (s *billingCacheMissStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { +func (s *billingCacheMissStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { return nil } +func (s *billingCacheMissStub) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) { + return nil, nil +} + +func (s *billingCacheMissStub) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error { + return nil +} + +func (s *billingCacheMissStub) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + return nil, nil +} + type balanceLoadUserRepoStub struct { mockUserRepo calls atomic.Int64 diff --git a/backend/internal/service/billing_cache_service_test.go b/backend/internal/service/billing_cache_service_test.go index bcd086fa5c1c..c344b4173849 100644 --- a/backend/internal/service/billing_cache_service_test.go +++ b/backend/internal/service/billing_cache_service_test.go @@ -80,10 +80,22 @@ func (b *billingCacheWorkerStub) DeleteUserPlatformQuotaCache(ctx context.Contex return nil } -func (b *billingCacheWorkerStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { +func (b *billingCacheWorkerStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { return nil } +func (b *billingCacheWorkerStub) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) { + return nil, nil +} + +func (b *billingCacheWorkerStub) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error { + return nil +} + +func (b *billingCacheWorkerStub) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + return nil, nil +} + func TestBillingCacheServiceQueueHighLoad(t *testing.T) { cache := &billingCacheWorkerStub{} svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil) diff --git a/backend/internal/service/billing_cache_service_user_platform_quota_test.go b/backend/internal/service/billing_cache_service_user_platform_quota_test.go index 57697ddb05b8..a82c0e00d75c 100644 --- a/backend/internal/service/billing_cache_service_user_platform_quota_test.go +++ b/backend/internal/service/billing_cache_service_user_platform_quota_test.go @@ -20,14 +20,15 @@ type fakeIncrCache struct { } type incrCall struct { - userID int64 - platform string - cost float64 - ttl time.Duration + userID int64 + platform string + cost float64 + ttl time.Duration + markDirty bool } -func (f *fakeIncrCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error { - f.calls = append(f.calls, incrCall{userID, platform, cost, ttl}) +func (f *fakeIncrCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error { + f.calls = append(f.calls, incrCall{userID, platform, cost, ttl, markDirty}) return nil } @@ -49,10 +50,10 @@ func TestIncrementUserPlatformQuotaUsage_SyncCallsCache(t *testing.T) { if len(fake.calls) != 2 { t.Fatalf("expected 2 incr calls, got %d", len(fake.calls)) } - if fake.calls[0] != (incrCall{101, "anthropic", 0.25, 120 * time.Second}) { + if fake.calls[0] != (incrCall{userID: 101, platform: "anthropic", cost: 0.25, ttl: 120 * time.Second, markDirty: false}) { t.Errorf("call[0] = %+v", fake.calls[0]) } - if fake.calls[1] != (incrCall{101, "openai", 0.50, 120 * time.Second}) { + if fake.calls[1] != (incrCall{userID: 101, platform: "openai", cost: 0.50, ttl: 120 * time.Second, markDirty: false}) { t.Errorf("call[1] = %+v", fake.calls[1]) } } @@ -88,13 +89,23 @@ func (f *fakeQuotaRepo) ResetExpiredWindow(_ context.Context, _ int64, _ string, return nil } -// fakeFullCache 同时支持 Get + Set + Incr + Delete。 +func (f *fakeQuotaRepo) BatchSnapshotUsage(_ context.Context, _ []UserPlatformQuotaSnapshot, _ time.Time) error { + return nil +} + +// fakeFullCache 同时支持 Get + Set + Incr + Delete + Pop/Readd/BatchGet（脏集读写）。 // mu 保护 entry 和 deleteCalls，防止异步 goroutine 与主 goroutine 之间的 data race。 type fakeFullCache struct { BillingCache mu sync.Mutex entry *UserPlatformQuotaCacheEntry deleteCalls int + setCalls int // SetUserPlatformQuotaCache 调用次数 + lastSetTTL time.Duration // 最近一次 Set 的 ttl + getErr error // 非 nil 时 Get 先返回 (nil,false,getErr) + setErr error // 非 nil 时 Set 返回该 err(setCalls 仍+1) + // dirty 模拟脏集，供 flusher 测试使用。 + dirty map[UserPlatformQuotaKey]struct{} } // getDeleteCalls 线程安全地读取 deleteCalls。 @@ -111,19 +122,41 @@ func (f *fakeFullCache) getEntry() *UserPlatformQuotaCacheEntry { return f.entry } +// getSetCalls 线程安全地读取 setCalls。 +func (f *fakeFullCache) getSetCalls() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.setCalls +} + +// getLastSetTTL 线程安全地读取 lastSetTTL。 +func (f *fakeFullCache) getLastSetTTL() time.Duration { + f.mu.Lock() + defer f.mu.Unlock() + return f.lastSetTTL +} + func (f *fakeFullCache) GetUserPlatformQuotaCache(_ context.Context, _ int64, _ string) (*UserPlatformQuotaCacheEntry, bool, error) { f.mu.Lock() defer f.mu.Unlock() + if f.getErr != nil { + return nil, false, f.getErr + } if f.entry == nil { return nil, false, nil } return f.entry, true, nil } -func (f *fakeFullCache) SetUserPlatformQuotaCache(_ context.Context, _ int64, _ string, e *UserPlatformQuotaCacheEntry, _ time.Duration) error { +func (f *fakeFullCache) SetUserPlatformQuotaCache(_ context.Context, _ int64, _ string, e *UserPlatformQuotaCacheEntry, ttl time.Duration) error { f.mu.Lock() defer f.mu.Unlock() + f.setCalls++ + if f.setErr != nil { + return f.setErr + } f.entry = e + f.lastSetTTL = ttl return nil } @@ -135,6 +168,48 @@ func (f *fakeFullCache) DeleteUserPlatformQuotaCache(_ context.Context, _ int64, return nil } +func (f *fakeFullCache) PopDirtyUserPlatformQuotaKeys(_ context.Context, n int) ([]UserPlatformQuotaKey, error) { + f.mu.Lock() + defer f.mu.Unlock() + if len(f.dirty) == 0 { + return nil, nil + } + keys := make([]UserPlatformQuotaKey, 0, n) + for k := range f.dirty { + if len(keys) >= n { + break + } + keys = append(keys, k) + delete(f.dirty, k) + } + return keys, nil +} + +func (f *fakeFullCache) ReaddDirtyUserPlatformQuotaKeys(_ context.Context, keys []UserPlatformQuotaKey) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.dirty == nil { + f.dirty = make(map[UserPlatformQuotaKey]struct{}) + } + for _, k := range keys { + f.dirty[k] = struct{}{} + } + return nil +} + +// BatchGetUserPlatformQuotaCache 对每个 key 返回 f.entry（MISS → nil）， +// 保持与输入 keys 顺序/长度对齐。注意此处所有 key 共享同一个 entry， +// 仅用于测试场景。 +func (f *fakeFullCache) BatchGetUserPlatformQuotaCache(_ context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + f.mu.Lock() + defer f.mu.Unlock() + results := make([]*UserPlatformQuotaCacheEntry, len(keys)) + for i := range keys { + results[i] = f.entry + } + return results, nil +} + func newServiceForPreflight(t *testing.T, repo UserPlatformQuotaRepository, cache BillingCache) *BillingCacheService { t.Helper() cfg := &config.Config{} @@ -593,3 +668,165 @@ func TestMonthlyQuotaWindowExpired_BoundaryTable(t *testing.T) { }) } } + +// TestCheckUserPlatformQuotaEligibility_NoRow_WritesSentinel 验证: +// cache MISS + DB 无行时,回填 sentinel entry(三 limit 全 nil,三 window_start 全 non-nil), +// TTL = UserPlatformQuotaSentinelTTLSeconds,函数返回 nil(fail-open)。 +func TestCheckUserPlatformQuotaEligibility_NoRow_WritesSentinel(t *testing.T) { + repo := &fakeQuotaRepo{rec: nil} // DB 无行 + cache := &fakeFullCache{} // entry=nil → Get 返回 MISS + svc := newServiceForPreflight(t, repo, cache) + svc.cfg.Billing.UserPlatformQuotaSentinelTTLSeconds = 3600 + + if err := svc.checkUserPlatformQuotaEligibility(context.Background(), 1, "anthropic"); err != nil { + t.Fatalf("expected nil (fail-open), got %v", err) + } + if cache.getSetCalls() != 1 { + t.Fatalf("expected 1 SetUserPlatformQuotaCache call for sentinel, got %d", cache.getSetCalls()) + } + sentinel := cache.getEntry() + if sentinel == nil { + t.Fatal("expected sentinel entry backfilled") + } + if sentinel.DailyLimitUSD != nil || sentinel.WeeklyLimitUSD != nil || sentinel.MonthlyLimitUSD != nil { + t.Errorf("sentinel must have all-nil limits") + } + if sentinel.DailyWindowStart == nil || sentinel.WeeklyWindowStart == nil || sentinel.MonthlyWindowStart == nil { + t.Errorf("sentinel must have non-nil window_start to avoid refresh churn") + } + if sentinel.SchemaVersion != UserPlatformQuotaCacheSchemaV1 { + t.Errorf("sentinel schema = %d, want V1", sentinel.SchemaVersion) + } + if cache.getLastSetTTL() != 3600*time.Second { + t.Errorf("sentinel ttl = %v, want 3600s", cache.getLastSetTTL()) + } +} + +// TestCheckUserPlatformQuotaEligibility_RedisGetError_NoSentinelBackfill 验证: +// Redis GET 故障(cacheErr!=nil)+ DB 无行时,不应回填 sentinel(与 "Redis 故障时不回填" 一致),且 fail-open。 +func TestCheckUserPlatformQuotaEligibility_RedisGetError_NoSentinelBackfill(t *testing.T) { + repo := &fakeQuotaRepo{rec: nil} + cache := &fakeFullCache{getErr: errors.New("redis get down")} + svc := newServiceForPreflight(t, repo, cache) + svc.cfg.Billing.UserPlatformQuotaSentinelTTLSeconds = 3600 + + if err := svc.checkUserPlatformQuotaEligibility(context.Background(), 1, "anthropic"); err != nil { + t.Fatalf("redis 故障应 fail-open, got %v", err) + } + if cache.getSetCalls() != 0 { + t.Errorf("redis-get-error 时不应回填 sentinel, got %d set calls", cache.getSetCalls()) + } +} + +// TestCheckUserPlatformQuotaEligibility_NoRow_SentinelSetFailsFailOpen 验证: +// sentinel SET 失败时 fail-open(返回 nil)且计 metric。 +func TestCheckUserPlatformQuotaEligibility_NoRow_SentinelSetFailsFailOpen(t *testing.T) { + before := userPlatformQuotaSentinelSetCacheErrorTotal.Load() + repo := &fakeQuotaRepo{rec: nil} + cache := &fakeFullCache{setErr: errors.New("redis set timeout")} + svc := newServiceForPreflight(t, repo, cache) + svc.cfg.Billing.UserPlatformQuotaSentinelTTLSeconds = 3600 + + if err := svc.checkUserPlatformQuotaEligibility(context.Background(), 1, "anthropic"); err != nil { + t.Fatalf("sentinel set 失败应 fail-open, got %v", err) + } + if cache.getSetCalls() != 1 { + t.Errorf("应尝试 set sentinel 恰好一次, got %d", cache.getSetCalls()) + } + if got := userPlatformQuotaSentinelSetCacheErrorTotal.Load() - before; got != 1 { + t.Errorf("set 失败应使 metric +1, got delta %d", got) + } +} + +// TestCheckUserPlatformQuotaEligibility_SentinelCrossDay_NoRefresh 验证: +// 命中 sentinel(三 limit 全 nil)且跨窗口(daily/weekly 过期)时,不应触发 refresh SetCache +// (否则会把短 sentinel TTL 误升级为 quota cache 默认 86400s)。 +func TestCheckUserPlatformQuotaEligibility_SentinelCrossDay_NoRefresh(t *testing.T) { + yesterday := timezone.StartOfDay(time.Now().AddDate(0, 0, -1)) + lastWeek := timezone.StartOfWeek(time.Now().AddDate(0, 0, -7)) + monthAgoOK := time.Now().AddDate(0, 0, -5) // <30d, monthly 不过期 + sentinel := &UserPlatformQuotaCacheEntry{ + SchemaVersion: UserPlatformQuotaCacheSchemaV1, + DailyWindowStart: &yesterday, // 跨日 → daily windowExpired = true + WeeklyWindowStart: &lastWeek, // 跨周 → weekly windowExpired = true + MonthlyWindowStart: &monthAgoOK, + // limits 全 nil → sentinel + } + cache := &fakeFullCache{entry: sentinel} // entry 非 nil → Get HIT + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, cache) + + if err := svc.checkUserPlatformQuotaEligibility(context.Background(), 1, "anthropic"); err != nil { + t.Fatalf("sentinel = no limit, expected nil, got %v", err) + } + if cache.getSetCalls() != 0 { + t.Errorf("sentinel cross-window must NOT trigger refresh SetCache, got %d calls", cache.getSetCalls()) + } +} + +// ── TestHasUserPlatformQuotaLimit ──────────────────────────────────────────── + +func TestHasUserPlatformQuotaLimit(t *testing.T) { + daily := 5.0 + + tests := []struct { + name string + setup func() *BillingCacheService + want bool + }{ + { + name: "has_limit", + setup: func() *BillingCacheService { + entry := &UserPlatformQuotaCacheEntry{DailyLimitUSD: &daily} + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{entry: entry}) + return svc + }, + want: true, + }, + { + name: "sentinel_no_limit", + setup: func() *BillingCacheService { + entry := &UserPlatformQuotaCacheEntry{} // 三个 limit 字段全 nil + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{entry: entry}) + return svc + }, + want: false, + }, + { + name: "cache_miss", + setup: func() *BillingCacheService { + // entry==nil → GetUserPlatformQuotaCache 返回 (nil,false,nil) + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{}) + return svc + }, + want: true, // fail-safe + }, + { + name: "redis_err", + setup: func() *BillingCacheService { + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{getErr: errors.New("redis down")}) + return svc + }, + want: true, // fail-safe + }, + { + name: "simple_mode", + setup: func() *BillingCacheService { + entry := &UserPlatformQuotaCacheEntry{DailyLimitUSD: &daily} + svc := newServiceForPreflight(t, &fakeQuotaRepo{}, &fakeFullCache{entry: entry}) + svc.cfg.RunMode = config.RunModeSimple + return svc + }, + want: false, // simple 模式始终跳过 + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + svc := tt.setup() + got := svc.HasUserPlatformQuotaLimit(context.Background(), 1, "anthropic") + if got != tt.want { + t.Errorf("HasUserPlatformQuotaLimit() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 7d3fa6ef8d11..98551ecd1717 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -21,6 +21,12 @@ type APIKeyRateLimitCacheData struct { Window7d int64 `json:"window_7d"` } +// UserPlatformQuotaKey 标识一个 user×platform，用于脏集出入与批量读。 +type UserPlatformQuotaKey struct { + UserID int64 + Platform string +} + // UserPlatformQuotaCacheEntry Redis hash 反序列化结果。 // // SchemaVersion 用于向后兼容： @@ -72,7 +78,13 @@ type BillingCache interface { SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *UserPlatformQuotaCacheEntry, ttl time.Duration) error DeleteUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) error // IncrUserPlatformQuotaUsageCache 在缓存命中时累加用量；缓存未命中（key 不存在）静默返回 nil。 - IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration) error + // markDirty=true 时将该 key 的 member 写入 Redis 脏集，供 flusher 批量回写 DB。 + IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error + + // 脏集读写，供 flusher 使用。 + PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) + ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error + BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) } // ModelPricing 模型价格配置（per-token价格，与LiteLLM格式一致） diff --git a/backend/internal/service/claude_code_validator.go b/backend/internal/service/claude_code_validator.go index 4e8ced67954b..2c5ded6f9b73 100644 --- a/backend/internal/service/claude_code_validator.go +++ b/backend/internal/service/claude_code_validator.go @@ -56,7 +56,7 @@ func NewClaudeCodeValidator() *ClaudeCodeValidator { // 采用与 claude-relay-service 完全一致的验证策略： // // Step 1: User-Agent 检查 (必需) - 必须是 claude-cli/x.x.x -// Step 2: 对于非 messages 路径，只要 UA 匹配就通过 +// Step 2: 对于非 messages 路径和 /messages/count_tokens，只要 UA 匹配就通过 // Step 3: 检查 max_tokens=1 + haiku 探测请求绕过（UA 已验证） // Step 4: 对于 messages 路径，进行严格验证： // - System prompt 相似度检查 @@ -71,12 +71,17 @@ func (v *ClaudeCodeValidator) Validate(r *http.Request, body map[string]any) boo return false } - // Step 2: 非 messages 路径，只要 UA 匹配就通过 + // Step 2: 非 messages 路径只要 UA 匹配就通过 path := r.URL.Path if !strings.Contains(path, "messages") { return true } + // count_tokens 是 Claude Code 官方辅助请求，通常不携带完整 messages system prompt。 + if isMessagesCountTokensPath(path) { + return true + } + // Step 3: 检查 max_tokens=1 + haiku 探测请求绕过 // 这类请求用于 Claude Code 验证 API 连通性，不携带 system prompt if isMaxTokensOneHaiku, ok := IsMaxTokensOneHaikuRequestFromContext(r.Context()); ok && isMaxTokensOneHaiku { @@ -128,6 +133,10 @@ func (v *ClaudeCodeValidator) Validate(r *http.Request, body map[string]any) boo return true } +func isMessagesCountTokensPath(path string) bool { + return strings.HasSuffix(path, "/messages/count_tokens") +} + // hasClaudeCodeSystemPrompt 检查请求是否包含 Claude Code 系统提示词 // 使用字符串相似度匹配（Dice coefficient） func (v *ClaudeCodeValidator) hasClaudeCodeSystemPrompt(body map[string]any) bool { diff --git a/backend/internal/service/claude_code_validator_test.go b/backend/internal/service/claude_code_validator_test.go index f87c56e839b7..a4b3050587fb 100644 --- a/backend/internal/service/claude_code_validator_test.go +++ b/backend/internal/service/claude_code_validator_test.go @@ -48,6 +48,97 @@ func TestClaudeCodeValidator_MessagesWithoutProbeStillNeedStrictValidation(t *te require.False(t, ok) } +func TestClaudeCodeValidator_CountTokensPathUAOnly(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages/count_tokens", nil) + req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + }) + require.True(t, ok) +} + +func TestClaudeCodeValidator_CountTokensPathRequiresUA(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages/count_tokens", nil) + req.Header.Set("User-Agent", "curl/8.0.0") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + }) + require.False(t, ok) +} + +func TestClaudeCodeValidator_MessagesPathFullValid(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil) + req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)") + req.Header.Set("X-App", "claude-code") + req.Header.Set("anthropic-beta", "claude-code-20250219") + req.Header.Set("anthropic-version", "2023-06-01") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + "stream": true, + "system": []any{ + map[string]any{ + "type": "text", + "text": "You are Claude Code, Anthropic's official CLI for Claude.", + }, + }, + "metadata": map[string]any{ + "user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", + }, + }) + require.True(t, ok) +} + +func TestClaudeCodeValidator_MessagesPathRejectsNonClaudeCodeUA(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil) + req.Header.Set("User-Agent", "curl/8.0.0") + req.Header.Set("X-App", "claude-code") + req.Header.Set("anthropic-beta", "claude-code-20250219") + req.Header.Set("anthropic-version", "2023-06-01") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + "stream": true, + "system": []any{ + map[string]any{ + "type": "text", + "text": "You are Claude Code, Anthropic's official CLI for Claude.", + }, + }, + "metadata": map[string]any{ + "user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", + }, + }) + require.False(t, ok) +} + +func TestClaudeCodeValidator_MessagesPathWithoutSystemPromptStillRejected(t *testing.T) { + validator := NewClaudeCodeValidator() + req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil) + req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)") + req.Header.Set("X-App", "claude-code") + req.Header.Set("anthropic-beta", "claude-code-20250219") + req.Header.Set("anthropic-version", "2023-06-01") + + ok := validator.Validate(req, map[string]any{ + "model": "claude-opus-4-8", + "stream": true, + "messages": []any{ + map[string]any{"role": "user", "content": "hello"}, + }, + "metadata": map[string]any{ + "user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", + }, + }) + require.False(t, ok) +} + func TestClaudeCodeValidator_NonMessagesPathUAOnly(t *testing.T) { validator := NewClaudeCodeValidator() req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/models", nil) diff --git a/backend/internal/service/content_moderation_test.go b/backend/internal/service/content_moderation_test.go index 96d58cafc309..3657430d7de5 100644 --- a/backend/internal/service/content_moderation_test.go +++ b/backend/internal/service/content_moderation_test.go @@ -277,6 +277,10 @@ func (r *contentModerationTestUserRepo) DisableTotp(ctx context.Context, userID panic("unexpected DisableTotp call") } +func (r *contentModerationTestUserRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*User, error) { + return r.GetByID(ctx, id) +} + type contentModerationTestAuthCacheInvalidator struct { userIDs []int64 } diff --git a/backend/internal/service/error_policy_test.go b/backend/internal/service/error_policy_test.go index 297a954c3422..2aa7a421d1f7 100644 --- a/backend/internal/service/error_policy_test.go +++ b/backend/internal/service/error_policy_test.go @@ -389,6 +389,60 @@ func TestApplyErrorPolicy(t *testing.T) { } } +func TestApplyErrorPolicy_GeminiRateLimitBypassesCustomSkip(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + cache := &stubSmartRetryCache{} + rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + svc := &AntigravityGatewayService{ + rateLimitService: rlSvc, + accountRepo: repo, + cache: cache, + } + + account := &Account{ + ID: 31, + Type: AccountTypeAPIKey, + Platform: PlatformAntigravity, + Credentials: map[string]any{ + "custom_error_codes_enabled": true, + "custom_error_codes": []any{float64(500)}, + }, + } + body := []byte(`{ + "error": { + "status": "RESOURCE_EXHAUSTED", + "details": [ + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"}, + {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"} + ] + } + }`) + p := antigravityRetryLoopParams{ + ctx: context.Background(), + prefix: "[test]", + account: account, + accountRepo: repo, + groupID: 42, + sessionHash: "gemini:sticky", + handleError: func(context.Context, string, *Account, int, http.Header, []byte, string, int64, string, bool) *handleModelRateLimitResult { + t.Fatal("model rate limit should be handled before custom error fallback") + return nil + }, + } + + handled, outStatus, retErr := svc.applyErrorPolicy(p, http.StatusTooManyRequests, http.Header{}, body) + + require.True(t, handled) + require.Equal(t, http.StatusTooManyRequests, outStatus) + require.NoError(t, retErr) + require.Len(t, repo.modelRateLimitCalls, 2) + require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey) + require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey) + require.Len(t, cache.deleteCalls, 1) + require.Equal(t, int64(42), cache.deleteCalls[0].groupID) + require.Equal(t, "gemini:sticky", cache.deleteCalls[0].sessionHash) +} + // --------------------------------------------------------------------------- // errorPolicyRepoStub — minimal AccountRepository stub for error policy tests // --------------------------------------------------------------------------- diff --git a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go index 5cb03f3017ce..e2da89b5db87 100644 --- a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go +++ b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go @@ -112,7 +112,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ForwardStreamPreservesBodyAnd body := []byte(`{"model":"claude-3-7-sonnet-20250219","stream":true,"system":[{"type":"text","text":"x-anthropic-billing-header keep"}],"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: "claude-3-7-sonnet-20250219", Stream: true, } @@ -202,7 +202,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ForwardCountTokensPreservesBo body := []byte(`{"model":"claude-3-5-sonnet-latest","messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}],"thinking":{"type":"enabled"}}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: "claude-3-5-sonnet-latest", } @@ -344,7 +344,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ModelMappingEdgeCases(t *test body := []byte(`{"model":"` + tt.model + `","messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: tt.model, } @@ -429,7 +429,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ModelMappingPreservesOtherFie // 包含复杂字段的请求体：system、thinking、messages body := []byte(`{"model":"claude-sonnet-4-20250514","system":[{"type":"text","text":"You are a helpful assistant."}],"messages":[{"role":"user","content":[{"type":"text","text":"hello world"}]}],"thinking":{"type":"enabled","budget_tokens":5000},"max_tokens":1024}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: "claude-sonnet-4-20250514", } @@ -476,6 +476,66 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ModelMappingPreservesOtherFie require.Equal(t, int64(1024), gjson.GetBytes(sentBody, "max_tokens").Int(), "max_tokens 不应被修改") } +func TestGatewayService_AnthropicAPIKeyPassthrough_CountTokensFiltersGenerationFields(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil) + + body := []byte(`{"model":"claude-sonnet-4-20250514","system":[{"type":"text","text":"sys"}],"messages":[{"role":"user","content":"hello"}],"tools":[{"name":"tool","input_schema":{"type":"object"}}],"temperature":0.7,"top_p":0.9,"top_k":40,"stream":true,"stop_sequences":["END"],"max_tokens":1024,"thinking":{"type":"enabled","budget_tokens":5000}}`) + parsed := &ParsedRequest{ + Body: NewRequestBodyRef(body), + Model: "claude-sonnet-4-20250514", + } + + upstreamRespBody := `{"input_tokens":42}` + upstream := &anthropicHTTPUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(upstreamRespBody)), + }, + } + + svc := &GatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}, + httpUpstream: upstream, + rateLimitService: &RateLimitService{}, + } + + account := &Account{ + ID: 302, + Name: "count-token-filter-test", + Platform: PlatformAnthropic, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "upstream-key", + "base_url": "https://api.anthropic.com", + }, + Extra: map[string]any{"anthropic_passthrough": true}, + Status: StatusActive, + Schedulable: true, + } + + err := svc.ForwardCountTokens(context.Background(), c, account, parsed) + require.NoError(t, err) + + sentBody := upstream.lastBody + require.False(t, gjson.GetBytes(sentBody, "temperature").Exists()) + require.False(t, gjson.GetBytes(sentBody, "top_p").Exists()) + require.False(t, gjson.GetBytes(sentBody, "top_k").Exists()) + require.False(t, gjson.GetBytes(sentBody, "stream").Exists()) + require.False(t, gjson.GetBytes(sentBody, "stop_sequences").Exists()) + require.Equal(t, "claude-sonnet-4-20250514", gjson.GetBytes(sentBody, "model").String()) + require.Equal(t, "sys", gjson.GetBytes(sentBody, "system.0.text").String()) + require.Equal(t, "hello", gjson.GetBytes(sentBody, "messages.0.content").String()) + require.Equal(t, "tool", gjson.GetBytes(sentBody, "tools.0.name").String()) + require.Equal(t, int64(1024), gjson.GetBytes(sentBody, "max_tokens").Int()) + require.Equal(t, "enabled", gjson.GetBytes(sentBody, "thinking.type").String()) +} + // TestGatewayService_AnthropicAPIKeyPassthrough_EmptyModelSkipsMapping // 确保空模型名不会触发映射逻辑 func TestGatewayService_AnthropicAPIKeyPassthrough_EmptyModelSkipsMapping(t *testing.T) { @@ -487,7 +547,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_EmptyModelSkipsMapping(t *tes body := []byte(`{"messages":[{"role":"user","content":"hello"}]}`) parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), Model: "", // 空模型 } @@ -576,7 +636,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_CountTokens404PassthroughNotE c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil) body := []byte(`{"model":"claude-sonnet-4-5-20250929","messages":[{"role":"user","content":"hi"}]}`) - parsed := &ParsedRequest{Body: body, Model: "claude-sonnet-4-5-20250929"} + parsed := &ParsedRequest{Body: NewRequestBodyRef(body), Model: "claude-sonnet-4-5-20250929"} upstream := &anthropicHTTPUpstreamRecorder{ resp: &http.Response{ @@ -653,7 +713,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_BuildRequestRejectsInvalidBas }, } - _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(context.Background(), c, account, []byte(`{}`), "k") + _, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(context.Background(), c, account, []byte(`{}`), "k") require.Error(t, err) } @@ -678,7 +738,7 @@ func TestGatewayService_AnthropicOAuth_NotAffectedByAPIKeyPassthroughToggle(t *t require.False(t, account.IsAnthropicAPIKeyPassthroughEnabled()) - req, err := svc.buildUpstreamRequest(context.Background(), c, account, []byte(`{"model":"claude-3-7-sonnet-20250219"}`), "oauth-token", "oauth", "claude-3-7-sonnet-20250219", true, false) + req, _, err := svc.buildUpstreamRequest(context.Background(), c, account, []byte(`{"model":"claude-3-7-sonnet-20250219"}`), "oauth-token", "oauth", "claude-3-7-sonnet-20250219", true, false) require.NoError(t, err) require.Equal(t, "Bearer oauth-token", getHeaderRaw(req.Header, "authorization")) require.Contains(t, getHeaderRaw(req.Header, "anthropic-beta"), claude.BetaOAuth, "OAuth 链路仍应按原逻辑补齐 oauth beta") @@ -707,7 +767,7 @@ func TestGatewayService_AnthropicOAuth_ForwardPreservesBillingHeaderSystemBlock( c, _ := gin.CreateTestContext(rec) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - parsed, err := ParseGatewayRequest([]byte(tt.body), PlatformAnthropic) + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), PlatformAnthropic) require.NoError(t, err) upstream := &anthropicHTTPUpstreamRecorder{ diff --git a/backend/internal/service/gateway_anthropic_vertex_service_account_test.go b/backend/internal/service/gateway_anthropic_vertex_service_account_test.go index 2f42b0ab4a9c..be8c58675f12 100644 --- a/backend/internal/service/gateway_anthropic_vertex_service_account_test.go +++ b/backend/internal/service/gateway_anthropic_vertex_service_account_test.go @@ -35,7 +35,7 @@ func TestGatewayService_BuildAnthropicVertexServiceAccountRequest(t *testing.T) body := []byte(`{"model":"claude-sonnet-4-5","stream":false,"max_tokens":32,"messages":[{"role":"user","content":"hello"}]}`) svc := &GatewayService{} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, @@ -87,7 +87,7 @@ func TestGatewayService_BuildAnthropicVertexServiceAccount_StripsContextManageme body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},"messages":[{"role":"user","content":"hi"}]}`) svc := &GatewayService{} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "vertex-token", "service_account", "claude-haiku-4-5@20251001", false, false, ) @@ -117,7 +117,7 @@ func TestGatewayService_BuildAnthropicVertexServiceAccount_PreservesContextManag body := []byte(`{"model":"claude-sonnet-4-6","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "vertex-token", "service_account", "claude-sonnet-4-6@20260218", false, false, ) diff --git a/backend/internal/service/gateway_context_management_test.go b/backend/internal/service/gateway_context_management_test.go index c2263bdc02b4..51b12809e794 100644 --- a/backend/internal/service/gateway_context_management_test.go +++ b/backend/internal/service/gateway_context_management_test.go @@ -364,7 +364,7 @@ func TestBuildUpstreamRequestAnthropicAPIKeyPassthrough_StripsContextManagementW body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough( + req, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough( context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token", ) require.NoError(t, err) @@ -381,7 +381,7 @@ func TestBuildUpstreamRequestAnthropicAPIKeyPassthrough_PreservesContextManageme body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough( + req, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough( context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token", ) require.NoError(t, err) @@ -427,7 +427,7 @@ func TestBuildUpstreamRequest_OAuthMimicHaiku_StripsContextManagementEndToEnd(t // body 必须 strip。 body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "oauth-tok", "oauth", "claude-haiku-4-5", false, true, // mimicClaudeCode=true ) @@ -457,7 +457,7 @@ func TestBuildUpstreamRequest_OAuthMimicNonHaiku_PreservesContextManagementEndTo // body 保留。 body := []byte(`{"model":"claude-sonnet-4-6","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "oauth-tok", "oauth", "claude-sonnet-4-6", false, true, ) @@ -488,7 +488,7 @@ func TestBuildUpstreamRequest_OAuthTransparentHaikuWithRealCCBeta_PreservesField } body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "oauth-tok", "oauth", "claude-haiku-4-5", false, false, // mimicClaudeCode=false（真 CC） ) @@ -580,7 +580,7 @@ func TestBuildCountTokensRequest_OAuthMimicHaiku_PreservesContextManagementEndTo } body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildCountTokensRequest( + req, _, err := svc.buildCountTokensRequest( context.Background(), c, account, body, "oauth-tok", "oauth", "claude-haiku-4-5", true, // mimicClaudeCode=true ) @@ -611,7 +611,7 @@ func TestBuildCountTokensRequest_APIKeyHaiku_StripsContextManagementEndToEnd(t * } body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildCountTokensRequest( + req, _, err := svc.buildCountTokensRequest( context.Background(), c, account, body, "sk-ant-xxx", "apikey", "claude-haiku-4-5", false, ) @@ -655,7 +655,7 @@ func TestBuildUpstreamRequest_APIKeyHaikuWithContextManagement_StripsField(t *te } body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[]},"messages":[]}`) svc := &GatewayService{cfg: &config.Config{}} - req, err := svc.buildUpstreamRequest( + req, _, err := svc.buildUpstreamRequest( context.Background(), c, account, body, "sk-ant-xxx", "apikey", "claude-haiku-4-5", false, false, ) diff --git a/backend/internal/service/gateway_forward_as_chat_completions.go b/backend/internal/service/gateway_forward_as_chat_completions.go index eaf67fab0ec8..1df450d6b4a7 100644 --- a/backend/internal/service/gateway_forward_as_chat_completions.go +++ b/backend/internal/service/gateway_forward_as_chat_completions.go @@ -119,7 +119,7 @@ func (s *GatewayService) ForwardAsChatCompletions( // 10. Build upstream request upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream) - upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode) + upstreamReq, _, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode) releaseUpstreamCtx() if err != nil { return nil, fmt.Errorf("build upstream request: %w", err) @@ -148,7 +148,7 @@ func (s *GatewayService) ForwardAsChatCompletions( // 12. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/gateway_forward_as_responses.go b/backend/internal/service/gateway_forward_as_responses.go index c55a5a984116..22951b88b1e7 100644 --- a/backend/internal/service/gateway_forward_as_responses.go +++ b/backend/internal/service/gateway_forward_as_responses.go @@ -116,7 +116,7 @@ func (s *GatewayService) ForwardAsResponses( // 10. Build upstream request upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream) - upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode) + upstreamReq, _, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode) releaseUpstreamCtx() if err != nil { return nil, fmt.Errorf("build upstream request: %w", err) @@ -145,7 +145,7 @@ func (s *GatewayService) ForwardAsResponses( // 12. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index a00b7dc1d50d..7a6acaaca8a2 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -1229,6 +1229,106 @@ func TestGatewayService_selectAccountWithMixedScheduling(t *testing.T) { require.Equal(t, int64(2), acc.ID, "应选择优先级最高的账户（包含启用混合调度的antigravity）") }) + t.Run("混合调度-Gemini家族限流后跳过Antigravity账户", func(t *testing.T) { + resetAt := time.Now().Add(10 * time.Minute).Format(time.RFC3339) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{ + { + ID: 1, + Platform: PlatformAntigravity, + Priority: 1, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": resetAt, + }, + }, + }, + }, + { + ID: 2, + Platform: PlatformAntigravity, + Priority: 1, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": resetAt, + }, + }, + }, + }, + { + ID: 3, + Platform: PlatformAntigravity, + Priority: 2, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{"mixed_scheduling": true}, + }, + }, + accountsByID: map[int64]*Account{}, + } + for i := range repo.accounts { + repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] + } + + svc := &GatewayService{ + accountRepo: repo, + cache: &mockGatewayCacheForPlatform{}, + cfg: testConfig(), + } + + acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "gemini-3-pro-preview", nil, PlatformGemini) + require.NoError(t, err) + require.NotNil(t, acc) + require.Equal(t, int64(3), acc.ID) + }) + + t.Run("混合调度-Gemini家族限流不影响Claude调度", func(t *testing.T) { + resetAt := time.Now().Add(10 * time.Minute).Format(time.RFC3339) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{ + { + ID: 1, + Platform: PlatformAntigravity, + Priority: 1, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": resetAt, + }, + }, + }, + }, + {ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true}, + }, + accountsByID: map[int64]*Account{}, + } + for i := range repo.accounts { + repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] + } + + svc := &GatewayService{ + accountRepo: repo, + cache: &mockGatewayCacheForPlatform{}, + cfg: testConfig(), + } + + acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "claude-sonnet-4-5", nil, PlatformAnthropic) + require.NoError(t, err) + require.NotNil(t, acc) + require.Equal(t, int64(1), acc.ID) + }) + t.Run("混合调度-路由优先选择路由账号", func(t *testing.T) { groupID := int64(30) requestedModel := "claude-sonnet-4-5" diff --git a/backend/internal/service/gateway_oauth_metadata_test.go b/backend/internal/service/gateway_oauth_metadata_test.go index ed6f1887145a..b172dc6eebdb 100644 --- a/backend/internal/service/gateway_oauth_metadata_test.go +++ b/backend/internal/service/gateway_oauth_metadata_test.go @@ -14,8 +14,6 @@ func TestBuildOAuthMetadataUserID_FallbackWithoutAccountUUID(t *testing.T) { Model: "claude-sonnet-4-5", Stream: true, MetadataUserID: "", - System: nil, - Messages: nil, } account := &Account{ diff --git a/backend/internal/service/gateway_request.go b/backend/internal/service/gateway_request.go index 91f7601cb1ac..dc59611cead5 100644 --- a/backend/internal/service/gateway_request.go +++ b/backend/internal/service/gateway_request.go @@ -51,6 +51,164 @@ type SessionContext struct { APIKeyID int64 } +type jsonRange struct { + start int // 原始请求体中的起始偏移（闭区间） + end int // 原始请求体中的结束偏移（开区间） + kind gjson.Type // JSON 值类型，用于调用方做轻量分支 +} + +type RequestBodyRef struct { + data []byte +} + +func NewRequestBodyRef(data []byte) *RequestBodyRef { + return &RequestBodyRef{data: data} +} + +func (b *RequestBodyRef) Bytes() []byte { + if b == nil { + return nil + } + return b.data +} + +func (b *RequestBodyRef) Len() int { + if b == nil { + return 0 + } + return len(b.data) +} + +func (b *RequestBodyRef) Replace(data []byte) { + if b == nil { + return + } + b.data = data +} + +func missingJSONRange() jsonRange { + return jsonRange{start: -1, end: -1} +} + +func rangeFromResult(r gjson.Result) jsonRange { + if r.Raw == "" || r.Index <= 0 { + return missingJSONRange() + } + end := r.Index + len(r.Raw) + if end < r.Index { + return missingJSONRange() + } + return jsonRange{start: r.Index, end: end, kind: r.Type} +} + +func (r jsonRange) exists() bool { + return r.start >= 0 && r.end >= r.start +} + +// clearGatewayRequestDerivedState 清空绑定当前 body 的轻量派生字段，防止 ReplaceBody 后读到旧值。 +func clearGatewayRequestDerivedState(parsed *ParsedRequest) { + if parsed == nil { + return + } + parsed.Model = "" + parsed.Stream = false + parsed.MetadataUserID = "" + parsed.HasSystem = false + parsed.ThinkingEnabled = false + parsed.OutputEffort = "" + parsed.MaxTokens = 0 + parsed.systemRange = missingJSONRange() + parsed.messagesRange = missingJSONRange() +} + +func clearGatewayRequestRanges(parsed *ParsedRequest) { + if parsed == nil { + return + } + parsed.HasSystem = false + parsed.systemRange = missingJSONRange() + parsed.messagesRange = missingJSONRange() +} + +func setGatewayRequestRanges(parsed *ParsedRequest, protocol string, jsonStr string) { + if parsed == nil { + return + } + switch protocol { + case domain.PlatformGemini: + if sysParts := gjson.Get(jsonStr, "systemInstruction.parts"); sysParts.Exists() && sysParts.IsArray() { + parsed.systemRange = rangeFromResult(sysParts) + } + if contents := gjson.Get(jsonStr, "contents"); contents.Exists() && contents.IsArray() { + parsed.messagesRange = rangeFromResult(contents) + } + default: + if sys := gjson.Get(jsonStr, "system"); sys.Exists() { + parsed.HasSystem = true + parsed.systemRange = rangeFromResult(sys) + } + if msgs := gjson.Get(jsonStr, "messages"); msgs.Exists() && msgs.IsArray() { + parsed.messagesRange = rangeFromResult(msgs) + } + } +} + +// parseGatewayRequestCurrentBody 只做标量和 raw range 轻量解析，不恢复 system/messages 对象图。 +func parseGatewayRequestCurrentBody(parsed *ParsedRequest, protocol string) error { + if parsed == nil || parsed.Body == nil { + return fmt.Errorf("empty request body") + } + + bodyBytes := parsed.Body.Bytes() + if !gjson.ValidBytes(bodyBytes) { + return fmt.Errorf("invalid json") + } + + // 只在当前函数内零拷贝读取 JSON 字段；ReplaceBody 后必须重新进入本函数刷新派生状态。 + jsonStr := *(*string)(unsafe.Pointer(&bodyBytes)) + clearGatewayRequestDerivedState(parsed) + parsed.protocol = protocol + + modelResult := gjson.Get(jsonStr, "model") + if modelResult.Exists() { + if modelResult.Type != gjson.String { + return fmt.Errorf("invalid model field type") + } + parsed.Model = modelResult.String() + } + + streamResult := gjson.Get(jsonStr, "stream") + if streamResult.Exists() { + if streamResult.Type != gjson.True && streamResult.Type != gjson.False { + return fmt.Errorf("invalid stream field type") + } + parsed.Stream = streamResult.Bool() + } + + parsed.MetadataUserID = gjson.Get(jsonStr, "metadata.user_id").String() + + thinkingType := gjson.Get(jsonStr, "thinking.type").String() + parsed.ThinkingEnabled = thinkingType == "enabled" || thinkingType == "adaptive" + + parsed.OutputEffort = strings.TrimSpace(gjson.Get(jsonStr, "output_config.effort").String()) + + maxTokensResult := gjson.Get(jsonStr, "max_tokens") + if maxTokensResult.Exists() && maxTokensResult.Type == gjson.Number { + f := maxTokensResult.Float() + if !math.IsNaN(f) && !math.IsInf(f, 0) && f == math.Trunc(f) && + f <= float64(math.MaxInt) && f >= float64(math.MinInt) { + parsed.MaxTokens = int(f) + } + } + + setGatewayRequestRanges(parsed, protocol, jsonStr) + return nil +} + +func refreshGatewayRequestRanges(parsed *ParsedRequest, protocol string) error { + return parseGatewayRequestCurrentBody(parsed, protocol) +} + // ParsedRequest 保存网关请求的预解析结果 // // 性能优化说明： @@ -64,18 +222,20 @@ type SessionContext struct { // 2. 将解析结果 ParsedRequest 传递给 Service 层 // 3. 避免重复 json.Unmarshal，减少 CPU 和内存开销 type ParsedRequest struct { - Body []byte // 原始请求体（保留用于转发） + Body *RequestBodyRef // 原始请求体引用（保留用于转发）；替换内容请走 ReplaceBody Model string // 请求的模型名称 Stream bool // 是否为流式请求 MetadataUserID string // metadata.user_id（用于会话亲和） - System any // system 字段内容 - Messages []any // messages 数组 HasSystem bool // 是否包含 system 字段（包含 null 也视为显式传入） ThinkingEnabled bool // 是否开启 thinking（部分平台会影响最终模型名） OutputEffort string // output_config.effort（Claude API 的推理强度控制） MaxTokens int // max_tokens 值（用于探测请求拦截） SessionContext *SessionContext // 可选：请求上下文区分因子（nil 时行为不变） + protocol string // 当前 Body 的协议格式，用于 Body 替换后刷新 raw range + systemRange jsonRange // system/systemInstruction.parts 的 raw JSON 范围，绑定 Body 当前内容 + messagesRange jsonRange // messages/contents 的 raw JSON 范围，绑定 Body 当前内容 + // GroupID 请求所属分组 ID（来自 API Key） GroupID *int64 @@ -130,117 +290,90 @@ func normalizeSessionUserAgentFallback(raw string) string { // ParseGatewayRequest 解析网关请求体并返回结构化结果。 // protocol 指定请求协议格式（domain.PlatformAnthropic / domain.PlatformGemini）， // 不同协议使用不同的 system/messages 字段名。 -func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) { - // 保持与旧实现一致：请求体必须是合法 JSON。 - // 注意：gjson.GetBytes 对非法 JSON 不会报错，因此需要显式校验。 - if !gjson.ValidBytes(body) { - return nil, fmt.Errorf("invalid json") +func ParseGatewayRequest(body *RequestBodyRef, protocol string) (*ParsedRequest, error) { + parsed := &ParsedRequest{Body: body} + if err := parseGatewayRequestCurrentBody(parsed, protocol); err != nil { + return nil, err } + return parsed, nil +} - // 性能： - // - gjson.GetBytes 会把匹配的 Raw/Str 安全复制成 string（对于巨大 messages 会产生额外拷贝）。 - // - 这里将 body 通过 unsafe 零拷贝视为 string，仅在本函数内使用，且 body 不会被修改。 - jsonStr := *(*string)(unsafe.Pointer(&body)) - - parsed := &ParsedRequest{ - Body: body, +func (p *ParsedRequest) raw(r jsonRange) []byte { + if p == nil || p.Body == nil || !r.exists() { + return nil } - - // --- gjson 提取简单字段（避免完整 Unmarshal） --- - - // model: 需要严格类型校验，非 string 返回错误 - modelResult := gjson.Get(jsonStr, "model") - if modelResult.Exists() { - if modelResult.Type != gjson.String { - return nil, fmt.Errorf("invalid model field type") - } - parsed.Model = modelResult.String() + body := p.Body.Bytes() + if r.end > len(body) { + return nil } + return body[r.start:r.end] +} - // stream: 需要严格类型校验，非 bool 返回错误 - streamResult := gjson.Get(jsonStr, "stream") - if streamResult.Exists() { - if streamResult.Type != gjson.True && streamResult.Type != gjson.False { - return nil, fmt.Errorf("invalid stream field type") - } - parsed.Stream = streamResult.Bool() - } +func (p *ParsedRequest) SystemRaw() []byte { + return p.raw(p.systemRange) +} - // metadata.user_id: 直接路径提取，不需要严格类型校验 - parsed.MetadataUserID = gjson.Get(jsonStr, "metadata.user_id").String() +func (p *ParsedRequest) MessagesRaw() []byte { + return p.raw(p.messagesRange) +} - // thinking.type: enabled/adaptive 都视为开启 - thinkingType := gjson.Get(jsonStr, "thinking.type").String() - if thinkingType == "enabled" || thinkingType == "adaptive" { - parsed.ThinkingEnabled = true +func (p *ParsedRequest) DecodeSystem(dst any) error { + raw := p.SystemRaw() + if len(raw) == 0 { + return nil } + return json.Unmarshal(raw, dst) +} - // output_config.effort: Claude API 的推理强度控制参数 - parsed.OutputEffort = strings.TrimSpace(gjson.Get(jsonStr, "output_config.effort").String()) - - // max_tokens: 仅接受整数值 - maxTokensResult := gjson.Get(jsonStr, "max_tokens") - if maxTokensResult.Exists() && maxTokensResult.Type == gjson.Number { - f := maxTokensResult.Float() - if !math.IsNaN(f) && !math.IsInf(f, 0) && f == math.Trunc(f) && - f <= float64(math.MaxInt) && f >= float64(math.MinInt) { - parsed.MaxTokens = int(f) - } +func (p *ParsedRequest) DecodeMessages(dst any) error { + raw := p.MessagesRaw() + if len(raw) == 0 { + return nil } + return json.Unmarshal(raw, dst) +} - // --- system/messages 提取 --- - // 避免把整个 body Unmarshal 到 map（会产生大量 map/接口分配）。 - // 使用 gjson 抽取目标字段的 Raw，再对该子树进行 Unmarshal。 - - switch protocol { - case domain.PlatformGemini: - // Gemini 原生格式: systemInstruction.parts / contents - if sysParts := gjson.Get(jsonStr, "systemInstruction.parts"); sysParts.Exists() && sysParts.IsArray() { - var parts []any - if err := json.Unmarshal(sliceRawFromBody(body, sysParts), &parts); err != nil { - return nil, err - } - parsed.System = parts - } - - if contents := gjson.Get(jsonStr, "contents"); contents.Exists() && contents.IsArray() { - var msgs []any - if err := json.Unmarshal(sliceRawFromBody(body, contents), &msgs); err != nil { - return nil, err - } - parsed.Messages = msgs - } - default: - // Anthropic / OpenAI 格式: system / messages - // system 字段只要存在就视为显式提供（即使为 null）， - // 以避免客户端传 null 时被默认 system 误注入。 - if sys := gjson.Get(jsonStr, "system"); sys.Exists() { - parsed.HasSystem = true - switch sys.Type { - case gjson.Null: - parsed.System = nil - case gjson.String: - // 与 encoding/json 的 Unmarshal 行为一致：返回解码后的字符串。 - parsed.System = sys.String() - default: - var system any - if err := json.Unmarshal(sliceRawFromBody(body, sys), &system); err != nil { - return nil, err - } - parsed.System = system - } - } +func (p *ParsedRequest) SystemValue() (any, bool) { + raw := p.SystemRaw() + if len(raw) == 0 { + return nil, false + } + var system any + if err := json.Unmarshal(raw, &system); err != nil { + return nil, false + } + return system, true +} - if msgs := gjson.Get(jsonStr, "messages"); msgs.Exists() && msgs.IsArray() { - var messages []any - if err := json.Unmarshal(sliceRawFromBody(body, msgs), &messages); err != nil { - return nil, err - } - parsed.Messages = messages - } +// CloneForBody 为单次账号尝试创建独立 body 视图，避免 failover 复用已改写的 ParsedRequest。 +func (p *ParsedRequest) CloneForBody(body []byte) (*ParsedRequest, error) { + if p == nil { + return nil, fmt.Errorf("parse request: empty request") + } + clone := *p + clone.Body = NewRequestBodyRef(body) + clone.OnUpstreamAccepted = nil + if err := refreshGatewayRequestRanges(&clone, clone.protocol); err != nil { + return nil, err } + return &clone, nil +} - return parsed, nil +// ReplaceBody 统一刷新当前 body 和 raw range，保证后续 helper 读取的是最新请求体。 +func (p *ParsedRequest) ReplaceBody(data []byte) error { + if p == nil { + return fmt.Errorf("parse request: empty request") + } + if p.Body == nil { + p.Body = NewRequestBodyRef(data) + } else { + p.Body.Replace(data) + } + if err := refreshGatewayRequestRanges(p, p.protocol); err != nil { + clearGatewayRequestRanges(p) + return err + } + return nil } // sliceRawFromBody 返回 Result.Raw 对应的原始字节切片。 diff --git a/backend/internal/service/gateway_request_test.go b/backend/internal/service/gateway_request_test.go index 045dc66c15bd..288c031c37d4 100644 --- a/backend/internal/service/gateway_request_test.go +++ b/backend/internal/service/gateway_request_test.go @@ -10,24 +10,25 @@ import ( "github.com/Wei-Shaw/sub2api/internal/domain" "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" ) func TestParseGatewayRequest(t *testing.T) { body := []byte(`{"model":"claude-3-7-sonnet","stream":true,"metadata":{"user_id":"session_123e4567-e89b-12d3-a456-426614174000"},"system":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}],"messages":[{"content":"hi"}]}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, "claude-3-7-sonnet", parsed.Model) require.True(t, parsed.Stream) require.Equal(t, "session_123e4567-e89b-12d3-a456-426614174000", parsed.MetadataUserID) require.True(t, parsed.HasSystem) - require.NotNil(t, parsed.System) - require.Len(t, parsed.Messages, 1) + require.NotEmpty(t, parsed.SystemRaw()) + require.NotEmpty(t, parsed.MessagesRaw()) require.False(t, parsed.ThinkingEnabled) } func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) { body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"enabled"},"messages":[{"content":"hi"}]}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, "claude-sonnet-4-5", parsed.Model) require.True(t, parsed.ThinkingEnabled) @@ -35,7 +36,7 @@ func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) { func TestParseGatewayRequest_ThinkingAdaptiveEnabled(t *testing.T) { body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"adaptive"},"messages":[{"content":"hi"}]}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, "claude-sonnet-4-5", parsed.Model) require.True(t, parsed.ThinkingEnabled) @@ -43,36 +44,36 @@ func TestParseGatewayRequest_ThinkingAdaptiveEnabled(t *testing.T) { func TestParseGatewayRequest_MaxTokens(t *testing.T) { body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, 1, parsed.MaxTokens) } func TestParseGatewayRequest_MaxTokensNonIntegralIgnored(t *testing.T) { body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1.5}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) require.Equal(t, 0, parsed.MaxTokens) } func TestParseGatewayRequest_SystemNull(t *testing.T) { body := []byte(`{"model":"claude-3","system":null}`) - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.NoError(t, err) // 显式传入 system:null 也应视为“字段已存在”，避免默认 system 被注入。 require.True(t, parsed.HasSystem) - require.Nil(t, parsed.System) + require.Equal(t, []byte("null"), parsed.SystemRaw()) } func TestParseGatewayRequest_InvalidModelType(t *testing.T) { body := []byte(`{"model":123}`) - _, err := ParseGatewayRequest(body, "") + _, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.Error(t, err) } func TestParseGatewayRequest_InvalidStreamType(t *testing.T) { body := []byte(`{"stream":"true"}`) - _, err := ParseGatewayRequest(body, "") + _, err := ParseGatewayRequest(NewRequestBodyRef(body), "") require.Error(t, err) } @@ -86,11 +87,11 @@ func TestParseGatewayRequest_GeminiContents(t *testing.T) { {"role": "user", "parts": [{"text": "How are you?"}]} ] }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) - require.Len(t, parsed.Messages, 3, "should parse contents as Messages") + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 3, "should parse contents as Messages") require.False(t, parsed.HasSystem, "Gemini format should not set HasSystem") - require.Nil(t, parsed.System, "no systemInstruction means nil System") + require.Nil(t, parsed.SystemRaw(), "no systemInstruction means nil System") } func TestParseGatewayRequest_GeminiSystemInstruction(t *testing.T) { @@ -102,16 +103,13 @@ func TestParseGatewayRequest_GeminiSystemInstruction(t *testing.T) { {"role": "user", "parts": [{"text": "Hello"}]} ] }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) - require.NotNil(t, parsed.System, "should parse systemInstruction.parts as System") - parts, ok := parsed.System.([]any) - require.True(t, ok) - require.Len(t, parts, 1) - partMap, ok := parts[0].(map[string]any) - require.True(t, ok) - require.Equal(t, "You are a helpful assistant.", partMap["text"]) - require.Len(t, parsed.Messages, 1) + system := gjson.ParseBytes(parsed.SystemRaw()) + require.True(t, system.IsArray(), "should parse systemInstruction.parts as System") + require.Len(t, system.Array(), 1) + require.Equal(t, "You are a helpful assistant.", system.Get("0.text").String()) + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 1) } func TestParseGatewayRequest_GeminiWithModel(t *testing.T) { @@ -119,10 +117,10 @@ func TestParseGatewayRequest_GeminiWithModel(t *testing.T) { "model": "gemini-2.5-pro", "contents": [{"role": "user", "parts": [{"text": "test"}]}] }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.Equal(t, "gemini-2.5-pro", parsed.Model) - require.Len(t, parsed.Messages, 1) + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 1) } func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) { @@ -132,25 +130,25 @@ func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) { "messages": [{"role": "user", "content": "ignored"}], "contents": [{"role": "user", "parts": [{"text": "real content"}]}] }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) require.False(t, parsed.HasSystem, "Gemini protocol should not parse Anthropic system field") - require.Nil(t, parsed.System, "no systemInstruction = nil System") - require.Len(t, parsed.Messages, 1, "should use contents, not messages") + require.Nil(t, parsed.SystemRaw(), "no systemInstruction = nil System") + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 1, "should use contents, not messages") } func TestParseGatewayRequest_GeminiEmptyContents(t *testing.T) { body := []byte(`{"contents": []}`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) - require.Empty(t, parsed.Messages) + require.Empty(t, gjson.ParseBytes(parsed.MessagesRaw()).Array()) } func TestParseGatewayRequest_GeminiNoContents(t *testing.T) { body := []byte(`{"model": "gemini-2.5-flash"}`) - parsed, err := ParseGatewayRequest(body, domain.PlatformGemini) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) require.NoError(t, err) - require.Nil(t, parsed.Messages) + require.Nil(t, parsed.MessagesRaw()) require.Equal(t, "gemini-2.5-flash", parsed.Model) } @@ -162,14 +160,13 @@ func TestParseGatewayRequest_AnthropicIgnoresGeminiFields(t *testing.T) { "contents": [{"role": "user", "parts": [{"text": "ignored"}]}], "systemInstruction": {"parts": [{"text": "ignored"}]} }`) - parsed, err := ParseGatewayRequest(body, domain.PlatformAnthropic) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic) require.NoError(t, err) require.True(t, parsed.HasSystem) - require.Equal(t, "real system", parsed.System) - require.Len(t, parsed.Messages, 1) - msg, ok := parsed.Messages[0].(map[string]any) - require.True(t, ok) - require.Equal(t, "real content", msg["content"]) + require.Equal(t, "real system", gjson.ParseBytes(parsed.SystemRaw()).String()) + messages := gjson.ParseBytes(parsed.MessagesRaw()).Array() + require.Len(t, messages, 1) + require.Equal(t, "real content", messages[0].Get("content").String()) } func TestFilterThinkingBlocks(t *testing.T) { @@ -897,7 +894,7 @@ func TestParseGatewayRequest_TypeValidation(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - _, err := ParseGatewayRequest([]byte(tt.body), "") + _, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "") if tt.wantErr { require.Error(t, err) if tt.errSubstr != "" { @@ -959,7 +956,7 @@ func TestParseGatewayRequest_OptionalFieldsMissing(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - parsed, err := ParseGatewayRequest([]byte(tt.body), "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "") require.NoError(t, err) require.Equal(t, tt.wantModel, parsed.Model) @@ -970,10 +967,10 @@ func TestParseGatewayRequest_OptionalFieldsMissing(t *testing.T) { require.Equal(t, tt.wantMaxTokens, parsed.MaxTokens) if tt.wantMessagesNil { - require.Nil(t, parsed.Messages) + require.Nil(t, parsed.MessagesRaw()) } if tt.wantMessagesLen > 0 { - require.Len(t, parsed.Messages, tt.wantMessagesLen) + require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), tt.wantMessagesLen) } }) } @@ -1023,7 +1020,7 @@ func TestParseGatewayRequest_MaxTokensBoundary(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - parsed, err := ParseGatewayRequest([]byte(tt.body), "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "") if tt.wantErr { require.Error(t, err) return @@ -1040,7 +1037,7 @@ func TestParseGatewayRequest_MaxTokensBoundary(t *testing.T) { // 核心路径：先 Unmarshal 到 map[string]any，再逐字段提取。 func parseGatewayRequestOld(body []byte, protocol string) (*ParsedRequest, error) { parsed := &ParsedRequest{ - Body: body, + Body: NewRequestBodyRef(body), } var req map[string]any @@ -1087,25 +1084,8 @@ func parseGatewayRequestOld(body []byte, protocol string) (*ParsedRequest, error } } - // system / messages（按协议分支） - switch protocol { - case domain.PlatformGemini: - if sysInst, ok := req["systemInstruction"].(map[string]any); ok { - if parts, ok := sysInst["parts"].([]any); ok { - parsed.System = parts - } - } - if contents, ok := req["contents"].([]any); ok { - parsed.Messages = contents - } - default: - if system, ok := req["system"]; ok { - parsed.HasSystem = true - parsed.System = system - } - if messages, ok := req["messages"].([]any); ok { - parsed.Messages = messages - } + if err := refreshGatewayRequestRanges(parsed, protocol); err != nil { + return nil, err } return parsed, nil @@ -1151,7 +1131,7 @@ func BenchmarkParseGatewayRequest_New_Small(b *testing.B) { b.SetBytes(int64(len(data))) b.ResetTimer() for i := 0; i < b.N; i++ { - _, _ = ParseGatewayRequest(data, "") + _, _ = ParseGatewayRequest(NewRequestBodyRef(data), "") } } @@ -1203,7 +1183,7 @@ func TestParseGatewayRequest_OutputEffort(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - parsed, err := ParseGatewayRequest([]byte(tt.body), "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "") require.NoError(t, err) require.Equal(t, tt.wantEffort, parsed.OutputEffort) }) @@ -1245,6 +1225,6 @@ func BenchmarkParseGatewayRequest_New_Large(b *testing.B) { b.SetBytes(int64(len(data))) b.ResetTimer() for i := 0; i < b.N; i++ { - _, _ = ParseGatewayRequest(data, "") + _, _ = ParseGatewayRequest(NewRequestBodyRef(data), "") } } diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 7c48f243d5ad..812780dc4b37 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -23,6 +23,7 @@ import ( "sync/atomic" "syscall" "time" + "unsafe" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/claude" @@ -56,6 +57,8 @@ const ( defaultModelsListCacheTTL = 15 * time.Second postUsageBillingTimeout = 15 * time.Second debugGatewayBodyEnv = "SUB2API_DEBUG_GATEWAY_BODY" + // 上游错误体只需要提取错误 JSON/日志摘要，默认 512KiB 避免错误风暴叠加大请求体。 + gatewayUpstreamErrorBodyReadLimit int64 = 512 << 10 ) const ( @@ -96,15 +99,20 @@ var ( modelsListCacheMissTotal atomic.Int64 modelsListCacheStoreTotal atomic.Int64 + // Deprecated: flusher_enabled=true 后不再增长(仅 flag=false 降级直写路径使用);新主路径见 FlusherMetrics。remove after 2026-09。 // userPlatformQuotaDBIncrErrorTotal 统计 finalizePostUsageBilling 异步 goroutine // 中 IncrementUsageWithReset 失败次数。Redis 已成功累加 + DB 写失败意味着 // Redis cache TTL 过期或被清后该笔 cost 会丢失（与实际消费偏差）。 // oncall 通过 GatewayUserPlatformQuotaIncrStats() 暴露给 ops 面板做阈值告警。 userPlatformQuotaDBIncrErrorTotal atomic.Int64 + // Deprecated: flusher_enabled=true 后不再增长(仅 flag=false 降级直写路径使用);新主路径见 FlusherMetrics。remove after 2026-09。 // userPlatformQuotaDBIncrLegacyErrorTotal 统计 legacy postUsageBilling // （applyUsageBilling 在 repo==nil 时 fallback）路径下的失败次数； // 与 DB Incr 失败分开计数，便于区分"主路径暂时故障"vs"基础设施长期未配齐"。 userPlatformQuotaDBIncrLegacyErrorTotal atomic.Int64 + // userPlatformQuotaSentinelSetCacheErrorTotal 统计 checkUserPlatformQuotaEligibility + // 在 DB 无行时回填 sentinel cache entry 写 Redis 失败的次数（phase A）。 + userPlatformQuotaSentinelSetCacheErrorTotal atomic.Int64 ) func GatewayWindowCostPrefetchStats() (cacheHit, cacheMiss, batchSQL, fallback, errCount int64) { @@ -127,13 +135,32 @@ func GatewayModelsListCacheStats() (cacheHit, cacheMiss, store int64) { return modelsListCacheHitTotal.Load(), modelsListCacheMissTotal.Load(), modelsListCacheStoreTotal.Load() } -// GatewayUserPlatformQuotaIncrStats 返回 (mainPathErr, legacyPathErr)。 +// GatewayUserPlatformQuotaIncrStats 返回 (mainPathErr, legacyPathErr, sentinelSetErr)。 // mainPathErr：finalizePostUsageBilling 异步 goroutine 写 DB 失败累计次数； -// legacyPathErr：postUsageBilling fallback 路径写 DB 失败累计次数。 +// legacyPathErr：postUsageBilling fallback 路径写 DB 失败累计次数； +// sentinelSetErr：DB 无行时回填 sentinel cache entry 写 Redis 失败累计次数。 // ops 监控面板可以按"持续上升斜率"做告警阈值。 -func GatewayUserPlatformQuotaIncrStats() (mainPathErr, legacyPathErr int64) { +func GatewayUserPlatformQuotaIncrStats() (mainPathErr, legacyPathErr, sentinelSetErr int64) { return userPlatformQuotaDBIncrErrorTotal.Load(), - userPlatformQuotaDBIncrLegacyErrorTotal.Load() + userPlatformQuotaDBIncrLegacyErrorTotal.Load(), + userPlatformQuotaSentinelSetCacheErrorTotal.Load() +} + +// GatewayUserPlatformQuotaFlusherStats 暴露 flusher 运行指标供 ops/health 面板查询。 +func GatewayUserPlatformQuotaFlusherStats(f *UserPlatformQuotaUsageFlusher) map[string]int64 { + if f == nil || f.metrics == nil { + return nil + } + m := f.metrics + return map[string]int64{ + "flush_success": m.FlushSuccessTotal.Load(), + "flush_error": m.FlushErrorTotal.Load(), + "flush_batch_size": m.FlushBatchSizeTotal.Load(), + "flush_latency_ms_max": m.FlushLatencyMsMax.Load(), + "dirty_readd": m.DirtyReaddTotal.Load(), + "dirty_lost": m.DirtyLostTotal.Load(), + "flush_fk_violation": m.FlushFKViolationTotal.Load(), + } } func openAIStreamEventIsTerminal(data string) bool { @@ -724,31 +751,10 @@ func (s *GatewayService) GenerateSessionHash(parsed *ParsedRequest) string { _, _ = combined.WriteString(strconv.FormatInt(parsed.SessionContext.APIKeyID, 10)) _, _ = combined.WriteString("|") } - if parsed.System != nil { - systemText := s.extractTextFromSystem(parsed.System) - if systemText != "" { - _, _ = combined.WriteString(systemText) - } - } - for _, msg := range parsed.Messages { - if m, ok := msg.(map[string]any); ok { - if content, exists := m["content"]; exists { - // Anthropic: messages[].content - if msgText := s.extractTextFromContent(content); msgText != "" { - _, _ = combined.WriteString(msgText) - } - } else if parts, ok := m["parts"].([]any); ok { - // Gemini: contents[].parts[].text - for _, part := range parts { - if partMap, ok := part.(map[string]any); ok { - if text, ok := partMap["text"].(string); ok { - _, _ = combined.WriteString(text) - } - } - } - } - } + if systemText := extractTextFromSystemRaw(parsed.SystemRaw()); systemText != "" { + _, _ = combined.WriteString(systemText) } + appendMessageTextsFromRaw(&combined, parsed.MessagesRaw()) if combined.Len() > 0 { hash := s.hashContent(combined.String()) slog.Info("sticky.hash_source", @@ -823,80 +829,133 @@ func (s *GatewayService) extractCacheableContent(parsed *ParsedRequest) string { return "" } - var builder strings.Builder + systemText := extractCacheableTextFromSystemRaw(parsed.SystemRaw()) + if messageText := extractCacheableTextFromMessagesRaw(parsed.MessagesRaw()); messageText != "" { + return messageText + } + return systemText +} - // 检查 system 中的 cacheable 内容 - if system, ok := parsed.System.([]any); ok { - for _, part := range system { - if partMap, ok := part.(map[string]any); ok { - if cc, ok := partMap["cache_control"].(map[string]any); ok { - if cc["type"] == "ephemeral" { - if text, ok := partMap["text"].(string); ok { - _, _ = builder.WriteString(text) - } - } - } - } - } +func parseRawJSONView(raw []byte) gjson.Result { + if len(raw) == 0 { + return gjson.Result{} } - systemText := builder.String() + // 这里只做同步只读解析，避免 gjson.ParseBytes 为大 messages/contents 复制整段 raw。 + return gjson.Parse(*(*string)(unsafe.Pointer(&raw))) +} - // 检查 messages 中的 cacheable 内容 - for _, msg := range parsed.Messages { - if msgMap, ok := msg.(map[string]any); ok { - if msgContent, ok := msgMap["content"].([]any); ok { - for _, part := range msgContent { - if partMap, ok := part.(map[string]any); ok { - if cc, ok := partMap["cache_control"].(map[string]any); ok { - if cc["type"] == "ephemeral" { - return s.extractTextFromContent(msgMap["content"]) - } - } - } - } - } +func extractTextFromSystemRaw(raw []byte) string { + system := parseRawJSONView(raw) + switch system.Type { + case gjson.String: + return system.String() + case gjson.JSON: + if !system.IsArray() { + return "" } + var builder strings.Builder + system.ForEach(func(_, part gjson.Result) bool { + if text := part.Get("text").String(); text != "" { + _, _ = builder.WriteString(text) + } + return true + }) + return builder.String() } - - return systemText + return "" } -func (s *GatewayService) extractTextFromSystem(system any) string { - switch v := system.(type) { - case string: - return v - case []any: - var texts []string - for _, part := range v { - if partMap, ok := part.(map[string]any); ok { - if text, ok := partMap["text"].(string); ok { - texts = append(texts, text) +func extractTextFromContentRaw(content gjson.Result) string { + switch content.Type { + case gjson.String: + return content.String() + case gjson.JSON: + if !content.IsArray() { + return "" + } + var builder strings.Builder + content.ForEach(func(_, part gjson.Result) bool { + if part.Get("type").String() == "text" { + if text := part.Get("text").String(); text != "" { + _, _ = builder.WriteString(text) } } - } - return strings.Join(texts, "") + return true + }) + return builder.String() } return "" } -func (s *GatewayService) extractTextFromContent(content any) string { - switch v := content.(type) { - case string: - return v - case []any: - var texts []string - for _, part := range v { - if partMap, ok := part.(map[string]any); ok { - if partMap["type"] == "text" { - if text, ok := partMap["text"].(string); ok { - texts = append(texts, text) - } +func appendMessageTextsFromRaw(builder *strings.Builder, raw []byte) { + if builder == nil || len(raw) == 0 { + return + } + messages := parseRawJSONView(raw) + if !messages.IsArray() { + return + } + messages.ForEach(func(_, msg gjson.Result) bool { + if content := msg.Get("content"); content.Exists() { + _, _ = builder.WriteString(extractTextFromContentRaw(content)) + return true + } + parts := msg.Get("parts") + if parts.IsArray() { + parts.ForEach(func(_, part gjson.Result) bool { + if text := part.Get("text").String(); text != "" { + _, _ = builder.WriteString(text) } + return true + }) + } + return true + }) +} + +func extractCacheableTextFromSystemRaw(raw []byte) string { + system := parseRawJSONView(raw) + if !system.IsArray() { + return "" + } + var builder strings.Builder + system.ForEach(func(_, part gjson.Result) bool { + if part.Get("cache_control.type").String() == "ephemeral" { + if text := part.Get("text").String(); text != "" { + _, _ = builder.WriteString(text) } } - return strings.Join(texts, "") + return true + }) + return builder.String() +} + +func extractCacheableTextFromMessagesRaw(raw []byte) string { + messages := parseRawJSONView(raw) + if !messages.IsArray() { + return "" } - return "" + var text string + messages.ForEach(func(_, msg gjson.Result) bool { + content := msg.Get("content") + if !content.IsArray() { + return true + } + found := false + content.ForEach(func(_, part gjson.Result) bool { + if part.Get("cache_control.type").String() == "ephemeral" { + found = true + return false + } + return true + }) + if found { + text = extractTextFromContentRaw(content) + return false + } + return true + }) + return text } func (s *GatewayService) hashContent(content string) string { @@ -1260,7 +1319,7 @@ func (s *GatewayService) applyClaudeCodeOAuthMimicryToBody( systemRewritten := false if !strings.Contains(strings.ToLower(model), "haiku") { - body = rewriteSystemForNonClaudeCode(body, systemRaw) + body = rewriteSystemForNonClaudeCode(body, normalizeSystemParam(systemRaw)) systemRewritten = true } @@ -4376,12 +4435,12 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } // Web Search 模拟：纯 web_search 请求时，直接调用搜索 API 构造响应 - if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body) { + if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body.Bytes()) { return s.handleWebSearchEmulation(ctx, c, account, parsed) } if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() { - passthroughBody := parsed.Body + passthroughBody := parsed.Body.Bytes() passthroughModel := parsed.Model if passthroughModel != "" { if mappedModel := account.GetMappedModel(passthroughModel); mappedModel != passthroughModel { @@ -4392,6 +4451,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } return s.forwardAnthropicAPIKeyPassthroughWithInput(ctx, c, account, anthropicPassthroughForwardInput{ Body: passthroughBody, + Parsed: parsed, RequestModel: passthroughModel, OriginalModel: parsed.Model, RequestStream: parsed.Stream, @@ -4417,7 +4477,14 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A c.Set(betaPolicyFilterSetKey, filterSet) } - body := parsed.Body + body := parsed.Body.Bytes() + replaceBody := func(next []byte) error { + if err := parsed.ReplaceBody(next); err != nil { + return fmt.Errorf("rewrite request body: %w", err) + } + body = parsed.Body.Bytes() + return nil + } reqModel := parsed.Model reqStream := parsed.Stream originalModel := reqModel @@ -4450,7 +4517,10 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // Parrot 的 transform_request 从不检查客户端 system 内容，直接覆盖。 systemRewritten := false if !strings.Contains(strings.ToLower(reqModel), "haiku") { - body = rewriteSystemForNonClaudeCode(body, parsed.System) + systemRaw, _ := parsed.SystemValue() + if err := replaceBody(rewriteSystemForNonClaudeCode(body, systemRaw)); err != nil { + return nil, err + } systemRewritten = true } @@ -4472,22 +4542,34 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } } - body, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + var normalizedBody []byte + normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + if err := replaceBody(normalizedBody); err != nil { + return nil, err + } // D/E/F: 可选 messages cache 策略 + 工具名混淆 + tools[-1] 断点 // 与 forward_as_chat_completions / forward_as_responses 路径对齐， // 原生 /v1/messages 路径也走同一套可配置字段级改写。 - body = s.rewriteMessageCacheControlIfEnabled(ctx, body) + if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil { + return nil, err + } if rw := buildToolNameRewriteFromBody(body); rw != nil { - body = applyToolNameRewriteToBody(body, rw) + if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil { + return nil, err + } c.Set(toolNameRewriteKey, rw) } else { - body = applyToolsLastCacheBreakpoint(body) + if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil { + return nil, err + } } } // 强制执行 cache_control 块数量限制（最多 4 个） - body = enforceCacheControlLimit(body) + if err := replaceBody(enforceCacheControlLimit(body)); err != nil { + return nil, err + } // 应用模型映射： // - APIKey 账号：使用账号级别的显式映射（如果配置），否则透传原始模型名 @@ -4521,13 +4603,18 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } if mappedModel != reqModel { // 替换请求体中的模型名 - body = s.replaceModelInBody(body, mappedModel) + if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil { + return nil, err + } reqModel = mappedModel + parsed.Model = mappedModel logger.LegacyPrintf("service.gateway", "Model mapping applied: %s -> %s (account: %s, source=%s)", originalModel, mappedModel, account.Name, mappingSource) } if s.shouldInjectAnthropicCacheTTL1h(ctx, account) { - body = injectAnthropicCacheControlTTL1h(body) + if err := replaceBody(injectAnthropicCacheControlTTL1h(body)); err != nil { + return nil, err + } } // 获取凭证 @@ -4551,19 +4638,24 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A logger.LegacyPrintf("service.gateway", "[Forward] Using account: ID=%d Name=%s Platform=%s Type=%s TLSFingerprint=%v Proxy=%s", account.ID, account.Name, account.Platform, account.Type, tlsProfile, proxyURL) // Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400. - body = StripEmptyTextBlocks(body) + if err := replaceBody(StripEmptyTextBlocks(body)); err != nil { + return nil, err + } // 重试循环 var resp *http.Response + lastWireBody := body retryStart := time.Now() for attempt := 1; attempt <= maxRetryAttempts; attempt++ { // 构建上游请求（每次重试需要重新构建，因为请求体需要重新读取） upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream) - upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + upstreamReq, wireBody, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) releaseUpstreamCtx() if err != nil { return nil, err } + // 记录本次实际发送的 wire body；只有请求成功后才写回 ParsedRequest，避免 400 retry 基于已签名 CCH 再改写。 + lastWireBody = wireBody // 发送请求 resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, tlsProfile) @@ -4595,7 +4687,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // 优先检测thinking block签名错误（400）并重试一次 if resp.StatusCode == 400 { - respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, readErr := s.readUpstreamErrorBody(resp) if readErr == nil { _ = resp.Body.Close() @@ -4641,18 +4733,24 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A filteredBody := FilterThinkingBlocksForRetry(body) retryCtx, releaseRetryCtx := detachStreamUpstreamContext(ctx, reqStream) - retryReq, buildErr := s.buildUpstreamRequest(retryCtx, c, account, filteredBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + retryReq, retryWireBody, buildErr := s.buildUpstreamRequest(retryCtx, c, account, filteredBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) releaseRetryCtx() if buildErr == nil { retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, tlsProfile) if retryErr == nil { if retryResp.StatusCode < 400 { + // 重试请求被上游接受后同步 ParsedRequest，保证 usage/日志看到真实请求体。 + lastWireBody = retryWireBody + if err := replaceBody(retryWireBody); err != nil { + _ = retryResp.Body.Close() + return nil, err + } logger.LegacyPrintf("service.gateway", "Account %d: thinking block retry succeeded (blocks downgraded)", account.ID) resp = retryResp break } - retryRespBody, retryReadErr := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20)) + retryRespBody, retryReadErr := s.readUpstreamErrorBody(retryResp) _ = retryResp.Body.Close() if retryReadErr == nil && retryResp.StatusCode == 400 && s.isSignatureErrorPattern(ctx, account, retryRespBody) { appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ @@ -4676,11 +4774,19 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A logger.LegacyPrintf("service.gateway", "Account %d: signature retry still failing and looks tool-related, retrying with tool blocks downgraded", account.ID) filteredBody2 := FilterSignatureSensitiveBlocksForRetry(body) retryCtx2, releaseRetryCtx2 := detachStreamUpstreamContext(ctx, reqStream) - retryReq2, buildErr2 := s.buildUpstreamRequest(retryCtx2, c, account, filteredBody2, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + retryReq2, retryWireBody2, buildErr2 := s.buildUpstreamRequest(retryCtx2, c, account, filteredBody2, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) releaseRetryCtx2() if buildErr2 == nil { retryResp2, retryErr2 := s.httpUpstream.DoWithTLS(retryReq2, proxyURL, account.ID, account.Concurrency, tlsProfile) if retryErr2 == nil { + if retryResp2.StatusCode < 400 { + // 二阶段工具块降级成功时也必须更新当前 body。 + lastWireBody = retryWireBody2 + if err := replaceBody(retryWireBody2); err != nil { + _ = retryResp2.Body.Close() + return nil, err + } + } resp = retryResp2 break } @@ -4747,11 +4853,19 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if applied && time.Since(retryStart) < maxRetryElapsed { logger.LegacyPrintf("service.gateway", "Account %d: detected budget_tokens constraint error, retrying with rectified budget (budget_tokens=%d, max_tokens=%d)", account.ID, BudgetRectifyBudgetTokens, BudgetRectifyMaxTokens) budgetRetryCtx, releaseBudgetRetryCtx := detachStreamUpstreamContext(ctx, reqStream) - budgetRetryReq, buildErr := s.buildUpstreamRequest(budgetRetryCtx, c, account, rectifiedBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + budgetRetryReq, budgetWireBody, buildErr := s.buildUpstreamRequest(budgetRetryCtx, c, account, rectifiedBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) releaseBudgetRetryCtx() if buildErr == nil { budgetRetryResp, retryErr := s.httpUpstream.DoWithTLS(budgetRetryReq, proxyURL, account.ID, account.Concurrency, tlsProfile) if retryErr == nil { + if budgetRetryResp.StatusCode < 400 { + // budget 修正请求成功后，ParsedRequest 也要描述被接受的修正版。 + lastWireBody = budgetWireBody + if err := replaceBody(budgetWireBody); err != nil { + _ = budgetRetryResp.Body.Close() + return nil, err + } + } resp = budgetRetryResp break } @@ -4786,7 +4900,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A break } - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, @@ -4833,7 +4947,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // 处理重试耗尽的情况 if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) { if s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -4868,7 +4982,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // 处理可切换账号的错误 if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -4900,7 +5014,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if resp.StatusCode >= 400 { // 可选：对部分 400 触发 failover（默认关闭以保持语义） if resp.StatusCode == 400 && s.cfg != nil && s.cfg.Gateway.FailoverOn400 { - respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, readErr := s.readUpstreamErrorBody(resp) if readErr != nil { // ReadAll failed, fall back to normal error handling without consuming the stream return s.handleErrorResponse(ctx, resp, c, account, reqModel) @@ -4948,6 +5062,13 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // 处理正常响应 + if !bytes.Equal(lastWireBody, body) { + // 成功后再同步最终 wire body，避免失败重试从已签名 CCH 的 body 继续派生。 + if err := replaceBody(lastWireBody); err != nil { + return nil, err + } + } + // 触发上游接受回调（提前释放串行锁，不等流完成） if parsed.OnUpstreamAccepted != nil { parsed.OnUpstreamAccepted() @@ -4990,6 +5111,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A type anthropicPassthroughForwardInput struct { Body []byte + Parsed *ParsedRequest RequestModel string OriginalModel string RequestStream bool @@ -5042,16 +5164,29 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( } // Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400. input.Body = StripEmptyTextBlocks(input.Body) + if input.Parsed != nil { + // 透传分支也会改写实际 wire body，成功 usage hash 依赖这里同步当前 body。 + if err := input.Parsed.ReplaceBody(input.Body); err != nil { + return nil, err + } + } var resp *http.Response retryStart := time.Now() for attempt := 1; attempt <= maxRetryAttempts; attempt++ { upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, input.RequestStream) - upstreamReq, err := s.buildUpstreamRequestAnthropicAPIKeyPassthrough(upstreamCtx, c, account, input.Body, token) + upstreamReq, wireBody, err := s.buildUpstreamRequestAnthropicAPIKeyPassthrough(upstreamCtx, c, account, input.Body, token) releaseUpstreamCtx() if err != nil { return nil, err } + if input.Parsed != nil && !bytes.Equal(wireBody, input.Body) { + // build 阶段会按 beta 能力清理 body，发送前同步到 ParsedRequest 当前视图。 + if err := input.Parsed.ReplaceBody(wireBody); err != nil { + return nil, err + } + input.Body = input.Parsed.Body.Bytes() + } resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) if err != nil { @@ -5097,7 +5232,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( break } - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, @@ -5135,7 +5270,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) { if s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -5169,7 +5304,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput( } if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -5243,13 +5378,13 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough( account *Account, body []byte, token string, -) (*http.Request, error) { +) (*http.Request, []byte, error) { targetURL := claudeAPIURL baseURL := account.GetBaseURL() if baseURL != "" { validatedURL, err := s.validateUpstreamBaseURL(baseURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = validatedURL + "/v1/messages?beta=true" } @@ -5267,7 +5402,7 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough( req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body)) if err != nil { - return nil, err + return nil, nil, err } if c != nil && c.Request != nil { @@ -5297,7 +5432,7 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough( setHeaderRaw(req.Header, "anthropic-version", "2023-06-01") } - return req, nil + return req, body, nil } func (s *GatewayService) handleStreamingResponseAnthropicAPIKeyPassthrough( @@ -5711,7 +5846,7 @@ func (s *GatewayService) forwardBedrock( ) (*ForwardResult, error) { reqModel := parsed.Model reqStream := parsed.Stream - body := parsed.Body + body := parsed.Body.Bytes() region := bedrockRuntimeRegion(account) mappedModel, ok := ResolveBedrockModelID(account, reqModel) @@ -5779,6 +5914,11 @@ func (s *GatewayService) forwardBedrock( return s.handleBedrockUpstreamErrors(ctx, resp, c, account) } + // Bedrock 分支绕过通用 Forward 成功路径，这里保持上游接受回调语义一致。 + if parsed.OnUpstreamAccepted != nil { + parsed.OnUpstreamAccepted() + } + // 响应处理 var usage *ClaudeUsage var firstTokenMs *int @@ -5882,7 +6022,7 @@ func (s *GatewayService) executeBedrockUpstream( break } - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, @@ -5927,7 +6067,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors( // retry exhausted + failover if s.shouldRetryUpstreamError(account, resp.StatusCode) { if s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -5954,7 +6094,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors( // non-retryable failover if s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -6055,9 +6195,10 @@ func (s *GatewayService) handleBedrockNonStreamingResponse( return usage, nil } -func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, error) { +func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, []byte, error) { if account.Platform == PlatformAnthropic && account.Type == AccountTypeServiceAccount { - return s.buildUpstreamRequestAnthropicVertex(ctx, c, account, body, token, modelID, reqStream) + req, err := s.buildUpstreamRequestAnthropicVertex(ctx, c, account, body, token, modelID, reqStream) + return req, body, err } // 确定目标URL @@ -6067,18 +6208,18 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex if baseURL != "" { validatedURL, err := s.validateUpstreamBaseURL(baseURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = validatedURL + "/v1/messages?beta=true" } } else if account.IsCustomBaseURLEnabled() { customURL := account.GetCustomBaseURL() if customURL == "" { - return nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) + return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) } validatedURL, err := s.validateUpstreamBaseURL(customURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages", account) } @@ -6153,7 +6294,7 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body)) if err != nil { - return nil, err + return nil, nil, err } // 设置认证头（保持原始大小写） @@ -6238,7 +6379,7 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode) } - return req, nil + return req, body, nil } func (s *GatewayService) buildUpstreamRequestAnthropicVertex( @@ -7096,8 +7237,19 @@ func isCountTokensUnsupported404(statusCode int, body []byte) bool { return strings.Contains(msg, "count_tokens") && strings.Contains(msg, "not found") } +func (s *GatewayService) readUpstreamErrorBody(resp *http.Response) ([]byte, error) { + if resp == nil || resp.Body == nil { + return nil, nil + } + limit := gatewayUpstreamErrorBodyReadLimit + if s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody && s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { + limit = int64(s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return io.ReadAll(io.LimitReader(resp.Body, limit)) +} + func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, requestedModel ...string) (*ForwardResult, error) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body, _ := s.readUpstreamErrorBody(resp) // 调试日志：打印上游错误响应 logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (non-retryable): Account=%d(%s) Status=%d RequestID=%s Body=%s", @@ -7250,7 +7402,7 @@ func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Res } func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body, _ := s.readUpstreamErrorBody(resp) statusCode := resp.StatusCode // OAuth/Setup Token 账号的 403：标记账号异常 @@ -7264,7 +7416,7 @@ func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, re } func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body, _ := s.readUpstreamErrorBody(resp) if len(requestedModel) > 0 { s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0]) return @@ -7277,7 +7429,7 @@ func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *ht // API Key 未配置错误码：仅返回错误，不标记账号 func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account) (*ForwardResult, error) { // Capture upstream error body before side-effects consume the stream. - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody, _ := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -8101,10 +8253,10 @@ func (s *GatewayService) getUserGroupRateMultiplier(ctx context.Context, userID, return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier) } -// RecordUsageInput 记录使用量的输入参数 +// RecordUsageInput 记录使用量的输入参数。 +// 异步 worker 只接收计费所需快照，不能持有 ParsedRequest/RequestBodyRef 这类大请求体引用。 type RecordUsageInput struct { Result *ForwardResult - ParsedRequest *ParsedRequest APIKey *APIKey User *User Account *Account @@ -8229,18 +8381,23 @@ func postUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *bill } } - // Platform quota DB-only 累加（与 finalizePostUsageBilling 行为对齐的兜底）： - // - 仅对 standard（余额）模式生效；订阅模式豁免 - // - 直接走 DB，不经 Redis Incr 队列：legacy 路径在 repo==nil（仓库未注入） - // 时被触发，此时整套 billing repo 都不可用，没有"双队列"风险 - // - 失败仅记 ALERT log + counter，不阻断主扣费流程；与正常路径一致 - // - // 历史背景：原 legacy path 完全跳过此累加，导致部署中如果 repo 偶然为 nil - // 时用户消费可绕过 platform quota，存在静默资金风险。 + // Platform quota 累加（legacy 兜底路径）：仅对 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写 + // - HasUserPlatformQuotaLimit 守卫:与正常路径对齐，无 limit 公司跳过 + // - 新增 Redis 同步写:enforcement 走 Redis，legacy 路径也必须同步写，否则 preflight 看不到消费 + // - flusher_enabled=false（降级）:保留原有同步直写 DB + // - flusher_enabled=true:跳过直写 DB，由 flusher 异步批量刷（markDirty 在 IncrementUserPlatformQuotaUsage 内部完成） + // - 失败仅记 ALERT log + counter，不阻断主扣费流程 if !p.IsSubscriptionBill && p.Platform != "" && cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil { - if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(billingCtx, p.User.ID, p.Platform, cost.ActualCost, time.Now().UTC()); err != nil { - userPlatformQuotaDBIncrLegacyErrorTotal.Add(1) - logger.LegacyPrintf("service.gateway", "ALERT: legacy incr user platform quota DB failed user=%d platform=%s cost=%f: %v", p.User.ID, p.Platform, cost.ActualCost, err) + if deps.billingCacheService.HasUserPlatformQuotaLimit(billingCtx, p.User.ID, p.Platform) { + deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, cost.ActualCost) + if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled { + // 降级路径:flusher 未启用时保留原有同步直写 DB + if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(billingCtx, p.User.ID, p.Platform, cost.ActualCost, time.Now().UTC()); err != nil { + userPlatformQuotaDBIncrLegacyErrorTotal.Add(1) + logger.LegacyPrintf("service.gateway", "ALERT: legacy incr user platform quota DB failed user=%d platform=%s cost=%f: %v", p.User.ID, p.Platform, cost.ActualCost, err) + } + } + // flusher_enabled=true:不直写 DB，flusher 异步批量刷 } } @@ -8390,30 +8547,38 @@ func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, de deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID) - // Platform quota 累加：仅在 standard（余额）模式生效；订阅模式豁免 - // Redis 同步写 + DB 异步持久化: + // Platform quota 累加：仅在 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写 + // Redis 同步写 + DB 异步持久化（flag=false 降级）或 flusher 异步刷（flag=true）: + // - HasUserPlatformQuotaLimit 守卫:无 limit 的公司跳过,避免无效写入 + 浪费 Redis 容量 // - Redis 同步:确保下次 preflight 立即看到最新 usage,把 TOCTOU 超支窗口 // 限制在并发 in-flight 请求数量内（旧实现的异步入队会让超支无限累积直到 worker 处理） - // - DB 异步:在独立 goroutine 中走 detached context,失败用 ALERT log 触发 oncall 对账 + // - DB 异步(flusher_enabled=false):在独立 goroutine 中走 detached context,失败用 ALERT log 触发 oncall 对账 + // - flusher_enabled=true:不直写 DB,由 flusher 异步批量刷（markDirty 已在 IncrementUserPlatformQuotaUsage 内部完成） if !p.IsSubscriptionBill && p.Platform != "" && p.Cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil { - deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, p.Cost.ActualCost) - dbCtx, dbCancel := detachUpstreamContext(ctx) - userID, platform, cost := p.User.ID, p.Platform, p.Cost.ActualCost - go func() { - defer func() { - if r := recover(); r != nil { - logger.LegacyPrintf("service.gateway", "ALERT: panic in user platform quota incr goroutine user=%d platform=%s: %v", userID, platform, r) - } - }() - defer dbCancel() - if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(dbCtx, userID, platform, cost, time.Now().UTC()); err != nil { - // 失败计数器:暴露给 GatewayUserPlatformQuotaIncrStats(),由 ops 面板做斜率告警。 - userPlatformQuotaDBIncrErrorTotal.Add(1) - // ALERT 级别:DB 持久化失败意味着 Redis cache 失效后该笔 cost 永久丢失, - // 用户配额视图与实际消费会偏差,oncall 需要据此对账或人工补录。 - logger.LegacyPrintf("service.gateway", "ALERT: incr user platform quota DB failed user=%d platform=%s cost=%f: %v", userID, platform, cost, err) + if deps.billingCacheService.HasUserPlatformQuotaLimit(ctx, p.User.ID, p.Platform) { + deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, p.Cost.ActualCost) + if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled { + // 降级路径:flusher 未启用时保留原有异步直写 DB + dbCtx, dbCancel := detachUpstreamContext(ctx) + userID, platform, cost := p.User.ID, p.Platform, p.Cost.ActualCost + go func() { + defer func() { + if r := recover(); r != nil { + logger.LegacyPrintf("service.gateway", "ALERT: panic in user platform quota incr goroutine user=%d platform=%s: %v", userID, platform, r) + } + }() + defer dbCancel() + if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(dbCtx, userID, platform, cost, time.Now().UTC()); err != nil { + // 失败计数器:暴露给 GatewayUserPlatformQuotaIncrStats(),由 ops 面板做斜率告警。 + userPlatformQuotaDBIncrErrorTotal.Add(1) + // ALERT 级别:DB 持久化失败意味着 Redis cache 失效后该笔 cost 永久丢失, + // 用户配额视图与实际消费会偏差,oncall 需要据此对账或人工补录。 + logger.LegacyPrintf("service.gateway", "ALERT: incr user platform quota DB failed user=%d platform=%s cost=%f: %v", userID, platform, cost, err) + } + }() } - }() + // flusher_enabled=true:不直写 DB,flusher 异步批量刷 + } } // Notification checks run async — all parameters are already captured, @@ -8528,6 +8693,7 @@ type billingDeps struct { deferredService *DeferredService balanceNotifyService *BalanceNotifyService userPlatformQuotaRepo UserPlatformQuotaRepository + cfg *config.Config } func (s *GatewayService) billingDeps() *billingDeps { @@ -8539,6 +8705,7 @@ func (s *GatewayService) billingDeps() *billingDeps { deferredService: s.deferredService, balanceNotifyService: s.balanceNotifyService, userPlatformQuotaRepo: s.userPlatformQuotaRepo, + cfg: s.cfg, } } @@ -8567,15 +8734,8 @@ func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usage } } -// recordUsageOpts 内部选项，参数化 RecordUsage 与 RecordUsageWithLongContext 的差异点。 +// recordUsageOpts 内部选项，参数化普通计费与长上下文计费的差异点。 type recordUsageOpts struct { - // Claude Max 策略所需的 ParsedRequest（可选，仅 Claude 路径传入） - ParsedRequest *ParsedRequest - - // EnableClaudePath 启用 Claude 路径特有逻辑： - // - Claude Max 缓存计费策略 - EnableClaudePath bool - // 长上下文计费（仅 Gemini 路径需要） LongContextThreshold int LongContextMultiplier float64 @@ -8598,9 +8758,7 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu APIKeyService: input.APIKeyService, QuotaPlatform: input.QuotaPlatform, ChannelUsageFields: input.ChannelUsageFields, - }, &recordUsageOpts{ - EnableClaudePath: true, - }) + }, &recordUsageOpts{}) } // RecordUsageLongContextInput 记录使用量的输入参数（支持长上下文双倍计费） @@ -8666,9 +8824,7 @@ type recordUsageCoreInput struct { } // recordUsageCore 是 RecordUsage 和 RecordUsageWithLongContext 的统一实现。 -// opts 中的字段控制两者之间的差异行为： -// - ParsedRequest != nil → 启用 Claude Max 缓存计费策略 -// - LongContextThreshold > 0 → Token 计费回退走 CalculateCostWithLongContext +// LongContextThreshold > 0 时 Token 计费回退走 CalculateCostWithLongContext。 func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsageCoreInput, opts *recordUsageOpts) error { result := input.Result apiKey := input.APIKey @@ -9133,7 +9289,7 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, } if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() { - passthroughBody := parsed.Body + passthroughBody := parsed.Body.Bytes() if reqModel := parsed.Model; reqModel != "" { if mappedModel := account.GetMappedModel(reqModel); mappedModel != reqModel { passthroughBody = s.replaceModelInBody(passthroughBody, mappedModel) @@ -9149,24 +9305,43 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, return nil } - body := parsed.Body + body := parsed.Body.Bytes() + replaceBody := func(next []byte) error { + if err := parsed.ReplaceBody(next); err != nil { + return fmt.Errorf("rewrite count_tokens body: %w", err) + } + body = parsed.Body.Bytes() + return nil + } reqModel := parsed.Model // Pre-filter: strip empty text blocks to prevent upstream 400. - body = StripEmptyTextBlocks(body) + if err := replaceBody(StripEmptyTextBlocks(body)); err != nil { + return err + } isClaudeCodeCT := IsClaudeCodeClient(ctx) || isClaudeCodeClient(c.GetHeader("User-Agent"), parsed.MetadataUserID) shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCodeCT if shouldMimicClaudeCode { normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: true} - body, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + var normalizedBody []byte + normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + if err := replaceBody(normalizedBody); err != nil { + return err + } - body = s.rewriteMessageCacheControlIfEnabled(ctx, body) + if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil { + return err + } if rw := buildToolNameRewriteFromBody(body); rw != nil { - body = applyToolNameRewriteToBody(body, rw) + if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil { + return err + } } else { - body = applyToolsLastCacheBreakpoint(body) + if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil { + return err + } } } @@ -9197,9 +9372,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, } } if mappedModel != reqModel { - body = s.replaceModelInBody(body, mappedModel) + originalReqModel := reqModel + if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil { + return err + } reqModel = mappedModel - logger.LegacyPrintf("service.gateway", "CountTokens model mapping applied: %s -> %s (account: %s, source=%s)", parsed.Model, mappedModel, account.Name, mappingSource) + parsed.Model = mappedModel + logger.LegacyPrintf("service.gateway", "CountTokens model mapping applied: %s -> %s (account: %s, source=%s)", originalReqModel, mappedModel, account.Name, mappingSource) } } @@ -9211,11 +9390,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, } // 构建上游请求 - upstreamReq, err := s.buildCountTokensRequest(ctx, c, account, body, token, tokenType, reqModel, shouldMimicClaudeCode) + upstreamReq, wireBody, err := s.buildCountTokensRequest(ctx, c, account, body, token, tokenType, reqModel, shouldMimicClaudeCode) if err != nil { s.countTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request") return err } + // 先记录首发 wire body；如果后面进入 400 retry，retry 会基于未签名的逻辑 body 重新构建。 + acceptedWireBody := wireBody // 获取代理URL（自定义 base URL 模式下，proxy 通过 buildCustomRelayURL 作为查询参数传递） proxyURL := "" @@ -9251,10 +9432,14 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, logger.LegacyPrintf("service.gateway", "Account %d: detected thinking block signature error on count_tokens, retrying with filtered thinking blocks", account.ID) filteredBody := FilterThinkingBlocksForRetry(body) - retryReq, buildErr := s.buildCountTokensRequest(ctx, c, account, filteredBody, token, tokenType, reqModel, shouldMimicClaudeCode) + retryReq, retryWireBody, buildErr := s.buildCountTokensRequest(ctx, c, account, filteredBody, token, tokenType, reqModel, shouldMimicClaudeCode) if buildErr == nil { retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) if retryErr == nil { + if retryResp.StatusCode < 400 { + // count_tokens 签名重试成功后记录最终 wire body，错误响应仍保留原 body 便于后续处理。 + acceptedWireBody = retryWireBody + } resp = retryResp respBody, err = ReadUpstreamResponseBody(resp.Body, s.cfg, c, countTokensTooLarge) _ = resp.Body.Close() @@ -9268,6 +9453,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, } } + if resp.StatusCode < 400 && !bytes.Equal(acceptedWireBody, body) { + // count_tokens 成功后再同步最终 wire body，避免 retry 从已签名 body 派生。 + if err := replaceBody(acceptedWireBody); err != nil { + return err + } + } + // 处理错误响应 if resp.StatusCode >= 400 { // 标记账号状态（429/529等） @@ -9448,6 +9640,7 @@ func (s *GatewayService) buildCountTokensRequestAnthropicAPIKeyPassthrough( } targetURL = validatedURL + "/v1/messages/count_tokens?beta=true" } + body = sanitizeCountTokensRequestBody(body) // 同 buildUpstreamRequestAnthropicAPIKeyPassthrough：能力维度 sanitize。 clientBeta := "" @@ -9493,7 +9686,7 @@ func (s *GatewayService) buildCountTokensRequestAnthropicAPIKeyPassthrough( } // buildCountTokensRequest 构建 count_tokens 上游请求 -func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, mimicClaudeCode bool) (*http.Request, error) { +func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, mimicClaudeCode bool) (*http.Request, []byte, error) { // 确定目标 URL targetURL := claudeAPICountTokensURL if account.Type == AccountTypeAPIKey { @@ -9501,18 +9694,18 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con if baseURL != "" { validatedURL, err := s.validateUpstreamBaseURL(baseURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = validatedURL + "/v1/messages/count_tokens?beta=true" } } else if account.IsCustomBaseURLEnabled() { customURL := account.GetCustomBaseURL() if customURL == "" { - return nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) + return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) } validatedURL, err := s.validateUpstreamBaseURL(customURL) if err != nil { - return nil, err + return nil, nil, err } targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages/count_tokens", account) } @@ -9564,10 +9757,11 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con if ctEnableCCH { body = signBillingHeaderCCH(body) } + body = sanitizeCountTokensRequestBody(body) req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body)) if err != nil { - return nil, err + return nil, nil, err } // 设置认证头（保持原始大小写） @@ -9631,7 +9825,26 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode) } - return req, nil + return req, body, nil +} + +func sanitizeCountTokensRequestBody(body []byte) []byte { + out := body + for _, path := range []string{ + "temperature", + "top_p", + "top_k", + "stream", + "stop_sequences", + "stop", + } { + if gjson.GetBytes(out, path).Exists() { + if next, ok := deleteJSONPathBytes(out, path); ok { + out = next + } + } + } + return out } // countTokensError 返回 count_tokens 错误响应 diff --git a/backend/internal/service/gateway_service_benchmark_test.go b/backend/internal/service/gateway_service_benchmark_test.go index c9c4d3dd8bab..42a711db9b62 100644 --- a/backend/internal/service/gateway_service_benchmark_test.go +++ b/backend/internal/service/gateway_service_benchmark_test.go @@ -2,10 +2,16 @@ package service import ( "strconv" + "strings" "testing" + + "github.com/Wei-Shaw/sub2api/internal/domain" ) -var benchmarkStringSink string +var ( + benchmarkStringSink string + benchmarkIntSink int +) // BenchmarkGenerateSessionHash_Metadata 关注 JSON 解析与正则匹配开销。 func BenchmarkGenerateSessionHash_Metadata(b *testing.B) { @@ -14,7 +20,7 @@ func BenchmarkGenerateSessionHash_Metadata(b *testing.B) { b.ReportAllocs() for i := 0; i < b.N; i++ { - parsed, err := ParseGatewayRequest(body, "") + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "") if err != nil { b.Fatalf("解析请求失败: %v", err) } @@ -22,6 +28,179 @@ func BenchmarkGenerateSessionHash_Metadata(b *testing.B) { } } +func BenchmarkParseGatewayRequest_LargeAnthropicMessages(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeAnthropicMessagesBody(size.bytes, false) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic) + if err != nil { + b.Fatalf("解析 Anthropic 请求失败: %v", err) + } + benchmarkIntSink = len(parsed.MessagesRaw()) + } + }) + } +} + +func BenchmarkParseGatewayRequest_LargeGeminiContents(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeGeminiContentsBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini) + if err != nil { + b.Fatalf("解析 Gemini 请求失败: %v", err) + } + benchmarkIntSink = len(parsed.MessagesRaw()) + } + }) + } +} + +func BenchmarkGenerateSessionHash_LargeAnthropicMessages(b *testing.B) { + svc := &GatewayService{} + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeAnthropicMessagesBody(size.bytes, true) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic) + if err != nil { + b.Fatalf("解析请求失败: %v", err) + } + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + benchmarkStringSink = svc.GenerateSessionHash(parsed) + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputMeta(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + model, stream, promptCacheKey := extractOpenAIRequestMetaFromBody(body) + benchmarkStringSink = model + promptCacheKey + if stream { + benchmarkIntSink++ + } + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputDecodeMap(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + reqBody, err := getOpenAIRequestBodyMap(nil, body) + if err != nil { + b.Fatalf("解析 OpenAI 请求失败: %v", err) + } + benchmarkIntSink = len(reqBody) + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputRawPatch(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + view := newOpenAIRequestView(body) + view.MarkPatchSet("instructions", "You are a helpful coding assistant.") + view.MarkPatchSet("reasoning.effort", "none") + patched, err := view.ApplyPatches() + if err != nil { + b.Fatalf("应用 OpenAI raw patch 失败: %v", err) + } + benchmarkIntSink = len(patched) + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputImageBillingRaw(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesImageToolBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + cfg, err := resolveOpenAIResponsesImageBillingConfigDetailedFromBody(body, "gpt-5.4") + if err != nil { + b.Fatalf("解析 OpenAI 图片计费配置失败: %v", err) + } + benchmarkStringSink = cfg.Model + cfg.SizeTier + cfg.InputSize + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputEmptyBase64Guard(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if openAIRequestBodyMayContainEmptyBase64InputImage(body) { + benchmarkIntSink++ + } + } + }) + } +} + +func BenchmarkOpenAIResponses_LargeInputFunctionCallValidation(b *testing.B) { + for _, size := range benchmarkBodySizes() { + b.Run(size.name, func(b *testing.B) { + body := buildLargeOpenAIResponsesToolContinuationBody(size.bytes) + + b.SetBytes(int64(len(body))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + validation := ValidateFunctionCallOutputContextBytes(body) + if !validation.HasFunctionCallOutput || !validation.HasItemReferenceForAllCallIDs { + b.Fatalf("工具续链校验结果异常: %+v", validation) + } + benchmarkIntSink++ + } + }) + } +} + // BenchmarkExtractCacheableContent_System 关注字符串拼接路径的性能。 func BenchmarkExtractCacheableContent_System(b *testing.B) { svc := &GatewayService{} @@ -33,18 +212,130 @@ func BenchmarkExtractCacheableContent_System(b *testing.B) { } } +func benchmarkBodySizes() []struct { + name string + bytes int +} { + return []struct { + name string + bytes int + }{ + {name: "4MB", bytes: 4 << 20}, + {name: "8MB", bytes: 8 << 20}, + {name: "16MB", bytes: 16 << 20}, + {name: "32MB", bytes: 32 << 20}, + } +} + func buildSystemCacheableRequest(parts int) *ParsedRequest { - systemParts := make([]any, 0, parts) + var builder strings.Builder + _, _ = builder.WriteString(`{"system":[`) for i := 0; i < parts; i++ { - systemParts = append(systemParts, map[string]any{ - "text": "system_part_" + strconv.Itoa(i), - "cache_control": map[string]any{ - "type": "ephemeral", - }, - }) + if i > 0 { + _ = builder.WriteByte(',') + } + _, _ = builder.WriteString(`{"text":"system_part_`) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`","cache_control":{"type":"ephemeral"}}`) + } + _, _ = builder.WriteString(`]}`) + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(builder.String())), "") + if err != nil { + panic(err) } - return &ParsedRequest{ - System: systemParts, - HasSystem: true, + return parsed +} + +func buildLargeAnthropicMessagesBody(targetBytes int, includeCacheControl bool) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + _, _ = builder.WriteString(`{"model":"claude-sonnet-4-5","stream":true,"system":[{"type":"text","text":"system seed"}],"messages":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + _ = builder.WriteByte(',') + } + _, _ = builder.WriteString(`{"role":"user","content":[{"type":"text","text":"`) + _, _ = builder.WriteString(strings.Repeat("anthropic payload ", 64)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _ = builder.WriteByte('"') + if includeCacheControl && i%32 == 0 { + _, _ = builder.WriteString(`,"cache_control":{"type":"ephemeral"}`) + } + _, _ = builder.WriteString(`}]}`) + } + _, _ = builder.WriteString(`]}`) + return []byte(builder.String()) +} + +func buildLargeGeminiContentsBody(targetBytes int) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + _, _ = builder.WriteString(`{"model":"gemini-2.5-pro","systemInstruction":{"parts":[{"text":"system seed"}]},"contents":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + _ = builder.WriteByte(',') + } + _, _ = builder.WriteString(`{"role":"user","parts":[{"text":"`) + _, _ = builder.WriteString(strings.Repeat("gemini payload ", 64)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`"}]}`) + } + _, _ = builder.WriteString(`]}`) + return []byte(builder.String()) +} + +func buildLargeOpenAIResponsesBody(targetBytes int) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + _, _ = builder.WriteString(`{"model":"gpt-5.4","stream":true,"prompt_cache_key":"session-benchmark","input":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + _ = builder.WriteByte(',') + } + _, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) + _, _ = builder.WriteString(strings.Repeat("openai responses payload ", 48)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`"}]}`) + } + _, _ = builder.WriteString(`],"tools":[{"type":"function","name":"lookup","parameters":{"type":"object","properties":{"query":{"type":"string"}}}}]}`) + return []byte(builder.String()) +} + +func buildLargeOpenAIResponsesToolContinuationBody(targetBytes int) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + _, _ = builder.WriteString(`{"model":"gpt-5.4","stream":true,"previous_response_id":"resp_benchmark","input":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + _ = builder.WriteByte(',') + } + callID := "call_" + strconv.Itoa(i) + _, _ = builder.WriteString(`{"type":"item_reference","id":"`) + _, _ = builder.WriteString(callID) + _, _ = builder.WriteString(`"},{"type":"function_call_output","call_id":"`) + _, _ = builder.WriteString(callID) + _, _ = builder.WriteString(`","output":"`) + _, _ = builder.WriteString(strings.Repeat("tool output payload ", 48)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`"}`) + } + _, _ = builder.WriteString(`]}`) + return []byte(builder.String()) +} + +func buildLargeOpenAIResponsesImageToolBody(targetBytes int) []byte { + var builder strings.Builder + builder.Grow(targetBytes + 1024) + _, _ = builder.WriteString(`{"model":"gpt-5.4","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[`) + for i := 0; builder.Len() < targetBytes; i++ { + if i > 0 { + _ = builder.WriteByte(',') + } + _, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) + _, _ = builder.WriteString(strings.Repeat("openai image billing payload ", 48)) + _, _ = builder.WriteString(strconv.Itoa(i)) + _, _ = builder.WriteString(`"}]}`) } + _, _ = builder.WriteString(`]}`) + return []byte(builder.String()) } diff --git a/backend/internal/service/gateway_websearch_emulation.go b/backend/internal/service/gateway_websearch_emulation.go index a42b5585e323..2f9c8e0c258b 100644 --- a/backend/internal/service/gateway_websearch_emulation.go +++ b/backend/internal/service/gateway_websearch_emulation.go @@ -150,7 +150,7 @@ func (s *GatewayService) handleWebSearchEmulation( parsed.OnUpstreamAccepted() } - query := extractSearchQueryFromBody(parsed.Body) + query := extractSearchQueryFromBody(parsed.Body.Bytes()) if query == "" { return nil, fmt.Errorf("web search emulation: no query found in messages") } diff --git a/backend/internal/service/gemini_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index dcc3213b1b69..ffea15950335 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -151,7 +151,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( } if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() if resp.StatusCode == http.StatusForbidden && isGeminiInsufficientScope(resp.Header, respBody) { resp = &http.Response{ @@ -207,7 +207,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( reasoningEffort := extractCCReasoningEffortFromBody(originalChatBody) if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) evBody := unwrapIfNeeded(account.Type == AccountTypeOAuth, respBody) diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 516556ca7de8..86073d9ca549 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -56,6 +56,18 @@ type GeminiMessagesCompatService struct { responseHeaderFilter *responseheaders.CompiledHeaderFilter } +func (s *GeminiMessagesCompatService) readUpstreamErrorBody(resp *http.Response) []byte { + if resp == nil || resp.Body == nil { + return nil + } + limit := gatewayUpstreamErrorBodyReadLimit + if s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody && s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { + limit = int64(s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + body, _ := io.ReadAll(io.LimitReader(resp.Body, limit)) + return body +} + func NewGeminiMessagesCompatService( accountRepo AccountRepository, groupRepo GroupRepository, @@ -789,7 +801,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex // Special-case: signature/thought_signature validation errors are not transient, but may be fixed by // downgrading Claude thinking/tool history to plain text (conservative two-stage retry). if resp.StatusCode == http.StatusBadRequest && signatureRetryStage < 2 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() if isGeminiSignatureRelatedError(respBody) { @@ -860,7 +872,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex } if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() // Don't treat insufficient-scope as transient. if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) { @@ -919,7 +931,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex defer func() { _ = resp.Body.Close() }() if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) // 统一错误策略：自定义错误码 + 临时不可调度 if s.rateLimitService != nil { switch s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody) { @@ -1329,7 +1341,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. } if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() // Don't treat insufficient-scope as transient. if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) { @@ -1410,7 +1422,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. isOAuth := account.Type == AccountTypeOAuth if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) // Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens. // This avoids Gemini SDKs failing hard during preflight token counting. // Checked before error policy so it always works regardless of custom error codes. @@ -1619,7 +1631,7 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop( if resp.StatusCode < 400 || s.rateLimitService == nil { return false, resp } - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() rebuilt = &http.Response{ StatusCode: resp.StatusCode, @@ -2031,6 +2043,22 @@ func (s *GeminiMessagesCompatService) handleStreamingResponse(c *gin.Context, re parts := extractGeminiParts(geminiResp) for _, part := range parts { if text, ok := part["text"].(string); ok && text != "" { + // Close an open tool_use block before starting text, mirroring + // the functionCall branch (which closes open text blocks) and + // the chat-completions sibling's closeOpenTool(). Otherwise a + // tool→text sequence keeps the tool_use block open while the + // text block starts, emitting overlapping Anthropic content + // blocks that violate the SSE contract. + if openToolIndex >= 0 { + writeSSE(c.Writer, "content_block_stop", map[string]any{ + "type": "content_block_stop", + "index": openToolIndex, + }) + openToolIndex = -1 + openToolName = "" + seenToolJSON = "" + } + delta, newSeen := computeGeminiTextDelta(seenText, text) seenText = newSeen if delta == "" { diff --git a/backend/internal/service/gemini_messages_compat_service_test.go b/backend/internal/service/gemini_messages_compat_service_test.go index d0560344dc47..79db633aae0f 100644 --- a/backend/internal/service/gemini_messages_compat_service_test.go +++ b/backend/internal/service/gemini_messages_compat_service_test.go @@ -832,3 +832,108 @@ func TestParseGeminiRateLimitResetTime(t *testing.T) { }) } } + +// TestGeminiMessagesHandleStreamingResponse_ClosesToolBlockBeforeText guards the +// tool→text ordering in the Gemini→Anthropic (messages) streaming bridge. When +// Gemini emits a functionCall part followed by a text part, the tool_use content +// block must be closed before the text block opens; otherwise the Anthropic SSE +// stream contains overlapping content blocks. The chat-completions sibling +// already enforces this via closeOpenTool(). +func TestGeminiMessagesHandleStreamingResponse_ClosesToolBlockBeforeText(t *testing.T) { + gin.SetMode(gin.TestMode) + + upstreamBody := `data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"get_weather","args":{"city":"SF"}}}]}}]}` + "\n\n" + + `data: {"candidates":[{"content":{"parts":[{"text":"All done."}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3}}` + "\n\n" + + "data: [DONE]\n\n" + + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstreamBody)), + } + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + + svc := &GeminiMessagesCompatService{} + result, err := svc.handleStreamingResponse(c, resp, time.Now(), "claude-3-5-sonnet") + require.NoError(t, err) + require.NotNil(t, result) + + events := parseAnthropicContentBlockEvents(t, rec.Body.String()) + + // Anthropic allows at most one content block open at a time: every + // content_block_start must be matched by a content_block_stop before the + // next start. Replay the lifecycle and assert there is no overlap. + open := -1 + blockTypes := map[int]string{} + textStarted := false + toolClosed := false + toolClosedBeforeText := false + for _, ev := range events { + switch ev.event { + case "content_block_start": + require.Equalf(t, -1, open, + "content block %d opened while block %d was still open (overlapping blocks)", ev.index, open) + open = ev.index + blockTypes[ev.index] = ev.blockType + if ev.blockType == "text" { + textStarted = true + if toolClosed { + toolClosedBeforeText = true + } + } + case "content_block_stop": + require.Equalf(t, open, ev.index, + "content_block_stop index %d does not match the open block %d", ev.index, open) + if blockTypes[ev.index] == "tool_use" { + toolClosed = true + } + open = -1 + } + } + + require.True(t, textStarted, "expected a text content block to be emitted after the tool call") + require.True(t, toolClosedBeforeText, "tool_use block must be closed before the text block starts") + require.Equal(t, -1, open, "stream ended with a content block still open") +} + +type anthropicContentBlockEvent struct { + event string + index int + blockType string +} + +// parseAnthropicContentBlockEvents extracts content_block_start/stop events (with +// their index and, for starts, the content block type) from an Anthropic SSE body. +func parseAnthropicContentBlockEvents(t *testing.T, raw string) []anthropicContentBlockEvent { + t.Helper() + var events []anthropicContentBlockEvent + for _, chunk := range strings.Split(raw, "\n\n") { + var eventName, dataLine string + for _, line := range strings.Split(chunk, "\n") { + switch { + case strings.HasPrefix(line, "event:"): + eventName = strings.TrimSpace(strings.TrimPrefix(line, "event:")) + case strings.HasPrefix(line, "data:"): + dataLine = strings.TrimSpace(strings.TrimPrefix(line, "data:")) + } + } + if eventName != "content_block_start" && eventName != "content_block_stop" { + continue + } + var payload struct { + Index int `json:"index"` + ContentBlock struct { + Type string `json:"type"` + } `json:"content_block"` + } + require.NoError(t, json.Unmarshal([]byte(dataLine), &payload)) + events = append(events, anthropicContentBlockEvent{ + event: eventName, + index: payload.Index, + blockType: payload.ContentBlock.Type, + }) + } + return events +} diff --git a/backend/internal/service/generate_session_hash_test.go b/backend/internal/service/generate_session_hash_test.go index 39679c3da9cf..5ed3f0ae99d5 100644 --- a/backend/internal/service/generate_session_hash_test.go +++ b/backend/internal/service/generate_session_hash_test.go @@ -3,12 +3,67 @@ package service import ( + "encoding/json" "testing" + "github.com/Wei-Shaw/sub2api/internal/domain" "github.com/stretchr/testify/require" ) -// ============ 基础优先级测试 ============ +func mustParseSessionHashRequest(t *testing.T, body string, ctx *SessionContext) *ParsedRequest { + t.Helper() + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(body)), domain.PlatformAnthropic) + require.NoError(t, err) + parsed.SessionContext = ctx + return parsed +} + +func mustParseGeminiSessionHashRequest(t *testing.T, body string, ctx *SessionContext) *ParsedRequest { + t.Helper() + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(body)), domain.PlatformGemini) + require.NoError(t, err) + parsed.SessionContext = ctx + return parsed +} + +func anthropicSessionBody(system any, messages []any, metadataUserID string) string { + body := map[string]any{} + if system != nil { + body["system"] = system + } + if messages != nil { + body["messages"] = messages + } + if metadataUserID != "" { + body["metadata"] = map[string]any{"user_id": metadataUserID} + } + data, _ := json.Marshal(body) + return string(data) +} + +func geminiSessionBody(systemParts []any, contents []any) string { + body := map[string]any{} + if systemParts != nil { + body["systemInstruction"] = map[string]any{"parts": systemParts} + } + if contents != nil { + body["contents"] = contents + } + data, _ := json.Marshal(body) + return string(data) +} + +func msg(role string, content any) map[string]any { + return map[string]any{"role": role, "content": content} +} + +func geminiMsg(role string, texts ...string) map[string]any { + parts := make([]any, 0, len(texts)) + for _, text := range texts { + parts = append(parts, map[string]any{"text": text}) + } + return map[string]any{"role": role, "parts": parts} +} func TestGenerateSessionHash_NilParsedRequest(t *testing.T) { svc := &GatewayService{} @@ -22,37 +77,17 @@ func TestGenerateSessionHash_EmptyRequest(t *testing.T) { func TestGenerateSessionHash_MetadataHasHighestPriority(t *testing.T) { svc := &GatewayService{} - - parsed := &ParsedRequest{ - MetadataUserID: "user_a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2_account__session_123e4567-e89b-12d3-a456-426614174000", - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + metadata := "user_a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2_account__session_123e4567-e89b-12d3-a456-426614174000" + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, metadata), nil) hash := svc.GenerateSessionHash(parsed) require.Equal(t, "123e4567-e89b-12d3-a456-426614174000", hash, "metadata session_id should have highest priority") } -// ============ System + Messages 基础测试 ============ - func TestGenerateSessionHash_SystemPlusMessages(t *testing.T) { svc := &GatewayService{} - - withSystem := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } - withoutSystem := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + withSystem := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, ""), nil) + withoutSystem := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "hello")}, ""), nil) h1 := svc.GenerateSessionHash(withSystem) h2 := svc.GenerateSessionHash(withoutSystem) @@ -63,32 +98,16 @@ func TestGenerateSessionHash_SystemPlusMessages(t *testing.T) { func TestGenerateSessionHash_SystemOnlyProducesHash(t *testing.T) { svc := &GatewayService{} + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", nil, ""), nil) - parsed := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - } hash := svc.GenerateSessionHash(parsed) require.NotEmpty(t, hash, "system prompt alone should produce a hash as part of full digest") } func TestGenerateSessionHash_DifferentSystemsSameMessages(t *testing.T) { svc := &GatewayService{} - - parsed1 := &ParsedRequest{ - System: "You are assistant A.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } - parsed2 := &ParsedRequest{ - System: "You are assistant B.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + parsed1 := mustParseSessionHashRequest(t, anthropicSessionBody("You are assistant A.", []any{msg("user", "hello")}, ""), nil) + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody("You are assistant B.", []any{msg("user", "hello")}, ""), nil) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) @@ -97,16 +116,8 @@ func TestGenerateSessionHash_DifferentSystemsSameMessages(t *testing.T) { func TestGenerateSessionHash_SameSystemSameMessages(t *testing.T) { svc := &GatewayService{} - mk := func() *ParsedRequest { - return &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "hi"}, - }, - } + return mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello"), msg("assistant", "hi")}, ""), nil) } h1 := svc.GenerateSessionHash(mk()) @@ -116,53 +127,19 @@ func TestGenerateSessionHash_SameSystemSameMessages(t *testing.T) { func TestGenerateSessionHash_DifferentMessagesProduceDifferentHash(t *testing.T) { svc := &GatewayService{} - - parsed1 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "help me with Go"}, - }, - } - parsed2 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "help me with Python"}, - }, - } + parsed1 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "help me with Go")}, ""), nil) + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "help me with Python")}, ""), nil) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h1, h2, "same system but different messages should produce different hashes") } -// ============ SessionContext 核心测试 ============ - func TestGenerateSessionHash_DifferentSessionContextProducesDifferentHash(t *testing.T) { svc := &GatewayService{} - - // 相同消息 + 不同 SessionContext → 不同 hash（解决碰撞问题的核心场景） - parsed1 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "192.168.1.1", - UserAgent: "Mozilla/5.0", - APIKeyID: 100, - }, - } - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "10.0.0.1", - UserAgent: "curl/7.0", - APIKeyID: 200, - }, - } + body := anthropicSessionBody(nil, []any{msg("user", "hello")}, "") + parsed1 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "192.168.1.1", UserAgent: "Mozilla/5.0", APIKeyID: 100}) + parsed2 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "10.0.0.1", UserAgent: "curl/7.0", APIKeyID: 200}) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) @@ -173,19 +150,9 @@ func TestGenerateSessionHash_DifferentSessionContextProducesDifferentHash(t *tes func TestGenerateSessionHash_SameSessionContextProducesSameHash(t *testing.T) { svc := &GatewayService{} - - mk := func() *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "192.168.1.1", - UserAgent: "Mozilla/5.0", - APIKeyID: 100, - }, - } - } + ctx := &SessionContext{ClientIP: "192.168.1.1", UserAgent: "Mozilla/5.0", APIKeyID: 100} + body := anthropicSessionBody(nil, []any{msg("user", "hello")}, "") + mk := func() *ParsedRequest { return mustParseSessionHashRequest(t, body, ctx) } h1 := svc.GenerateSessionHash(mk()) h2 := svc.GenerateSessionHash(mk()) @@ -194,35 +161,17 @@ func TestGenerateSessionHash_SameSessionContextProducesSameHash(t *testing.T) { func TestGenerateSessionHash_MetadataOverridesSessionContext(t *testing.T) { svc := &GatewayService{} - - parsed := &ParsedRequest{ - MetadataUserID: "user_a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2_account__session_123e4567-e89b-12d3-a456-426614174000", - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "192.168.1.1", - UserAgent: "Mozilla/5.0", - APIKeyID: 100, - }, - } + metadata := "user_a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2_account__session_123e4567-e89b-12d3-a456-426614174000" + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "hello")}, metadata), &SessionContext{ClientIP: "192.168.1.1", UserAgent: "Mozilla/5.0", APIKeyID: 100}) hash := svc.GenerateSessionHash(parsed) - require.Equal(t, "123e4567-e89b-12d3-a456-426614174000", hash, - "metadata session_id should take priority over SessionContext") + require.Equal(t, "123e4567-e89b-12d3-a456-426614174000", hash, "metadata session_id should take priority over SessionContext") } func TestGenerateSessionHash_MetadataJSON_HasHighestPriority(t *testing.T) { svc := &GatewayService{} - - parsed := &ParsedRequest{ - MetadataUserID: `{"device_id":"a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2","account_uuid":"","session_id":"c72554f2-1234-5678-abcd-123456789abc"}`, - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + metadata := `{"device_id":"a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2","account_uuid":"","session_id":"c72554f2-1234-5678-abcd-123456789abc"}` + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, metadata), nil) hash := svc.GenerateSessionHash(parsed) require.Equal(t, "c72554f2-1234-5678-abcd-123456789abc", hash, "JSON format metadata session_id should have highest priority") @@ -230,69 +179,25 @@ func TestGenerateSessionHash_MetadataJSON_HasHighestPriority(t *testing.T) { func TestGenerateSessionHash_NilSessionContextBackwardCompatible(t *testing.T) { svc := &GatewayService{} - - withCtx := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: nil, - } - withoutCtx := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - } + body := anthropicSessionBody(nil, []any{msg("user", "hello")}, "") + withCtx := mustParseSessionHashRequest(t, body, nil) + withoutCtx := mustParseSessionHashRequest(t, body, nil) h1 := svc.GenerateSessionHash(withCtx) h2 := svc.GenerateSessionHash(withoutCtx) require.Equal(t, h1, h2, "nil SessionContext should produce same hash as no SessionContext") } -// ============ 多轮连续会话测试 ============ - func TestGenerateSessionHash_ContinuousConversation_HashChangesWithMessages(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 模拟连续会话：每增加一轮对话，hash 应该不同（内容累积变化） - round1 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: ctx, - } - - round2 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "Hi there!"}, - map[string]any{"role": "user", "content": "How are you?"}, - }, - SessionContext: ctx, - } - - round3 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "Hi there!"}, - map[string]any{"role": "user", "content": "How are you?"}, - map[string]any{"role": "assistant", "content": "I'm doing well!"}, - map[string]any{"role": "user", "content": "Tell me a joke"}, - }, - SessionContext: ctx, - } + round1 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, ""), ctx) + round2 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello"), msg("assistant", "Hi there!"), msg("user", "How are you?")}, ""), ctx) + round3 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello"), msg("assistant", "Hi there!"), msg("user", "How are you?"), msg("assistant", "I'm doing well!"), msg("user", "Tell me a joke")}, ""), ctx) h1 := svc.GenerateSessionHash(round1) h2 := svc.GenerateSessionHash(round2) h3 := svc.GenerateSessionHash(round3) - require.NotEmpty(t, h1) require.NotEmpty(t, h2) require.NotEmpty(t, h3) @@ -303,62 +208,20 @@ func TestGenerateSessionHash_ContinuousConversation_HashChangesWithMessages(t *t func TestGenerateSessionHash_ContinuousConversation_SameRoundSameHash(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 同一轮对话重复请求（如重试）应产生相同 hash - mk := func() *ParsedRequest { - return &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - map[string]any{"role": "assistant", "content": "Hi there!"}, - map[string]any{"role": "user", "content": "How are you?"}, - }, - SessionContext: ctx, - } - } + body := anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello"), msg("assistant", "Hi there!"), msg("user", "How are you?")}, "") + mk := func() *ParsedRequest { return mustParseSessionHashRequest(t, body, ctx) } h1 := svc.GenerateSessionHash(mk()) h2 := svc.GenerateSessionHash(mk()) require.Equal(t, h1, h2, "same conversation state should produce identical hash on retry") } -// ============ 消息回退测试 ============ - func TestGenerateSessionHash_MessageRollback(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 模拟消息回退：用户删掉最后一轮再重发 - original := &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "msg1"}, - map[string]any{"role": "assistant", "content": "reply1"}, - map[string]any{"role": "user", "content": "msg2"}, - map[string]any{"role": "assistant", "content": "reply2"}, - map[string]any{"role": "user", "content": "msg3"}, - }, - SessionContext: ctx, - } - - // 回退到 msg2 后，用新的 msg3 替代 - rollback := &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "msg1"}, - map[string]any{"role": "assistant", "content": "reply1"}, - map[string]any{"role": "user", "content": "msg2"}, - map[string]any{"role": "assistant", "content": "reply2"}, - map[string]any{"role": "user", "content": "different msg3"}, - }, - SessionContext: ctx, - } + original := mustParseSessionHashRequest(t, anthropicSessionBody("System prompt", []any{msg("user", "msg1"), msg("assistant", "reply1"), msg("user", "msg2"), msg("assistant", "reply2"), msg("user", "msg3")}, ""), ctx) + rollback := mustParseSessionHashRequest(t, anthropicSessionBody("System prompt", []any{msg("user", "msg1"), msg("assistant", "reply1"), msg("user", "msg2"), msg("assistant", "reply2"), msg("user", "different msg3")}, ""), ctx) hOrig := svc.GenerateSessionHash(original) hRollback := svc.GenerateSessionHash(rollback) @@ -367,58 +230,19 @@ func TestGenerateSessionHash_MessageRollback(t *testing.T) { func TestGenerateSessionHash_MessageRollbackSameContent(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 回退后重新发送相同内容 → 相同 hash（合理的粘性恢复） - mk := func() *ParsedRequest { - return &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "msg1"}, - map[string]any{"role": "assistant", "content": "reply1"}, - map[string]any{"role": "user", "content": "msg2"}, - }, - SessionContext: ctx, - } - } + body := anthropicSessionBody("System prompt", []any{msg("user", "msg1"), msg("assistant", "reply1"), msg("user", "msg2")}, "") + mk := func() *ParsedRequest { return mustParseSessionHashRequest(t, body, ctx) } h1 := svc.GenerateSessionHash(mk()) h2 := svc.GenerateSessionHash(mk()) require.Equal(t, h1, h2, "rollback and resend same content should produce same hash") } -// ============ 相同 System、不同用户消息 ============ - func TestGenerateSessionHash_SameSystemDifferentUsers(t *testing.T) { svc := &GatewayService{} - - // 两个不同用户使用相同 system prompt 但发送不同消息 - user1 := &ParsedRequest{ - System: "You are a code reviewer.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "Review this Go code"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "vscode", - APIKeyID: 1, - }, - } - user2 := &ParsedRequest{ - System: "You are a code reviewer.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "Review this Python code"}, - }, - SessionContext: &SessionContext{ - ClientIP: "2.2.2.2", - UserAgent: "vscode", - APIKeyID: 2, - }, - } + user1 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a code reviewer.", []any{msg("user", "Review this Go code")}, ""), &SessionContext{ClientIP: "1.1.1.1", UserAgent: "vscode", APIKeyID: 1}) + user2 := mustParseSessionHashRequest(t, anthropicSessionBody("You are a code reviewer.", []any{msg("user", "Review this Python code")}, ""), &SessionContext{ClientIP: "2.2.2.2", UserAgent: "vscode", APIKeyID: 2}) h1 := svc.GenerateSessionHash(user1) h2 := svc.GenerateSessionHash(user2) @@ -427,55 +251,20 @@ func TestGenerateSessionHash_SameSystemDifferentUsers(t *testing.T) { func TestGenerateSessionHash_SameSystemSameMessageDifferentContext(t *testing.T) { svc := &GatewayService{} - - // 这是修复的核心场景：两个不同用户发送完全相同的 system + messages（如 "hello"） - // 有了 SessionContext 后应该产生不同 hash - user1 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "Mozilla/5.0", - APIKeyID: 10, - }, - } - user2 := &ParsedRequest{ - System: "You are a helpful assistant.", - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "2.2.2.2", - UserAgent: "Mozilla/5.0", - APIKeyID: 20, - }, - } + body := anthropicSessionBody("You are a helpful assistant.", []any{msg("user", "hello")}, "") + user1 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: "Mozilla/5.0", APIKeyID: 10}) + user2 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "2.2.2.2", UserAgent: "Mozilla/5.0", APIKeyID: 20}) h1 := svc.GenerateSessionHash(user1) h2 := svc.GenerateSessionHash(user2) require.NotEqual(t, h1, h2, "CRITICAL: same system+messages but different users should get different hashes") } -// ============ SessionContext 各字段独立影响测试 ============ - func TestGenerateSessionHash_SessionContext_IPDifference(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(ip string) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: ip, - UserAgent: "same-ua", - APIKeyID: 1, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: ip, UserAgent: "same-ua", APIKeyID: 1}) } h1 := svc.GenerateSessionHash(base("1.1.1.1")) @@ -485,18 +274,9 @@ func TestGenerateSessionHash_SessionContext_IPDifference(t *testing.T) { func TestGenerateSessionHash_SessionContext_UADifference(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(ua string) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: ua, - APIKeyID: 1, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: ua, APIKeyID: 1}) } h1 := svc.GenerateSessionHash(base("Mozilla/5.0")) @@ -506,18 +286,9 @@ func TestGenerateSessionHash_SessionContext_UADifference(t *testing.T) { func TestGenerateSessionHash_SessionContext_UAVersionNoiseIgnored(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(ua string) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: ua, - APIKeyID: 1, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: ua, APIKeyID: 1}) } h1 := svc.GenerateSessionHash(base("Mozilla/5.0 codex_cli_rs/0.1.0")) @@ -527,18 +298,9 @@ func TestGenerateSessionHash_SessionContext_UAVersionNoiseIgnored(t *testing.T) func TestGenerateSessionHash_SessionContext_FreeformUAVersionNoiseIgnored(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(ua string) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: ua, - APIKeyID: 1, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: ua, APIKeyID: 1}) } h1 := svc.GenerateSessionHash(base("Codex CLI 0.1.0")) @@ -548,18 +310,9 @@ func TestGenerateSessionHash_SessionContext_FreeformUAVersionNoiseIgnored(t *tes func TestGenerateSessionHash_SessionContext_APIKeyIDDifference(t *testing.T) { svc := &GatewayService{} - + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") base := func(keyID int64) *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "same-ua", - APIKeyID: keyID, - }, - } + return mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: "same-ua", APIKeyID: keyID}) } h1 := svc.GenerateSessionHash(base(1)) @@ -567,24 +320,12 @@ func TestGenerateSessionHash_SessionContext_APIKeyIDDifference(t *testing.T) { require.NotEqual(t, h1, h2, "different APIKeyID should produce different hash") } -// ============ 多用户并发相同消息场景 ============ - func TestGenerateSessionHash_MultipleUsersSameFirstMessage(t *testing.T) { svc := &GatewayService{} - - // 模拟 5 个不同用户同时发送 "hello" → 应该产生 5 个不同的 hash hashes := make(map[string]bool) + body := anthropicSessionBody(nil, []any{msg("user", "hello")}, "") for i := 0; i < 5; i++ { - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "192.168.1." + string(rune('1'+i)), - UserAgent: "client-" + string(rune('A'+i)), - APIKeyID: int64(i + 1), - }, - } + parsed := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "192.168.1." + string(rune('1'+i)), UserAgent: "client-" + string(rune('A'+i)), APIKeyID: int64(i + 1)}) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h) require.False(t, hashes[h], "hash collision detected for user %d", i) @@ -593,134 +334,56 @@ func TestGenerateSessionHash_MultipleUsersSameFirstMessage(t *testing.T) { require.Len(t, hashes, 5, "5 different users should produce 5 unique hashes") } -// ============ 连续会话粘性：多轮对话同一用户 ============ - func TestGenerateSessionHash_SameUserGrowingConversation(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "browser", APIKeyID: 42} - - // 模拟同一用户的连续会话，每轮 hash 不同但同用户重试保持一致 - messages := []map[string]any{ - {"role": "user", "content": "msg1"}, - {"role": "assistant", "content": "reply1"}, - {"role": "user", "content": "msg2"}, - {"role": "assistant", "content": "reply2"}, - {"role": "user", "content": "msg3"}, - {"role": "assistant", "content": "reply3"}, - {"role": "user", "content": "msg4"}, + messages := []any{ + msg("user", "msg1"), msg("assistant", "reply1"), msg("user", "msg2"), msg("assistant", "reply2"), + msg("user", "msg3"), msg("assistant", "reply3"), msg("user", "msg4"), } prevHash := "" for round := 1; round <= len(messages); round += 2 { - // 构建前 round 条消息 - msgs := make([]any, round) - for j := 0; j < round; j++ { - msgs[j] = messages[j] - } - parsed := &ParsedRequest{ - System: "System", - HasSystem: true, - Messages: msgs, - SessionContext: ctx, - } + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("System", messages[:round], ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "round %d hash should not be empty", round) - if prevHash != "" { require.NotEqual(t, prevHash, h, "round %d hash should differ from previous round", round) } prevHash = h - - // 同一轮重试应该相同 h2 := svc.GenerateSessionHash(parsed) require.Equal(t, h, h2, "retry of round %d should produce same hash", round) } } -// ============ 多轮消息内容结构化测试 ============ - func TestGenerateSessionHash_MultipleUserMessages(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 5 条用户消息（无 assistant 回复） - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "first"}, - map[string]any{"role": "user", "content": "second"}, - map[string]any{"role": "user", "content": "third"}, - map[string]any{"role": "user", "content": "fourth"}, - map[string]any{"role": "user", "content": "fifth"}, - }, - SessionContext: ctx, - } - + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "first"), msg("user", "second"), msg("user", "third"), msg("user", "fourth"), msg("user", "fifth")}, ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h) - // 修改中间一条消息应该改变 hash - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "first"}, - map[string]any{"role": "user", "content": "CHANGED"}, - map[string]any{"role": "user", "content": "third"}, - map[string]any{"role": "user", "content": "fourth"}, - map[string]any{"role": "user", "content": "fifth"}, - }, - SessionContext: ctx, - } - + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "first"), msg("user", "CHANGED"), msg("user", "third"), msg("user", "fourth"), msg("user", "fifth")}, ""), ctx) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h, h2, "changing any message should change the hash") } func TestGenerateSessionHash_MessageOrderMatters(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - parsed1 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "alpha"}, - map[string]any{"role": "user", "content": "beta"}, - }, - SessionContext: ctx, - } - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "beta"}, - map[string]any{"role": "user", "content": "alpha"}, - }, - SessionContext: ctx, - } + parsed1 := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "alpha"), msg("user", "beta")}, ""), ctx) + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", "beta"), msg("user", "alpha")}, ""), ctx) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h1, h2, "message order should affect the hash") } -// ============ 复杂内容格式测试 ============ - func TestGenerateSessionHash_StructuredContent(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 结构化 content（数组形式） - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{"type": "text", "text": "Look at this"}, - map[string]any{"type": "text", "text": "And this too"}, - }, - }, - }, - SessionContext: ctx, - } + content := []any{map[string]any{"type": "text", "text": "Look at this"}, map[string]any{"type": "text", "text": "And this too"}} + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{msg("user", content)}, ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "structured content should produce a hash") @@ -728,100 +391,37 @@ func TestGenerateSessionHash_StructuredContent(t *testing.T) { func TestGenerateSessionHash_ArraySystemPrompt(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 数组格式的 system prompt - parsed := &ParsedRequest{ - System: []any{ - map[string]any{"type": "text", "text": "You are a helpful assistant."}, - map[string]any{"type": "text", "text": "Be concise."}, - }, - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: ctx, - } + system := []any{map[string]any{"type": "text", "text": "You are a helpful assistant."}, map[string]any{"type": "text", "text": "Be concise."}} + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(system, []any{msg("user", "hello")}, ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "array system prompt should produce a hash") } -// ============ SessionContext 与 cache_control 优先级 ============ - func TestGenerateSessionHash_CacheControlOverridesSessionContext(t *testing.T) { svc := &GatewayService{} - - // 当有 cache_control: ephemeral 时，使用第 2 级优先级 - // SessionContext 不应影响结果 - parsed1 := &ParsedRequest{ - System: []any{ - map[string]any{ - "type": "text", - "text": "You are a tool-specific assistant.", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - }, - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "ua1", - APIKeyID: 100, - }, - } - parsed2 := &ParsedRequest{ - System: []any{ - map[string]any{ - "type": "text", - "text": "You are a tool-specific assistant.", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - }, - HasSystem: true, - Messages: []any{ - map[string]any{"role": "user", "content": "hello"}, - }, - SessionContext: &SessionContext{ - ClientIP: "2.2.2.2", - UserAgent: "ua2", - APIKeyID: 200, - }, - } + system := []any{map[string]any{"type": "text", "text": "You are a tool-specific assistant.", "cache_control": map[string]any{"type": "ephemeral"}}} + body := anthropicSessionBody(system, []any{msg("user", "hello")}, "") + parsed1 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: "ua1", APIKeyID: 100}) + parsed2 := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "2.2.2.2", UserAgent: "ua2", APIKeyID: 200}) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) require.Equal(t, h1, h2, "cache_control ephemeral has higher priority, SessionContext should not affect result") } -// ============ 边界情况 ============ - func TestGenerateSessionHash_EmptyMessages(t *testing.T) { svc := &GatewayService{} + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{}, ""), &SessionContext{ClientIP: "1.1.1.1", UserAgent: "test", APIKeyID: 1}) - parsed := &ParsedRequest{ - Messages: []any{}, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "test", - APIKeyID: 1, - }, - } - - // 空 messages + 只有 SessionContext 时，combined.Len() > 0 因为有 context 写入 h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "empty messages with SessionContext should still produce a hash from context") } func TestGenerateSessionHash_EmptyMessagesNoContext(t *testing.T) { svc := &GatewayService{} - - parsed := &ParsedRequest{ - Messages: []any{}, - } + parsed := mustParseSessionHashRequest(t, anthropicSessionBody(nil, []any{}, ""), nil) h := svc.GenerateSessionHash(parsed) require.Empty(t, h, "empty messages without SessionContext should produce empty hash") @@ -829,98 +429,37 @@ func TestGenerateSessionHash_EmptyMessagesNoContext(t *testing.T) { func TestGenerateSessionHash_SessionContextWithEmptyFields(t *testing.T) { svc := &GatewayService{} - - // SessionContext 字段为空字符串和零值时仍应影响 hash - withEmptyCtx := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - SessionContext: &SessionContext{ - ClientIP: "", - UserAgent: "", - APIKeyID: 0, - }, - } - withoutCtx := &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "test"}, - }, - } + body := anthropicSessionBody(nil, []any{msg("user", "test")}, "") + withEmptyCtx := mustParseSessionHashRequest(t, body, &SessionContext{ClientIP: "", UserAgent: "", APIKeyID: 0}) + withoutCtx := mustParseSessionHashRequest(t, body, nil) h1 := svc.GenerateSessionHash(withEmptyCtx) h2 := svc.GenerateSessionHash(withoutCtx) - // 有 SessionContext（即使字段为空）仍然会写入分隔符 "::" 等 require.NotEqual(t, h1, h2, "empty-field SessionContext should still differ from nil SessionContext") } -// ============ 长对话历史测试 ============ - func TestGenerateSessionHash_LongConversation(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "test", APIKeyID: 1} - - // 构建 20 轮对话 messages := make([]any, 0, 40) for i := 0; i < 20; i++ { - messages = append(messages, map[string]any{ - "role": "user", - "content": "user message " + string(rune('A'+i)), - }) - messages = append(messages, map[string]any{ - "role": "assistant", - "content": "assistant reply " + string(rune('A'+i)), - }) - } - - parsed := &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: messages, - SessionContext: ctx, + messages = append(messages, msg("user", "user message "+string(rune('A'+i)))) + messages = append(messages, msg("assistant", "assistant reply "+string(rune('A'+i)))) } + parsed := mustParseSessionHashRequest(t, anthropicSessionBody("System prompt", messages, ""), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h) - // 再加一轮应该不同 - moreMessages := make([]any, len(messages)+2) - copy(moreMessages, messages) - moreMessages[len(messages)] = map[string]any{"role": "user", "content": "one more"} - moreMessages[len(messages)+1] = map[string]any{"role": "assistant", "content": "ok"} - - parsed2 := &ParsedRequest{ - System: "System prompt", - HasSystem: true, - Messages: moreMessages, - SessionContext: ctx, - } - + moreMessages := append(append([]any{}, messages...), msg("user", "one more"), msg("assistant", "ok")) + parsed2 := mustParseSessionHashRequest(t, anthropicSessionBody("System prompt", moreMessages, ""), ctx) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h, h2, "adding more messages to long conversation should change hash") } -// ============ Gemini 原生格式 session hash 测试 ============ - func TestGenerateSessionHash_GeminiContentsProducesHash(t *testing.T) { svc := &GatewayService{} - - // Gemini 格式: contents[].parts[].text - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Hello from Gemini"}, - }, - }, - }, - SessionContext: &SessionContext{ - ClientIP: "1.2.3.4", - UserAgent: "gemini-cli", - APIKeyID: 1, - }, - } + parsed := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Hello from Gemini")}), &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1}) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "Gemini contents with parts should produce a non-empty hash") @@ -928,31 +467,9 @@ func TestGenerateSessionHash_GeminiContentsProducesHash(t *testing.T) { func TestGenerateSessionHash_GeminiDifferentContentsDifferentHash(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - parsed1 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Hello"}, - }, - }, - }, - SessionContext: ctx, - } - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Goodbye"}, - }, - }, - }, - SessionContext: ctx, - } + parsed1 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Hello")}), ctx) + parsed2 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Goodbye")}), ctx) h1 := svc.GenerateSessionHash(parsed1) h2 := svc.GenerateSessionHash(parsed2) @@ -961,28 +478,9 @@ func TestGenerateSessionHash_GeminiDifferentContentsDifferentHash(t *testing.T) func TestGenerateSessionHash_GeminiSameContentsSameHash(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - mk := func() *ParsedRequest { - return &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Hello"}, - }, - }, - map[string]any{ - "role": "model", - "parts": []any{ - map[string]any{"text": "Hi there!"}, - }, - }, - }, - SessionContext: ctx, - } - } + body := geminiSessionBody(nil, []any{geminiMsg("user", "Hello"), geminiMsg("model", "Hi there!")}) + mk := func() *ParsedRequest { return mustParseGeminiSessionHashRequest(t, body, ctx) } h1 := svc.GenerateSessionHash(mk()) h2 := svc.GenerateSessionHash(mk()) @@ -991,36 +489,9 @@ func TestGenerateSessionHash_GeminiSameContentsSameHash(t *testing.T) { func TestGenerateSessionHash_GeminiMultiTurnHashChanges(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - round1 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: ctx, - } - - round2 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - map[string]any{ - "role": "model", - "parts": []any{map[string]any{"text": "Hi!"}}, - }, - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "How are you?"}}, - }, - }, - SessionContext: ctx, - } + round1 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "hello")}), ctx) + round2 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "hello"), geminiMsg("model", "Hi!"), geminiMsg("user", "How are you?")}), ctx) h1 := svc.GenerateSessionHash(round1) h2 := svc.GenerateSessionHash(round2) @@ -1031,34 +502,9 @@ func TestGenerateSessionHash_GeminiMultiTurnHashChanges(t *testing.T) { func TestGenerateSessionHash_GeminiDifferentUsersSameContentDifferentHash(t *testing.T) { svc := &GatewayService{} - - // 核心场景：两个不同用户发送相同 Gemini 格式消息应得到不同 hash - user1 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: &SessionContext{ - ClientIP: "1.1.1.1", - UserAgent: "gemini-cli", - APIKeyID: 10, - }, - } - user2 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: &SessionContext{ - ClientIP: "2.2.2.2", - UserAgent: "gemini-cli", - APIKeyID: 20, - }, - } + body := geminiSessionBody(nil, []any{geminiMsg("user", "hello")}) + user1 := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "1.1.1.1", UserAgent: "gemini-cli", APIKeyID: 10}) + user2 := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "2.2.2.2", UserAgent: "gemini-cli", APIKeyID: 20}) h1 := svc.GenerateSessionHash(user1) h2 := svc.GenerateSessionHash(user2) @@ -1067,31 +513,9 @@ func TestGenerateSessionHash_GeminiDifferentUsersSameContentDifferentHash(t *tes func TestGenerateSessionHash_GeminiSystemInstructionAffectsHash(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - // systemInstruction 经 ParseGatewayRequest 解析后存入 parsed.System - withSys := &ParsedRequest{ - System: []any{ - map[string]any{"text": "You are a coding assistant."}, - }, - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: ctx, - } - withoutSys := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{map[string]any{"text": "hello"}}, - }, - }, - SessionContext: ctx, - } + withSys := mustParseGeminiSessionHashRequest(t, geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "hello")}), ctx) + withoutSys := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "hello")}), ctx) h1 := svc.GenerateSessionHash(withSys) h2 := svc.GenerateSessionHash(withoutSys) @@ -1100,64 +524,21 @@ func TestGenerateSessionHash_GeminiSystemInstructionAffectsHash(t *testing.T) { func TestGenerateSessionHash_GeminiMultiPartMessage(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - // 多 parts 的消息 - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Part 1"}, - map[string]any{"text": "Part 2"}, - map[string]any{"text": "Part 3"}, - }, - }, - }, - SessionContext: ctx, - } - + parsed := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Part 1", "Part 2", "Part 3")}), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "multi-part Gemini message should produce a hash") - // 不同内容的多 parts - parsed2 := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Part 1"}, - map[string]any{"text": "CHANGED"}, - map[string]any{"text": "Part 3"}, - }, - }, - }, - SessionContext: ctx, - } - + parsed2 := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, []any{geminiMsg("user", "Part 1", "CHANGED", "Part 3")}), ctx) h2 := svc.GenerateSessionHash(parsed2) require.NotEqual(t, h, h2, "changing a part should change the hash") } func TestGenerateSessionHash_GeminiNonTextPartsIgnored(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "1.2.3.4", UserAgent: "gemini-cli", APIKeyID: 1} - - // 含非 text 类型 parts（如 inline_data），应被跳过但不报错 - parsed := &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "parts": []any{ - map[string]any{"text": "Describe this image"}, - map[string]any{"inline_data": map[string]any{"mime_type": "image/png", "data": "base64..."}}, - }, - }, - }, - SessionContext: ctx, - } + content := []any{map[string]any{"role": "user", "parts": []any{map[string]any{"text": "Describe this image"}, map[string]any{"inline_data": map[string]any{"mime_type": "image/png", "data": "base64..."}}}}} + parsed := mustParseGeminiSessionHashRequest(t, geminiSessionBody(nil, content), ctx) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "Gemini message with mixed parts should still produce a hash from text parts") @@ -1165,107 +546,41 @@ func TestGenerateSessionHash_GeminiNonTextPartsIgnored(t *testing.T) { func TestGenerateSessionHash_GeminiMultiTurnHashNotSticky(t *testing.T) { svc := &GatewayService{} - ctx := &SessionContext{ClientIP: "10.0.0.1", UserAgent: "gemini-cli", APIKeyID: 42} + rounds := []string{ + geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "Write a Go function")}), + geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "Write a Go function"), geminiMsg("model", "func hello() {}"), geminiMsg("user", "Add error handling")}), + geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "Write a Go function"), geminiMsg("model", "func hello() {}"), geminiMsg("user", "Add error handling"), geminiMsg("model", "func hello() error { return nil }"), geminiMsg("user", "Now add tests")}), + } - // 模拟同一 Gemini 会话的三轮请求，每轮 contents 累积增长。 - // 验证预期行为：每轮 hash 都不同，即 GenerateSessionHash 不具备跨轮粘性。 - // 这是 by-design 的——Gemini 的跨轮粘性由 Digest Fallback（BuildGeminiDigestChain）负责。 - round1Body := []byte(`{ - "systemInstruction": {"parts": [{"text": "You are a coding assistant."}]}, - "contents": [ - {"role": "user", "parts": [{"text": "Write a Go function"}]} - ] - }`) - round2Body := []byte(`{ - "systemInstruction": {"parts": [{"text": "You are a coding assistant."}]}, - "contents": [ - {"role": "user", "parts": [{"text": "Write a Go function"}]}, - {"role": "model", "parts": [{"text": "func hello() {}"}]}, - {"role": "user", "parts": [{"text": "Add error handling"}]} - ] - }`) - round3Body := []byte(`{ - "systemInstruction": {"parts": [{"text": "You are a coding assistant."}]}, - "contents": [ - {"role": "user", "parts": [{"text": "Write a Go function"}]}, - {"role": "model", "parts": [{"text": "func hello() {}"}]}, - {"role": "user", "parts": [{"text": "Add error handling"}]}, - {"role": "model", "parts": [{"text": "func hello() error { return nil }"}]}, - {"role": "user", "parts": [{"text": "Now add tests"}]} - ] - }`) - - hashes := make([]string, 3) - for i, body := range [][]byte{round1Body, round2Body, round3Body} { - parsed, err := ParseGatewayRequest(body, "gemini") - require.NoError(t, err) - parsed.SessionContext = ctx + hashes := make([]string, len(rounds)) + for i, body := range rounds { + parsed := mustParseGeminiSessionHashRequest(t, body, ctx) hashes[i] = svc.GenerateSessionHash(parsed) require.NotEmpty(t, hashes[i], "round %d hash should not be empty", i+1) } - - // 每轮 hash 都不同——这是预期行为 require.NotEqual(t, hashes[0], hashes[1], "round 1 vs 2 hash should differ (contents grow)") require.NotEqual(t, hashes[1], hashes[2], "round 2 vs 3 hash should differ (contents grow)") require.NotEqual(t, hashes[0], hashes[2], "round 1 vs 3 hash should differ") - // 同一轮重试应产生相同 hash - parsed1Again, err := ParseGatewayRequest(round2Body, "gemini") - require.NoError(t, err) - parsed1Again.SessionContext = ctx - h2Again := svc.GenerateSessionHash(parsed1Again) + parsedAgain := mustParseGeminiSessionHashRequest(t, rounds[1], ctx) + h2Again := svc.GenerateSessionHash(parsedAgain) require.Equal(t, hashes[1], h2Again, "retry of same round should produce same hash") } func TestGenerateSessionHash_GeminiEndToEnd(t *testing.T) { svc := &GatewayService{} - - // 端到端测试：模拟 ParseGatewayRequest + GenerateSessionHash 完整流程 - body := []byte(`{ - "model": "gemini-2.5-pro", - "systemInstruction": { - "parts": [{"text": "You are a coding assistant."}] - }, - "contents": [ - {"role": "user", "parts": [{"text": "Write a Go function"}]}, - {"role": "model", "parts": [{"text": "Here is a function..."}]}, - {"role": "user", "parts": [{"text": "Now add error handling"}]} - ] - }`) - - parsed, err := ParseGatewayRequest(body, "gemini") - require.NoError(t, err) - parsed.SessionContext = &SessionContext{ - ClientIP: "10.0.0.1", - UserAgent: "gemini-cli/1.0", - APIKeyID: 42, - } + body := geminiSessionBody([]any{map[string]any{"text": "You are a coding assistant."}}, []any{geminiMsg("user", "Write a Go function"), geminiMsg("model", "Here is a function..."), geminiMsg("user", "Now add error handling")}) + parsed := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "10.0.0.1", UserAgent: "gemini-cli/1.0", APIKeyID: 42}) h := svc.GenerateSessionHash(parsed) require.NotEmpty(t, h, "end-to-end Gemini flow should produce a hash") - // 同一请求再次解析应产生相同 hash - parsed2, err := ParseGatewayRequest(body, "gemini") - require.NoError(t, err) - parsed2.SessionContext = &SessionContext{ - ClientIP: "10.0.0.1", - UserAgent: "gemini-cli/1.0", - APIKeyID: 42, - } - + parsed2 := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "10.0.0.1", UserAgent: "gemini-cli/1.0", APIKeyID: 42}) h2 := svc.GenerateSessionHash(parsed2) require.Equal(t, h, h2, "same request should produce same hash") - // 不同用户发送相同请求应产生不同 hash - parsed3, err := ParseGatewayRequest(body, "gemini") - require.NoError(t, err) - parsed3.SessionContext = &SessionContext{ - ClientIP: "10.0.0.2", - UserAgent: "gemini-cli/1.0", - APIKeyID: 99, - } - + parsed3 := mustParseGeminiSessionHashRequest(t, body, &SessionContext{ClientIP: "10.0.0.2", UserAgent: "gemini-cli/1.0", APIKeyID: 99}) h3 := svc.GenerateSessionHash(parsed3) require.NotEqual(t, h, h3, "different user with same Gemini request should get different hash") } diff --git a/backend/internal/service/image_generation_intent.go b/backend/internal/service/image_generation_intent.go index 4aca1239ff97..80b6c66d56a7 100644 --- a/backend/internal/service/image_generation_intent.go +++ b/backend/internal/service/image_generation_intent.go @@ -1,7 +1,6 @@ package service import ( - "encoding/json" "strings" "github.com/tidwall/gjson" @@ -91,7 +90,7 @@ func openAIJSONToolsContainImageGeneration(tools gjson.Result) bool { } found := false tools.ForEach(func(_, item gjson.Result) bool { - if strings.TrimSpace(item.Get("type").String()) == "image_generation" { + if openAIJSONString(item.Get("type")) == "image_generation" { found = true return false } @@ -100,6 +99,36 @@ func openAIJSONToolsContainImageGeneration(tools gjson.Result) bool { return found } +func openAIRequestBodyHasImageGenerationTool(body []byte) bool { + if len(body) == 0 || !gjson.ValidBytes(body) { + return false + } + return openAIJSONToolsContainImageGeneration(gjson.GetBytes(body, "tools")) +} + +func openAIRequestBodyImageGenerationToolNeedsNormalization(body []byte) bool { + if len(body) == 0 || !gjson.ValidBytes(body) { + return false + } + tools := gjson.GetBytes(body, "tools") + if !tools.IsArray() { + return false + } + needsNormalization := false + tools.ForEach(func(_, item gjson.Result) bool { + if openAIJSONString(item.Get("type")) != "image_generation" { + return true + } + // 只有旧字段需要迁移时才进入 map 修改，纯计费读取保持 raw 路径。 + if item.Get("format").Exists() || item.Get("compression").Exists() { + needsNormalization = true + return false + } + return true + }) + return needsNormalization +} + func openAIJSONToolChoiceSelectsImageGeneration(choice gjson.Result) bool { if !choice.Exists() { return false @@ -159,17 +188,6 @@ func apiKeyGroup(apiKey *APIKey) *Group { return apiKey.Group } -func cloneRequestMapForImageIntent(body []byte) map[string]any { - if len(body) == 0 { - return nil - } - var out map[string]any - if err := json.Unmarshal(body, &out); err != nil { - return nil - } - return out -} - type OpenAIResponsesImageBillingConfig struct { Model string SizeTier string @@ -225,8 +243,43 @@ func resolveOpenAIResponsesImageBillingConfigFromBody(body []byte, fallbackModel } func resolveOpenAIResponsesImageBillingConfigDetailedFromBody(body []byte, fallbackModel string) (OpenAIResponsesImageBillingConfig, error) { - reqBody := cloneRequestMapForImageIntent(body) - return resolveOpenAIResponsesImageBillingConfigDetailed(reqBody, fallbackModel) + imageModel := "" + imageSize := "" + hasImageTool := false + if len(body) > 0 && gjson.ValidBytes(body) { + tools := gjson.GetBytes(body, "tools") + if tools.IsArray() { + tools.ForEach(func(_, item gjson.Result) bool { + if openAIJSONString(item.Get("type")) != "image_generation" { + return true + } + hasImageTool = true + imageModel = openAIJSONString(item.Get("model")) + imageSize = openAIJSONString(item.Get("size")) + return false + }) + } + if imageSize == "" { + imageSize = openAIJSONString(gjson.GetBytes(body, "size")) + } + if imageModel == "" { + bodyModel := openAIJSONString(gjson.GetBytes(body, "model")) + if isOpenAIImageBillingModelAlias(bodyModel) || !hasImageTool { + imageModel = bodyModel + } + } + } + if imageModel == "" && hasImageTool { + imageModel = "gpt-image-2" + } + if imageModel == "" { + imageModel = strings.TrimSpace(fallbackModel) + } + return OpenAIResponsesImageBillingConfig{ + Model: imageModel, + SizeTier: normalizeOpenAIImageSizeTier(imageSize), + InputSize: imageSize, + }, nil } func isOpenAIImageBillingModelAlias(model string) bool { @@ -236,3 +289,10 @@ func isOpenAIImageBillingModelAlias(model string) bool { } return isOpenAIImageGenerationModel(normalized) || strings.Contains(normalized, "image") } + +func openAIJSONString(value gjson.Result) string { + if value.Type != gjson.String { + return "" + } + return strings.TrimSpace(value.String()) +} diff --git a/backend/internal/service/image_generation_intent_test.go b/backend/internal/service/image_generation_intent_test.go index 4621e9d98572..59aab39c4b4b 100644 --- a/backend/internal/service/image_generation_intent_test.go +++ b/backend/internal/service/image_generation_intent_test.go @@ -84,6 +84,17 @@ func TestResolveOpenAIResponsesImageBillingConfigToolModelWins(t *testing.T) { require.Equal(t, "2K", imageSize) } +func TestResolveOpenAIResponsesImageBillingConfigFromBodyIgnoresUnrelatedLargeInput(t *testing.T) { + cfg, err := resolveOpenAIResponsesImageBillingConfigDetailedFromBody( + []byte(`{"model":"mapped-text-model","tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[{"type":"message","content":[{"type":"input_text","text":"hi","nonce":1e1000000}]}]}`), + "requested-model", + ) + require.NoError(t, err) + require.Equal(t, "gpt-image-2", cfg.Model) + require.Equal(t, "2K", cfg.SizeTier) + require.Equal(t, "2048x1152", cfg.InputSize) +} + func TestResolveOpenAIResponsesImageBillingConfigSupportsOfficialAndCustomSizes(t *testing.T) { tests := []struct { name string diff --git a/backend/internal/service/model_rate_limit.go b/backend/internal/service/model_rate_limit.go index e4dadbe94bfc..6697a8450e93 100644 --- a/backend/internal/service/model_rate_limit.go +++ b/backend/internal/service/model_rate_limit.go @@ -6,7 +6,10 @@ import ( "time" ) -const modelRateLimitsKey = "model_rate_limits" +const ( + modelRateLimitsKey = "model_rate_limits" + antigravityGeminiModelRateLimitKey = "antigravity:gemini" +) type modelRateLimitExtraScopesContextKey struct{} @@ -62,6 +65,9 @@ func (a *Account) isModelRateLimitedWithContext(ctx context.Context, requestedMo modelKey := a.GetMappedModel(requestedModel) if a.Platform == PlatformAntigravity { modelKey = resolveFinalAntigravityModelKey(ctx, a, requestedModel) + if isAntigravityGeminiModel(modelKey) && a.isRateLimitActiveForKey(antigravityGeminiModelRateLimitKey) { + return true + } } modelKey = strings.TrimSpace(modelKey) if modelKey == "" { @@ -103,6 +109,11 @@ func (a *Account) GetModelRateLimitRemainingTimeWithContext(ctx context.Context, remaining = scopeRemaining } } + if a.Platform == PlatformAntigravity && isAntigravityGeminiModel(modelKey) { + if familyRemaining := a.getRateLimitRemainingForKey(antigravityGeminiModelRateLimitKey); familyRemaining > remaining { + remaining = familyRemaining + } + } return remaining } @@ -118,6 +129,22 @@ func resolveFinalAntigravityModelKey(ctx context.Context, account *Account, requ return modelKey } +func isAntigravityGeminiModel(model string) bool { + return strings.HasPrefix(normalizeAntigravityModelName(model), "gemini-") +} + +func antigravityModelRateLimitKeys(model string) []string { + model = strings.TrimSpace(model) + if model == "" { + return nil + } + keys := []string{model} + if isAntigravityGeminiModel(model) && model != antigravityGeminiModelRateLimitKey { + keys = append(keys, antigravityGeminiModelRateLimitKey) + } + return keys +} + func (a *Account) modelRateLimitResetAt(scope string) *time.Time { if a == nil || a.Extra == nil || scope == "" { return nil diff --git a/backend/internal/service/model_rate_limit_test.go b/backend/internal/service/model_rate_limit_test.go index c0b8f83c7532..e1df57c55f95 100644 --- a/backend/internal/service/model_rate_limit_test.go +++ b/backend/internal/service/model_rate_limit_test.go @@ -121,6 +121,36 @@ func TestIsModelRateLimited(t *testing.T) { requestedModel: "gemini-3-pro-preview", expected: true, }, + { + name: "antigravity platform - gemini family rate limit blocks mapped preview", + account: &Account{ + Platform: PlatformAntigravity, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": future, + }, + }, + }, + }, + requestedModel: "gemini-3-pro-preview", + expected: true, + }, + { + name: "antigravity platform - gemini family rate limit does not block claude", + account: &Account{ + Platform: PlatformAntigravity, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": future, + }, + }, + }, + }, + requestedModel: "claude-sonnet-4-5", + expected: false, + }, { name: "non-antigravity platform - gemini-3-pro-preview NOT mapped", account: &Account{ @@ -329,6 +359,38 @@ func TestGetModelRateLimitRemainingTime(t *testing.T) { minExpected: 4 * time.Minute, maxExpected: 6 * time.Minute, }, + { + name: "antigravity platform - gemini family rate limit remaining", + account: &Account{ + Platform: PlatformAntigravity, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": future10m, + }, + }, + }, + }, + requestedModel: "gemini-3-pro-preview", + minExpected: 9 * time.Minute, + maxExpected: 11 * time.Minute, + }, + { + name: "antigravity platform - gemini family remaining ignored for claude", + account: &Account{ + Platform: PlatformAntigravity, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": future10m, + }, + }, + }, + }, + requestedModel: "claude-sonnet-4-5", + minExpected: 0, + maxExpected: 0, + }, } for _, tt := range tests { diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 09ce9c1111e6..f0ebbc61434c 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -372,7 +372,6 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash( _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) return nil, nil } - result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency) if acquireErr == nil && result != nil && result.Acquired { _ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL()) @@ -980,6 +979,13 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatible(ctx context.C if !openAIAccountTypeCompatible(account, req.RequiredAccountType) { return false } + // Quota auto-pause must be evaluated during the initial filter too. Without it the + // TopK candidate pool can be filled with paused accounts and the later fresh/DB + // rechecks won't reach healthy accounts that fell outside TopK, manifesting as + // "no available accounts" even though healthy ones exist. + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { + return false + } if req.RequestedModel != "" && !account.IsModelSupported(req.RequestedModel) { return false } @@ -1253,6 +1259,7 @@ func (s *OpenAIGatewayService) selectAccountWithScheduler( requiredAccountType string, requireCompact bool, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { + ctx = s.withOpenAIQuotaAutoPauseContext(ctx) decision := OpenAIAccountScheduleDecision{} scheduler := s.getOpenAIAccountScheduler(ctx) if scheduler == nil { diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index bbaa931e13fd..0e3af717b566 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -747,6 +747,224 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyRateLimite require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) } +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy5hThreshold(t *testing.T) { + ctx := context.Background() + primary := Account{ + ID: 35001, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 95.0, + "auto_pause_5h_threshold": 0.95, + }, + } + secondary := Account{ID: 35002, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35002), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AllowsBelow5hThreshold(t *testing.T) { + ctx := context.Background() + primary := Account{ + ID: 35101, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 80.0, + "auto_pause_5h_threshold": 0.95, + }, + } + secondary := Account{ID: 35102, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35101), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy7dThreshold(t *testing.T) { + ctx := context.Background() + primary := Account{ + ID: 35201, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_7d_used_percent": 95.0, + "auto_pause_7d_threshold": 0.95, + }, + } + secondary := Account{ID: 35202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35202), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UnconfiguredThresholdKeepsLegacyBehavior(t *testing.T) { + ctx := context.Background() + primary := Account{ID: 35301, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, Extra: map[string]any{"codex_5h_used_percent": 99.0, "codex_7d_used_percent": 99.0}} + secondary := Account{ID: 35302, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35301), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UsesGlobalDefaultThreshold(t *testing.T) { + ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold5h: 0.95}) + primary := Account{ + ID: 35401, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 95.0, + }, + } + secondary := Account{ID: 35402, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35402), account.ID) +} + +// Regression: a per-account explicit-disable flag exempts the account from auto-pause +// even when a global default threshold is set. Without this, "leave threshold blank" +// silently falls back to global default and admins have no way to whitelist a single +// account. +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_PerAccountDisableOverridesGlobalDefault(t *testing.T) { + ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold5h: 0.95}) + // Account has high usage AND no per-account threshold (would normally fall back to + // the global default and get paused), but the explicit disable flag is set. + primary := Account{ + ID: 35701, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 99.0, + "auto_pause_5h_disabled": true, + }, + } + secondary := Account{ID: 35702, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35701), account.ID) +} + +// Disable is per-window: disabling only 5h must still allow 7d auto-pause to fire. +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_PerWindowDisableScoped(t *testing.T) { + ctx := context.Background() + primary := Account{ + ID: 35801, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 99.0, + "codex_7d_used_percent": 99.0, + "auto_pause_5h_disabled": true, + "auto_pause_7d_threshold": 0.95, + }, + } + secondary := Account{ID: 35802, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35802), account.ID, "7d auto-pause must still fire even though 5h is disabled") +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_StaleUsageWindowResetSkipsPause(t *testing.T) { + ctx := context.Background() + // Usage is over threshold but the window's reset time has already passed, so the + // cached percentage is stale (the real window rolled over) and the account must NOT + // stay paused — otherwise it could be skipped forever with no traffic to refresh it. + primary := Account{ + ID: 35501, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 99.0, + "auto_pause_5h_threshold": 0.95, + "codex_5h_reset_at": time.Now().Add(-time.Minute).Format(time.RFC3339), + }, + } + secondary := Account{ID: 35502, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35501), account.ID) +} + +func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_FreshUsageWindowStillPauses(t *testing.T) { + ctx := context.Background() + // Same as above but the window has not reset yet, so the account stays paused. + primary := Account{ + ID: 35601, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 99.0, + "auto_pause_5h_threshold": 0.95, + "codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339), + }, + } + secondary := Account{ID: 35602, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5} + svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}} + + account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil) + require.NoError(t, err) + require.NotNil(t, account) + require.Equal(t, int64(35602), account.ID) +} + func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_SkipsFreshlyRateLimitedSnapshotCandidate(t *testing.T) { ctx := context.Background() groupID := int64(10102) @@ -1492,6 +1710,85 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceTopKFallback } } +// Regression: TopK initial filter must drop quota-auto-paused accounts. Otherwise +// the candidate pool is filled with paused accounts, healthy accounts fall outside +// TopK, and the scheduler returns "no available accounts" even though healthy ones +// exist. +func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceTopKExcludesQuotaPaused(t *testing.T) { + ctx := context.Background() + groupID := int64(110) + accounts := []Account{ + { + ID: 37001, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + "codex_5h_used_percent": 96.0, + "auto_pause_5h_threshold": 0.95, + }, + }, + { + ID: 37002, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 5, + }, + } + + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.LBTopK = 1 // TopK=1 makes the bug fatal: paused account would crowd out the healthy one entirely + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 0.4 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1.0 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1.0 + + concurrencyCache := schedulerTestConcurrencyCache{ + loadMap: map[int64]*AccountLoadInfo{ + 37001: {AccountID: 37001, LoadRate: 5, WaitingCount: 0}, + 37002: {AccountID: 37002, LoadRate: 5, WaitingCount: 0}, + }, + acquireResults: map[int64]bool{ + 37002: true, + }, + } + + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: &schedulerTestGatewayCache{}, + cfg: cfg, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"), + concurrencyService: NewConcurrencyService(concurrencyCache), + } + + selection, decision, err := svc.SelectAccountWithScheduler( + ctx, + &groupID, + "", + "", + "gpt-5.1", + nil, + OpenAIUpstreamTransportAny, + false, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(37002), selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + // Only the healthy account should ever enter the candidate pool; the paused one + // must be filtered out at the initial-filter stage. + require.Equal(t, 1, decision.CandidateCount) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics(t *testing.T) { ctx := context.Background() groupID := int64(12) diff --git a/backend/internal/service/openai_embeddings.go b/backend/internal/service/openai_embeddings.go index 359df3bb80a2..7c710259f3fb 100644 --- a/backend/internal/service/openai_embeddings.go +++ b/backend/internal/service/openai_embeddings.go @@ -104,7 +104,7 @@ func (s *OpenAIGatewayService) ForwardEmbeddings( defer func() { _ = resp.Body.Close() }() if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_failover_cached_body_test.go b/backend/internal/service/openai_failover_cached_body_test.go new file mode 100644 index 000000000000..776f1b1ac56c --- /dev/null +++ b/backend/internal/service/openai_failover_cached_body_test.go @@ -0,0 +1,134 @@ +package service + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestOpenAIGatewayService_Forward_FailoverReparsesCachedBodyForNextAccount(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + requestModel string + firstMapping map[string]any + secondMapping map[string]any + wantFirst string + wantSecond string + }{ + { + name: "both accounts have mapping", + firstMapping: map[string]any{"alias-model": "base-model-a"}, + secondMapping: map[string]any{"alias-model": "base-model-b"}, + wantFirst: "base-model-a", + wantSecond: "base-model-b", + }, + { + name: "first account has mapping second account has none", + requestModel: "gpt-5.4-high", + firstMapping: map[string]any{"gpt-5.4-high": "gpt-5.4"}, + wantFirst: "gpt-5.4", + wantSecond: "gpt-5.4", + }, + { + name: "first account has no mapping second account has mapping", + secondMapping: map[string]any{"alias-model": "base-model-b"}, + wantFirst: "alias-model", + wantSecond: "base-model-b", + }, + { + name: "legacy context cache is ignored when mappings differ", + firstMapping: map[string]any{"alias-model": "base-model-a"}, + secondMapping: map[string]any{"alias-model": "base-model-b"}, + wantFirst: "base-model-a", + wantSecond: "base-model-b", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + requestModel := tt.requestModel + if requestModel == "" { + requestModel = "alias-model" + } + body := []byte(`{"model":"` + requestModel + `","stream":false,"instructions":"cache-test","input":"hello"}`) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + { + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid-failover-a"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"type":"rate_limit_error","message":"rate limited"}}`)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid-ok-b"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_123","status":"completed","model":"ok","output":[],"usage":{"input_tokens":1,"output_tokens":1}}`)), + }, + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + firstAccount := openAIFailoverCachedBodyTestAccount(1, "account-a", tt.firstMapping) + secondAccount := openAIFailoverCachedBodyTestAccount(2, "account-b", tt.secondMapping) + + _, err := svc.Forward(context.Background(), c, firstAccount, body) + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.True(t, errors.As(err, &failoverErr)) + require.Len(t, upstream.bodies, 1) + require.Equal(t, tt.wantFirst, gjson.GetBytes(upstream.bodies[0], "model").String()) + + c.Set("openai_parsed_request_body", map[string]any{"model": tt.wantFirst, "stream": true}) + result, err := svc.Forward(context.Background(), c, secondAccount, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.bodies, 2) + require.Equal(t, tt.wantSecond, gjson.GetBytes(upstream.bodies[1], "model").String()) + }) + } +} + +func TestGetOpenAIRequestBodyMap_IgnoresLegacyContextCache(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Set("openai_parsed_request_body", map[string]any{"model": "base-model-a", "stream": true}) + + got, err := getOpenAIRequestBodyMap(c, []byte(`{"model":"alias-model","stream":false}`)) + require.NoError(t, err) + require.Equal(t, "alias-model", got["model"]) + require.Equal(t, false, got["stream"]) +} + +func openAIFailoverCachedBodyTestAccount(id int64, name string, mapping map[string]any) *Account { + credentials := map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-account"} + if mapping != nil { + credentials["model_mapping"] = mapping + } + return &Account{ + ID: id, + Name: name, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: credentials, + Status: StatusActive, + Schedulable: true, + RateMultiplier: f64p(1), + } +} diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index 807ff43afc8a..6e91d85c83a3 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -241,7 +241,7 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions( // 8. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index e351fa75ba7a..3ff6fac4aa8e 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -183,7 +183,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( // 7. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index 291c217e85df..4398bd27cb38 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -300,7 +300,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // 8. Handle error response with failover if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_gateway_record_usage_test.go b/backend/internal/service/openai_gateway_record_usage_test.go index e330a0fd68c7..7c1f959096b9 100644 --- a/backend/internal/service/openai_gateway_record_usage_test.go +++ b/backend/internal/service/openai_gateway_record_usage_test.go @@ -726,6 +726,37 @@ func TestOpenAIGatewayServiceRecordUsage_PrefersClientRequestIDOverUpstreamReque require.Equal(t, "client:openai-client-stable-123", usageRepo.lastLog.RequestID) } +func TestOpenAIGatewayServiceRecordUsage_WSModePrefersUpstreamRequestIDOverClientRequestID(t *testing.T) { + usageRepo := &openAIRecordUsageLogRepoStub{} + billingRepo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: true}} + userRepo := &openAIRecordUsageUserRepoStub{} + subRepo := &openAIRecordUsageSubRepoStub{} + svc := newOpenAIRecordUsageServiceWithBillingRepoForTest(usageRepo, billingRepo, userRepo, subRepo, nil) + + ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "openai-ws-connection-123") + err := svc.RecordUsage(ctx, &OpenAIRecordUsageInput{ + Result: &OpenAIForwardResult{ + RequestID: "resp_openai_ws_turn_456", + OpenAIWSMode: true, + Usage: OpenAIUsage{ + InputTokens: 8, + OutputTokens: 4, + }, + Model: "gpt-5.1", + Duration: time.Second, + }, + APIKey: &APIKey{ID: 10050}, + User: &User{ID: 20050}, + Account: &Account{ID: 30050}, + }) + + require.NoError(t, err) + require.NotNil(t, billingRepo.lastCmd) + require.Equal(t, "resp_openai_ws_turn_456", billingRepo.lastCmd.RequestID) + require.NotNil(t, usageRepo.lastLog) + require.Equal(t, "resp_openai_ws_turn_456", usageRepo.lastLog.RequestID) +} + func TestOpenAIGatewayServiceRecordUsage_GeneratesRequestIDWhenAllSourcesMissing(t *testing.T) { usageRepo := &openAIRecordUsageLogRepoStub{} billingRepo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: true}} diff --git a/backend/internal/service/openai_gateway_responses_chat_fallback.go b/backend/internal/service/openai_gateway_responses_chat_fallback.go index cfab389a3151..205b27f7766f 100644 --- a/backend/internal/service/openai_gateway_responses_chat_fallback.go +++ b/backend/internal/service/openai_gateway_responses_chat_fallback.go @@ -163,7 +163,7 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions( defer func() { _ = resp.Body.Close() }() if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 10eb4eb6f7d8..042f8fb9e30b 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -47,11 +47,11 @@ const ( // codex_cli_only 拒绝时单个请求头日志长度上限（字符） codexCLIOnlyHeaderValueMaxBytes = 256 - // OpenAIParsedRequestBodyKey 缓存 handler 侧已解析的请求体，避免重复解析。 - OpenAIParsedRequestBodyKey = "openai_parsed_request_body" // OpenAI WS Mode 失败后的重连次数上限（不含首次尝试）。 // 与 Codex 客户端保持一致：失败后最多重连 5 次。 openAIWSReconnectRetryLimit = 5 + // 上游错误体只需要提取错误 JSON/日志摘要，默认 512KiB 避免错误风暴叠加大请求体。 + openAIUpstreamErrorBodyReadLimit int64 = 512 << 10 // OpenAI WS Mode 重连退避默认值（可由配置覆盖）。 openAIWSRetryBackoffInitialDefault = 120 * time.Millisecond openAIWSRetryBackoffMaxDefault = 2 * time.Second @@ -127,6 +127,19 @@ type NormalizedCodexLimits struct { Window7dMinutes *int } +func normalizeCodexFiveHourUsedPercent(raw *float64) *float64 { + if raw == nil { + return nil + } + // OpenAI's 5h Codex quota header is remaining%, despite the upstream header + // name saying "used"; the canonical codex_5h_used_percent field stores used%. + used := 100 - *raw + if used < 0 { + used = 0 + } + return &used +} + // Normalize converts primary/secondary fields to canonical 5h/7d fields. // Strategy: Compare window_minutes to determine which is 5h vs 7d. // Returns nil if snapshot is nil or has no useful data. @@ -185,7 +198,7 @@ func (s *OpenAICodexUsageSnapshot) Normalize() *NormalizedCodexLimits { // Assign values if use5hFromPrimary { - result.Used5hPercent = s.PrimaryUsedPercent + result.Used5hPercent = normalizeCodexFiveHourUsedPercent(s.PrimaryUsedPercent) result.Reset5hSeconds = s.PrimaryResetAfterSeconds result.Window5hMinutes = s.PrimaryWindowMinutes result.Used7dPercent = s.SecondaryUsedPercent @@ -195,7 +208,7 @@ func (s *OpenAICodexUsageSnapshot) Normalize() *NormalizedCodexLimits { result.Used7dPercent = s.PrimaryUsedPercent result.Reset7dSeconds = s.PrimaryResetAfterSeconds result.Window7dMinutes = s.PrimaryWindowMinutes - result.Used5hPercent = s.SecondaryUsedPercent + result.Used5hPercent = normalizeCodexFiveHourUsedPercent(s.SecondaryUsedPercent) result.Reset5hSeconds = s.SecondaryResetAfterSeconds result.Window5hMinutes = s.SecondaryWindowMinutes } @@ -248,6 +261,9 @@ type OpenAIForwardResult struct { ImageOutputSizes []string ImageSizeSource string ImageSizeBreakdown map[string]int + + wsReplayInput []json.RawMessage + wsReplayInputExists bool } type OpenAIWSRetryMetricsSnapshot struct { @@ -1295,7 +1311,7 @@ func (s *OpenAIGatewayService) SelectAccountForModel(ctx context.Context, groupI // SelectAccountForModelWithExclusions selects an account supporting the requested model while excluding specified accounts. // SelectAccountForModelWithExclusions 选择支持指定模型的账号，同时排除指定的账号。 func (s *OpenAIGatewayService) SelectAccountForModelWithExclusions(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*Account, error) { - return s.selectAccountForModelWithExclusions(ctx, groupID, sessionHash, requestedModel, excludedIDs, false, 0, "") + return s.selectAccountForModelWithExclusions(s.withOpenAIQuotaAutoPauseContext(ctx), groupID, sessionHash, requestedModel, excludedIDs, false, 0, "") } // noAvailableOpenAISelectionError builds the standard "no account available" error @@ -1332,6 +1348,17 @@ func isOpenAIAccountEligibleForRequest(ctx context.Context, account *Account, re if account == nil || !account.IsOpenAI() || !account.IsSchedulableForModelWithContext(ctx, requestedModel) { return false } + if paused, reason := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { + // Debug level: this fires per-candidate on the scheduling hot path, so Info + // would amplify into log spam once several accounts cross the threshold. + slog.Debug("account_auto_paused_by_quota", + "account_id", account.ID, + "window", reason.window, + "threshold", reason.threshold, + "utilization", reason.utilization, + ) + return false + } if requestedModel != "" && !account.IsModelSupported(requestedModel) { return false } @@ -1344,6 +1371,201 @@ func isOpenAIAccountEligibleForRequest(ctx context.Context, account *Account, re return true } +type openAIQuotaAutoPauseDecision struct { + window string + threshold float64 + utilization float64 +} + +func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account) (bool, openAIQuotaAutoPauseDecision) { + if account == nil || !account.IsOpenAI() { + return false, openAIQuotaAutoPauseDecision{} + } + // Per-account explicit-disable flags must take precedence over the global default. + // Without these, leaving the account threshold blank means "use global default", + // so an admin has no way to exempt a single account from auto-pause once a global + // default exists. The disable flag is per-window so an account can opt out of + // only 5h or only 7d auto-pause. + disabled5h := resolveAccountExtraBool(account.Extra, "auto_pause_5h_disabled") + disabled7d := resolveAccountExtraBool(account.Extra, "auto_pause_7d_disabled") + threshold5h, threshold7d := resolveOpenAIQuotaAutoPauseThresholds(ctx, account) + now := time.Now() + if !disabled5h && threshold5h > 0 { + if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "5h", now); ok && utilization >= threshold5h { + return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: threshold5h, utilization: utilization} + } + } + if !disabled7d && threshold7d > 0 { + if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "7d", now); ok && utilization >= threshold7d { + return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: threshold7d, utilization: utilization} + } + } + return false, openAIQuotaAutoPauseDecision{} +} + +// resolveAccountExtraBool reads a bool-like value from account extra, tolerating +// the few shapes JSON unmarshalling may produce (real bool, "true"/"false" +// strings, 0/1 numbers). +func resolveAccountExtraBool(extra map[string]any, key string) bool { + if len(extra) == 0 { + return false + } + value, ok := extra[key] + if !ok || value == nil { + return false + } + switch v := value.(type) { + case bool: + return v + case string: + parsed, err := strconv.ParseBool(strings.TrimSpace(v)) + return err == nil && parsed + case float64: + return v != 0 + case float32: + return v != 0 + case int: + return v != 0 + case int64: + return v != 0 + case json.Number: + if i, err := v.Int64(); err == nil { + return i != 0 + } + } + return false +} + +func resolveOpenAIQuotaAutoPauseThresholds(ctx context.Context, account *Account) (float64, float64) { + threshold5h, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_5h_threshold") + threshold7d, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_7d_threshold") + threshold5h = clamp01(threshold5h) + threshold7d = clamp01(threshold7d) + if threshold5h > 0 && threshold7d > 0 { + return threshold5h, threshold7d + } + settings := openAIQuotaAutoPauseSettingsFromContext(ctx) + if threshold5h <= 0 { + threshold5h = clamp01(settings.DefaultThreshold5h) + } + if threshold7d <= 0 { + threshold7d = clamp01(settings.DefaultThreshold7d) + } + return threshold5h, threshold7d +} + +func resolveAccountExtraNumber(extra map[string]any, keys ...string) (float64, bool) { + if len(extra) == 0 { + return 0, false + } + for _, key := range keys { + value, ok := extra[key] + if !ok || value == nil { + continue + } + switch v := value.(type) { + case float64: + return v, true + case float32: + return float64(v), true + case int: + return float64(v), true + case int64: + return float64(v), true + case json.Number: + parsed, err := v.Float64() + if err == nil { + return parsed, true + } + case string: + parsed, err := strconv.ParseFloat(strings.TrimSpace(v), 64) + if err == nil { + return parsed, true + } + } + } + return 0, false +} + +// resolveOpenAIQuotaUtilization returns the current utilization ratio (0..1) for the +// given Codex usage window. ok=false means there is no usable signal to pause on: +// either no snapshot exists, or the window has already rolled over so the cached +// percentage is stale. The stale guard matters because a paused account stops +// receiving requests, so its snapshot is never refreshed from upstream headers — +// without this check an old used_percent would keep the account paused forever even +// after the real window reset. +func resolveOpenAIQuotaUtilization(extra map[string]any, window string, now time.Time) (float64, bool) { + usedPercent := readOpenAIQuotaUsedPercent(extra, window) + if usedPercent <= 0 { + return 0, false + } + if openAIQuotaWindowReset(extra, window, now) { + return 0, false + } + return usedPercent / 100, true +} + +// openAIQuotaWindowReset reports whether the Codex usage window's reset time has +// already passed relative to now. It prefers the absolute codex__reset_at +// timestamp and falls back to codex__reset_after_seconds anchored at +// codex_usage_updated_at, mirroring AccountUsageService's window-progress logic. +func openAIQuotaWindowReset(extra map[string]any, window string, now time.Time) bool { + if len(extra) == 0 { + return false + } + if resetAtRaw, ok := extra["codex_"+window+"_reset_at"]; ok { + if resetAt, err := parseTime(fmt.Sprint(resetAtRaw)); err == nil { + return !now.Before(resetAt) + } + } + resetAfter := parseExtraInt(extra["codex_"+window+"_reset_after_seconds"]) + if resetAfter <= 0 { + return false + } + base := now + if updatedRaw, ok := extra["codex_usage_updated_at"]; ok { + if updatedAt, err := parseTime(fmt.Sprint(updatedRaw)); err == nil { + base = updatedAt + } + } + resetAt := base.Add(time.Duration(resetAfter) * time.Second) + return !now.Before(resetAt) +} + +func readOpenAIQuotaUsedPercent(extra map[string]any, window string) float64 { + if len(extra) == 0 { + return 0 + } + if value, ok := resolveAccountExtraNumber(extra, "codex_"+window+"_used_percent"); ok { + return value + } + return 0 +} + +type openAIQuotaAutoPauseCtxKey struct{} + +func withOpenAIQuotaAutoPauseSettings(ctx context.Context, settings OpsOpenAIAccountQuotaAutoPauseSettings) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, openAIQuotaAutoPauseCtxKey{}, settings) +} + +func openAIQuotaAutoPauseSettingsFromContext(ctx context.Context) OpsOpenAIAccountQuotaAutoPauseSettings { + if ctx == nil { + return OpsOpenAIAccountQuotaAutoPauseSettings{} + } + settings, _ := ctx.Value(openAIQuotaAutoPauseCtxKey{}).(OpsOpenAIAccountQuotaAutoPauseSettings) + return settings +} + +func (s *OpenAIGatewayService) withOpenAIQuotaAutoPauseContext(ctx context.Context) context.Context { + if s == nil || s.settingService == nil { + return ctx + } + return withOpenAIQuotaAutoPauseSettings(ctx, s.settingService.GetOpenAIQuotaAutoPauseSettings(ctx)) +} + // prioritizeOpenAICompactAccounts re-orders a slice so that accounts with known // compact support are tried first, followed by unknown, then explicitly unsupported. // The relative order within each tier is preserved. @@ -1592,7 +1814,7 @@ func (s *OpenAIGatewayService) isBetterAccount(candidate, current *Account) bool // SelectAccountWithLoadAwareness selects an account with load-awareness and wait plan. func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*AccountSelectionResult, error) { - return s.selectAccountWithLoadAwareness(ctx, groupID, sessionHash, requestedModel, excludedIDs, false, "") + return s.selectAccountWithLoadAwareness(s.withOpenAIQuotaAutoPauseContext(ctx), groupID, sessionHash, requestedModel, excludedIDs, false, "") } func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, requireCompact bool, requiredCapability OpenAIEndpointCapability) (*AccountSelectionResult, error) { @@ -2083,8 +2305,42 @@ func (s *OpenAIGatewayService) shouldFailoverOpenAIUpstreamResponse(statusCode i return isOpenAITransientProcessingError(statusCode, upstreamMsg, upstreamBody) } +func marshalOpenAIUpstreamJSON(v any) ([]byte, error) { + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + if err := enc.Encode(v); err != nil { + return nil, err + } + out := buf.Bytes() + if len(out) > 0 && out[len(out)-1] == '\n' { + out = out[:len(out)-1] + } + return out, nil +} + +func openAIUpstreamErrorBodyReadLimitForConfig(cfg *config.Config) int64 { + limit := openAIUpstreamErrorBodyReadLimit + if cfg != nil && cfg.Gateway.LogUpstreamErrorBody && cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { + limit = int64(cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return limit +} + +func (s *OpenAIGatewayService) readUpstreamErrorBody(resp *http.Response) []byte { + if resp == nil || resp.Body == nil { + return nil + } + cfg := (*config.Config)(nil) + if s != nil { + cfg = s.cfg + } + body, _ := io.ReadAll(io.LimitReader(resp.Body, openAIUpstreamErrorBodyReadLimitForConfig(cfg))) + return body +} + func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) if len(requestedModel) > 0 { s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0]) return @@ -2111,7 +2367,8 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } originalBody := body - reqModel, reqStream, promptCacheKey := extractOpenAIRequestMetaFromBody(body) + requestView := newOpenAIRequestView(body) + reqModel, reqStream, promptCacheKey := requestView.Model, requestView.Stream, requestView.PromptCacheKey originalModel := reqModel if account.Type == AccountTypeAPIKey && !openai_compat.ShouldUseResponsesAPI(account.Extra) { @@ -2161,172 +2418,83 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco return s.forwardOpenAIPassthrough(ctx, c, account, originalBody, reqModel, reasoningEffort, reqStream, startTime) } - reqBody, err := getOpenAIRequestBodyMap(c, body) - if err != nil { - return nil, err - } - - if v, ok := reqBody["model"].(string); ok { - reqModel = v - originalModel = reqModel - } - if v, ok := reqBody["stream"].(bool); ok { - reqStream = v - } - if promptCacheKey == "" { - if v, ok := reqBody["prompt_cache_key"].(string); ok { - promptCacheKey = strings.TrimSpace(v) - } - } - apiKey := getAPIKeyFromContext(c) - imageGenerationAllowed := GroupAllowsImageGeneration(nil) - if apiKey != nil { - imageGenerationAllowed = GroupAllowsImageGeneration(apiKey.Group) - } - codexImageGenerationBridgeEnabled := isCodexCLI && imageGenerationAllowed && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) - if IsImageGenerationIntentMap(openAIResponsesEndpoint, reqModel, reqBody) && !imageGenerationAllowed { - MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) - c.JSON(http.StatusForbidden, gin.H{ - "error": gin.H{ - "type": "permission_error", - "message": ImageGenerationPermissionMessage(), - }, - }) - return nil, errors.New("image generation disabled for group") - } - - // Track if body needs re-serialization bodyModified := false - // 单字段补丁快速路径：只要整个变更集最终可归约为同一路径的 set/delete，就避免全量 Marshal。 - patchDisabled := false - patchHasOp := false - patchDelete := false - patchPath := "" - var patchValue any - markPatchSet := func(path string, value any) { - if strings.TrimSpace(path) == "" { - patchDisabled = true - return + var reqBody map[string]any + ensureReqBody := func() (map[string]any, error) { + if requestView.HasPatches() { + patchedBody, patchErr := requestView.ApplyPatches() + if patchErr != nil { + return nil, patchErr + } + body = patchedBody + requestView = newOpenAIRequestView(body) + reqBody = nil + bodyModified = false } - if patchDisabled { - return + if reqBody != nil { + return reqBody, nil } - if !patchHasOp { - patchHasOp = true - patchDelete = false - patchPath = path - patchValue = value - return + decoded, decodeErr := requestView.Decode(c) + if decodeErr != nil { + return nil, decodeErr } - if patchDelete || patchPath != path { - patchDisabled = true + reqBody = decoded + return reqBody, nil + } + markPatchSet := func(path string, value any) { + bodyModified = true + if requestView.patchesDisabled { + if reqBody != nil { + setOpenAIRequestMapPath(reqBody, path, value) + } return } - patchValue = value + requestView.MarkPatchSet(path, value) } markPatchDelete := func(path string) { - if strings.TrimSpace(path) == "" { - patchDisabled = true - return - } - if patchDisabled { - return - } - if !patchHasOp { - patchHasOp = true - patchDelete = true - patchPath = path + bodyModified = true + if requestView.patchesDisabled { + if reqBody != nil { + deleteOpenAIRequestMapPath(reqBody, path) + } return } - if !patchDelete || patchPath != path { - patchDisabled = true - } + requestView.MarkPatchDelete(path) } disablePatch := func() { - patchDisabled = true + requestView.DisablePatches() } - - // 非透传模式下，instructions 为空时注入默认指令。 - if isInstructionsEmpty(reqBody) && !compatMessagesBridge { - reqBody["instructions"] = "You are a helpful coding assistant." - bodyModified = true - markPatchSet("instructions", "You are a helpful coding assistant.") - } - - if codexImageGenerationBridgeEnabled && ensureOpenAIResponsesImageGenerationTool(reqBody) { + markDecodedModified := func() { bodyModified = true disablePatch() - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Injected /responses image_generation tool for Codex client") } - if normalizeOpenAIResponsesImageGenerationTools(reqBody) { - bodyModified = true - disablePatch() - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized /responses image_generation tool payload") + apiKey := getAPIKeyFromContext(c) + imageGenerationAllowed := GroupAllowsImageGeneration(nil) + if apiKey != nil { + imageGenerationAllowed = GroupAllowsImageGeneration(apiKey.Group) } - if codexImageGenerationBridgeEnabled && applyCodexImageGenerationBridgeInstructions(reqBody) { - bodyModified = true - disablePatch() - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Added Codex image_generation bridge instructions") + codexImageGenerationBridgeEnabled := isCodexCLI && imageGenerationAllowed && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) + imageIntent := IsImageGenerationIntent(openAIResponsesEndpoint, reqModel, body) + if imageIntent && !imageGenerationAllowed { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) + c.JSON(http.StatusForbidden, gin.H{"error": gin.H{"type": "permission_error", "message": ImageGenerationPermissionMessage()}}) + return nil, errors.New("image generation disabled for group") + } + + instructions := gjson.GetBytes(body, "instructions") + instructionsEmpty := !instructions.Exists() || instructions.Type != gjson.String || strings.TrimSpace(instructions.String()) == "" + if instructionsEmpty && !compatMessagesBridge { + markPatchSet("instructions", "You are a helpful coding assistant.") } - // 对所有请求执行模型映射（包含 Codex CLI）。 billingModel := account.GetMappedModel(reqModel) if billingModel != reqModel { logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Model mapping applied: %s -> %s (account: %s, isCodexCLI: %v)", reqModel, billingModel, account.Name, isCodexCLI) - reqBody["model"] = billingModel - bodyModified = true + reqModel = billingModel markPatchSet("model", billingModel) } upstreamModel := billingModel - if imageGenerationAllowed && normalizeOpenAIResponsesImageOnlyModel(reqBody) { - bodyModified = true - disablePatch() - if model, ok := reqBody["model"].(string); ok { - upstreamModel = strings.TrimSpace(model) - } - logger.LegacyPrintf( - "service.openai_gateway", - "[OpenAI] Normalized /responses image-only model request inbound_model=%s image_model=%s upstream_model=%s", - reqModel, - billingModel, - upstreamModel, - ) - } - if err := validateOpenAIResponsesImageModel(reqBody, upstreamModel); err != nil { - setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") - c.JSON(http.StatusBadRequest, gin.H{ - "error": gin.H{ - "type": "invalid_request_error", - "message": err.Error(), - "param": "model", - }, - }) - return nil, err - } - if hasOpenAIImageGenerationTool(reqBody) { - logger.LegacyPrintf( - "service.openai_gateway", - "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", - reqModel, - upstreamModel, - account.Type, - ) - } - if err := validateCodexSparkInput(reqBody, upstreamModel); err != nil { - setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") - c.JSON(http.StatusBadRequest, gin.H{ - "error": gin.H{ - "type": "invalid_request_error", - "message": err.Error(), - "param": "input", - }, - }) - return nil, err - } - - // Compact-only model 映射：仅在 /responses/compact 路径生效，且优先级高于 - // OAuth 模型规范化（避免 OAuth 规范化覆盖 compact-only 自定义模型）。 isCompactRequest := isOpenAIResponsesCompactPath(c) compactMapped := false if isCompactRequest { @@ -2334,65 +2502,100 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco if compactMappedModel != "" && compactMappedModel != billingModel { compactMapped = true upstreamModel = compactMappedModel - reqBody["model"] = compactMappedModel - bodyModified = true + reqModel = compactMappedModel markPatchSet("model", compactMappedModel) logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Compact model mapping applied: %s -> %s (account: %s, isCodexCLI: %v)", billingModel, compactMappedModel, account.Name, isCodexCLI) } } - - // OpenAI OAuth 账号走 ChatGPT internal Codex endpoint，需要将模型名规范化为 - // 上游可识别的 Codex/GPT 系列。API Key 账号则应保留原始/映射后的模型名， - // 以兼容自定义 base_url 的 OpenAI-compatible 上游。 - if model, ok := reqBody["model"].(string); ok { - if !compactMapped { - upstreamModel = normalizeOpenAIModelForUpstream(account, model) - if upstreamModel != "" && upstreamModel != model { - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Upstream model resolved: %s -> %s (account: %s, type: %s, isCodexCLI: %v)", - model, upstreamModel, account.Name, account.Type, isCodexCLI) - reqBody["model"] = upstreamModel - bodyModified = true - markPatchSet("model", upstreamModel) - } + if !compactMapped { + modelForNormalize := reqModel + if modelForNormalize == "" { + modelForNormalize = requestView.Model } + upstreamModel = normalizeOpenAIModelForUpstream(account, modelForNormalize) + if upstreamModel != "" && upstreamModel != modelForNormalize { + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Upstream model resolved: %s -> %s (account: %s, type: %s, isCodexCLI: %v)", modelForNormalize, upstreamModel, account.Name, account.Type, isCodexCLI) + reqModel = upstreamModel + markPatchSet("model", upstreamModel) + } + } + if strings.TrimSpace(gjson.GetBytes(body, "reasoning.effort").String()) == "minimal" { + markPatchSet("reasoning.effort", "none") + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized reasoning.effort: minimal -> none (account: %s)", account.Name) + } + + imageIntent = imageIntent || IsImageGenerationIntent(openAIResponsesEndpoint, reqModel, nil) || isOpenAIImageGenerationModel(upstreamModel) + if imageIntent && !imageGenerationAllowed { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) + c.JSON(http.StatusForbidden, gin.H{"error": gin.H{"type": "permission_error", "message": ImageGenerationPermissionMessage()}}) + return nil, errors.New("image generation disabled for group") + } - // 移除 gpt-5.2-codex 以下的版本 verbosity 参数 - // 确保高版本模型向低版本模型映射不报错 - if !SupportsVerbosity(upstreamModel) { - if text, ok := reqBody["text"].(map[string]any); ok { - delete(text, "verbosity") + if imageGenerationAllowed && (codexImageGenerationBridgeEnabled || isOpenAIImageGenerationModel(requestView.Model) || openAIRequestBodyImageGenerationToolNeedsNormalization(body) || isOpenAIImageGenerationModel(upstreamModel)) { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + if codexImageGenerationBridgeEnabled && ensureOpenAIResponsesImageGenerationTool(decoded) { + markDecodedModified() + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Injected /responses image_generation tool for Codex client") + } + if normalizeOpenAIResponsesImageGenerationTools(decoded) { + markDecodedModified() + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized /responses image_generation tool payload") + } + if normalizeOpenAIResponsesImageOnlyModel(decoded) { + markDecodedModified() + if model, ok := decoded["model"].(string); ok { + upstreamModel = strings.TrimSpace(model) } + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized /responses image-only model request inbound_model=%s image_model=%s upstream_model=%s", requestView.Model, billingModel, upstreamModel) + } + if err := validateOpenAIResponsesImageModel(decoded, upstreamModel); err != nil { + setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"type": "invalid_request_error", "message": err.Error(), "param": "model"}}) + return nil, err + } + if hasOpenAIImageGenerationTool(decoded) { + imageIntent = true + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", requestView.Model, upstreamModel, account.Type) + } + if codexImageGenerationBridgeEnabled && applyCodexImageGenerationBridgeInstructions(decoded) { + markDecodedModified() + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Added Codex image_generation bridge instructions") } + } else if imageGenerationAllowed && imageIntent && openAIRequestBodyHasImageGenerationTool(body) { + // 完整 image_generation tool 只做 raw 计费读取，校验/桥接/旧字段迁移命中时才展开大 input map。 + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", requestView.Model, upstreamModel, account.Type) } - // 规范化 reasoning.effort 参数（minimal -> none），与上游允许值对齐。 - if reasoning, ok := reqBody["reasoning"].(map[string]any); ok { - if effort, ok := reasoning["effort"].(string); ok && effort == "minimal" { - reasoning["effort"] = "none" - bodyModified = true - markPatchSet("reasoning.effort", "none") - logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Normalized reasoning.effort: minimal -> none (account: %s)", account.Name) + if isCodexSparkModel(upstreamModel) && openAIRequestBodyMayContainImageInput(body) { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + if err := validateCodexSparkInput(decoded, upstreamModel); err != nil { + setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"type": "invalid_request_error", "message": err.Error(), "param": "input"}}) + return nil, err } } if account.Type == AccountTypeOAuth { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } codexResult := codexTransformResult{} if compatMessagesBridge { - codexResult = applyCodexOAuthTransformWithOptions(reqBody, codexOAuthTransformOptions{ - IsCodexCLI: isCodexCLI, - IsCompact: isCompactRequest, - SkipDefaultInstructions: true, - PreserveToolCallIDs: true, - }) - ensureCodexOAuthInstructionsField(reqBody) - bodyModified = true - disablePatch() + codexResult = applyCodexOAuthTransformWithOptions(decoded, codexOAuthTransformOptions{IsCodexCLI: isCodexCLI, IsCompact: isCompactRequest, SkipDefaultInstructions: true, PreserveToolCallIDs: true}) + ensureCodexOAuthInstructionsField(decoded) + markDecodedModified() } else { - codexResult = applyCodexOAuthTransform(reqBody, isCodexCLI, isCompactRequest) + codexResult = applyCodexOAuthTransform(decoded, isCodexCLI, isCompactRequest) } if codexResult.Modified { - bodyModified = true - disablePatch() + markDecodedModified() } if codexResult.NormalizedModel != "" { upstreamModel = codexResult.NormalizedModel @@ -2402,90 +2605,57 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } } - // Handle max_output_tokens based on platform and account type + if !SupportsVerbosity(upstreamModel) && gjson.GetBytes(body, "text.verbosity").Exists() { + markPatchDelete("text.verbosity") + } + if !isCodexCLI { - if maxOutputTokens, hasMaxOutputTokens := reqBody["max_output_tokens"]; hasMaxOutputTokens { + maxOutputTokens := gjson.GetBytes(body, "max_output_tokens") + if maxOutputTokens.Exists() { switch account.Platform { case PlatformOpenAI: - // For OpenAI API Key, remove max_output_tokens (not supported) - // For OpenAI OAuth (Responses API), keep it (supported) if account.Type == AccountTypeAPIKey { - delete(reqBody, "max_output_tokens") - bodyModified = true markPatchDelete("max_output_tokens") } case PlatformAnthropic: - // For Anthropic (Claude), convert to max_tokens - delete(reqBody, "max_output_tokens") - markPatchDelete("max_output_tokens") - if _, hasMaxTokens := reqBody["max_tokens"]; !hasMaxTokens { - reqBody["max_tokens"] = maxOutputTokens - disablePatch() + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + delete(decoded, "max_output_tokens") + if _, hasMaxTokens := decoded["max_tokens"]; !hasMaxTokens { + decoded["max_tokens"] = maxOutputTokens.Value() } - bodyModified = true + markDecodedModified() case PlatformGemini: - // For Gemini, remove (will be handled by Gemini-specific transform) - delete(reqBody, "max_output_tokens") - bodyModified = true markPatchDelete("max_output_tokens") default: - // For unknown platforms, remove to be safe - delete(reqBody, "max_output_tokens") - bodyModified = true markPatchDelete("max_output_tokens") } } - - // Also handle max_completion_tokens (similar logic) - if _, hasMaxCompletionTokens := reqBody["max_completion_tokens"]; hasMaxCompletionTokens { - if account.Type == AccountTypeAPIKey || account.Platform != PlatformOpenAI { - delete(reqBody, "max_completion_tokens") - bodyModified = true - markPatchDelete("max_completion_tokens") - } + if gjson.GetBytes(body, "max_completion_tokens").Exists() && (account.Type == AccountTypeAPIKey || account.Platform != PlatformOpenAI) { + markPatchDelete("max_completion_tokens") } - - // Remove unsupported fields (not supported by upstream OpenAI API) - unsupportedFields := []string{"prompt_cache_retention", "safety_identifier"} - for _, unsupportedField := range unsupportedFields { - if _, has := reqBody[unsupportedField]; has { - delete(reqBody, unsupportedField) - bodyModified = true + for _, unsupportedField := range []string{"prompt_cache_retention", "safety_identifier"} { + if gjson.GetBytes(body, unsupportedField).Exists() { markPatchDelete(unsupportedField) } } } - - // 仅在 WSv2 模式保留 previous_response_id，其他模式（HTTP/WSv1）统一过滤。 - // 注意：该规则同样适用于 Codex CLI 请求，避免 WSv1 向上游透传不支持字段。 - if wsDecision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { - if _, has := reqBody["previous_response_id"]; has { - delete(reqBody, "previous_response_id") - bodyModified = true - markPatchDelete("previous_response_id") - } + if wsDecision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 && gjson.GetBytes(body, "previous_response_id").Exists() { + markPatchDelete("previous_response_id") } - - if sanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(reqBody) { - bodyModified = true - disablePatch() + if openAIRequestBodyMayContainEmptyBase64InputImage(body) { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + if sanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(decoded) { + markDecodedModified() + } } - // Apply OpenAI fast policy (参照 Claude BetaPolicy 的 fast-mode 过滤)： - // 针对 body 的 service_tier 字段（"priority" 即 fast，"flex"），按策略 - // 执行 filter（删除字段）或 block（拒绝请求）。对 gpt-5.5 等模型屏蔽 - // fast 时在此生效。 - // - // 注意： - // 1. 此处统一使用 upstreamModel（已经过 GetMappedModel + - // normalizeOpenAIModelForUpstream + Codex OAuth normalize），与 - // chat-completions / messages 入口保持一致，避免不同入口因为模型 - // 维度不同而出现 whitelist 命中差异。 - // 2. action=pass 时也要把 raw "fast" 归一化为 "priority" 写回 body， - // 否则 native /responses 入口透传 "fast" 给上游会被拒。chat- - // completions 入口由 normalizeResponsesBodyServiceTier 完成同一 - // 行为，这里手工实现等效逻辑。 - if rawTier, ok := reqBody["service_tier"].(string); ok { + if rawTier := requestView.ServiceTier; rawTier != "" { if normTier := normalizedOpenAIServiceTierValue(rawTier); normTier != "" { action, errMsg := s.evaluateOpenAIFastPolicy(ctx, account, upstreamModel, normTier) switch action { @@ -2498,46 +2668,51 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco writeOpenAIFastPolicyBlockedResponse(c, blocked) return nil, blocked case BetaPolicyActionFilter: - delete(reqBody, "service_tier") - bodyModified = true - disablePatch() + markPatchDelete("service_tier") default: - // pass：若客户端传的是别名 "fast"，归一化为 "priority" - // 后写回 body，确保上游收到的是其能识别的规范值。 if normTier != rawTier { - reqBody["service_tier"] = normTier - bodyModified = true markPatchSet("service_tier", normTier) } } } } - if IsImageGenerationIntentMap(openAIResponsesEndpoint, reqModel, reqBody) && !imageGenerationAllowed { - MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) - c.JSON(http.StatusForbidden, gin.H{ - "error": gin.H{ - "type": "permission_error", - "message": ImageGenerationPermissionMessage(), - }, - }) - return nil, errors.New("image generation disabled for group") + if bodyModified { + if requestView.HasPatches() { + if patchedBody, patchErr := requestView.ApplyPatches(); patchErr == nil { + body = patchedBody + requestView = newOpenAIRequestView(body) + reqBody = nil + bodyModified = false + } + } + if bodyModified { + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + var marshalErr error + body, marshalErr = marshalOpenAIUpstreamJSON(decoded) + if marshalErr != nil { + return nil, fmt.Errorf("serialize request body: %w", marshalErr) + } + requestView = newOpenAIRequestView(body) + } } imageBillingModel := "" imageSizeTier := "" imageInputSize := "" - if IsImageGenerationIntentMap(openAIResponsesEndpoint, reqModel, reqBody) { + if imageIntent { + var imageCfg OpenAIResponsesImageBillingConfig var imageCfgErr error - imageCfg, imageCfgErr := resolveOpenAIResponsesImageBillingConfigDetailed(reqBody, billingModel) + if reqBody != nil { + imageCfg, imageCfgErr = resolveOpenAIResponsesImageBillingConfigDetailed(reqBody, billingModel) + } else { + imageCfg, imageCfgErr = resolveOpenAIResponsesImageBillingConfigDetailedFromBody(body, billingModel) + } if imageCfgErr != nil { setOpsUpstreamError(c, http.StatusBadRequest, imageCfgErr.Error(), "") - c.JSON(http.StatusBadRequest, gin.H{ - "error": gin.H{ - "type": "invalid_request_error", - "message": imageCfgErr.Error(), - "param": "size", - }, - }) + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"type": "invalid_request_error", "message": imageCfgErr.Error(), "param": "size"}}) return nil, imageCfgErr } imageBillingModel = imageCfg.Model @@ -2545,29 +2720,6 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco imageInputSize = imageCfg.InputSize } - // Re-serialize body only if modified - if bodyModified { - serializedByPatch := false - if !patchDisabled && patchHasOp { - var patchErr error - if patchDelete { - body, patchErr = sjson.DeleteBytes(body, patchPath) - } else { - body, patchErr = sjson.SetBytes(body, patchPath, patchValue) - } - if patchErr == nil { - serializedByPatch = true - } - } - if !serializedByPatch { - var marshalErr error - body, marshalErr = json.Marshal(reqBody) - if marshalErr != nil { - return nil, fmt.Errorf("serialize request body: %w", marshalErr) - } - } - } - // Get access token token, _, err := s.GetAccessToken(ctx, account) if err != nil { @@ -2576,12 +2728,10 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco // 命中 WS 时仅走 WebSocket Mode；不再自动回退 HTTP。 if wsDecision.Transport == OpenAIUpstreamTransportResponsesWebsocketV2 { - wsReqBody := reqBody - if len(reqBody) > 0 { - wsReqBody = make(map[string]any, len(reqBody)) - for k, v := range reqBody { - wsReqBody[k] = v - } + // WS 分支需要结构化 payload 与重连恢复，命中后再触发 full-map decode。 + wsReqBody, err := ensureReqBody() + if err != nil { + return nil, err } _, hasPreviousResponseID := wsReqBody["previous_response_id"] logOpenAIWSModeDebug( @@ -2833,7 +2983,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco // Handle error response if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) @@ -2841,8 +2991,12 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) upstreamCode := extractUpstreamErrorCode(respBody) if !httpInvalidEncryptedContentRetryTried && resp.StatusCode == http.StatusBadRequest && upstreamCode == "invalid_encrypted_content" { - if trimOpenAIEncryptedReasoningItems(reqBody) { - body, err = json.Marshal(reqBody) + decoded, decodeErr := ensureReqBody() + if decodeErr != nil { + return nil, decodeErr + } + if trimOpenAIEncryptedReasoningItems(decoded) { + body, err = marshalOpenAIUpstreamJSON(decoded) if err != nil { return nil, fmt.Errorf("serialize invalid_encrypted_content retry body: %w", err) } @@ -2883,9 +3037,10 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } defer func() { _ = resp.Body.Close() }() - reasoningEffort := extractOpenAIReasoningEffort(reqBody, originalModel) - serviceTier := extractOpenAIServiceTier(reqBody) - releaseOpenAIParsedRequestBody(c) + reasoningEffort := extractOpenAIReasoningEffortFromBody(body, originalModel) + serviceTier := extractOpenAIServiceTierFromBody(body) + // 上游接受后只保留计费需要的标量，避免响应处理期间继续保活完整 input/tools map。 + reqBody = nil // Handle normal response var usage *OpenAIUsage @@ -3350,7 +3505,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( account *Account, requestBody []byte, ) error { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) @@ -3392,7 +3547,7 @@ func (s *OpenAIGatewayService) handleErrorResponsePassthrough( account *Account, requestBody []byte, ) error { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) @@ -4085,7 +4240,7 @@ func (s *OpenAIGatewayService) handleErrorResponse( requestBody []byte, requestedModel ...string, ) (*OpenAIForwardResult, error) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) @@ -4246,7 +4401,7 @@ func (s *OpenAIGatewayService) handleCompatErrorResponse( writeError compatErrorWriter, requestedModel ...string, ) (*OpenAIForwardResult, error) { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body := s.readUpstreamErrorBody(resp) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) if upstreamMsg == "" { @@ -5611,6 +5766,11 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec durationMs := int(result.Duration.Milliseconds()) accountRateMultiplier := account.BillingRateMultiplier() requestID := resolveUsageBillingRequestID(ctx, result.RequestID) + if result.OpenAIWSMode { + if upstreamRequestID := strings.TrimSpace(result.RequestID); upstreamRequestID != "" { + requestID = upstreamRequestID + } + } // 确定 RequestedModel（渠道映射前的原始模型） requestedModel := result.Model @@ -6094,15 +6254,163 @@ func deriveOpenAIReasoningEffortFromModel(model string) string { return normalizeOpenAIReasoningEffort(parts[len(parts)-1]) } -func extractOpenAIRequestMetaFromBody(body []byte) (model string, stream bool, promptCacheKey string) { +type openAIRequestView struct { + body []byte + Model string + Stream bool + PromptCacheKey string + PreviousResponseID string + ServiceTier string + ReasoningEffort string + patches []openAIRequestPatch + patchesDisabled bool +} + +type openAIRequestPatch struct { + path string + delete bool + value any +} + +func newOpenAIRequestView(body []byte) openAIRequestView { if len(body) == 0 { - return "", false, "" + 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()), + } +} + +// Decode 保留阶段一既有 full-map 行为；后续阶段会把调用点下沉到复杂分支。 +func (v openAIRequestView) Decode(c *gin.Context) (map[string]any, error) { + return getOpenAIRequestBodyMap(c, v.body) +} + +func (v *openAIRequestView) MarkPatchSet(path string, value any) { + if v == nil || v.patchesDisabled { + return + } + path = strings.TrimSpace(path) + if !isSimpleOpenAIRequestPatchPath(path) { + v.DisablePatches() + return } + v.patches = append(v.patches, openAIRequestPatch{path: path, value: value}) +} - model = strings.TrimSpace(gjson.GetBytes(body, "model").String()) - stream = gjson.GetBytes(body, "stream").Bool() - promptCacheKey = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()) - return model, stream, promptCacheKey +func (v *openAIRequestView) MarkPatchDelete(path string) { + if v == nil || v.patchesDisabled { + return + } + path = strings.TrimSpace(path) + if !isSimpleOpenAIRequestPatchPath(path) { + v.DisablePatches() + return + } + v.patches = append(v.patches, openAIRequestPatch{path: path, delete: true}) +} + +func isSimpleOpenAIRequestPatchPath(path string) bool { + if path == "" || strings.ContainsRune(path, '\\') { + return false + } + for _, part := range strings.Split(path, ".") { + if strings.TrimSpace(part) == "" { + return false + } + } + return true +} + +func (v *openAIRequestView) DisablePatches() { + if v == nil { + return + } + v.patchesDisabled = true + v.patches = nil +} + +func (v openAIRequestView) HasPatches() bool { + return !v.patchesDisabled && len(v.patches) > 0 +} + +func (v openAIRequestView) ApplyPatches() ([]byte, error) { + if v.patchesDisabled || len(v.patches) == 0 { + return nil, errors.New("openai request patches disabled") + } + body := v.body + for _, patch := range v.patches { + var err error + if patch.delete { + body, err = sjson.DeleteBytes(body, patch.path) + } else { + body, err = sjson.SetBytes(body, patch.path, patch.value) + } + if err != nil { + return nil, err + } + } + return body, nil +} + +func setOpenAIRequestMapPath(reqBody map[string]any, path string, value any) { + path = strings.TrimSpace(path) + if reqBody == nil || path == "" { + return + } + parts := strings.Split(path, ".") + current := reqBody + for _, part := range parts[:len(parts)-1] { + part = strings.TrimSpace(part) + if part == "" { + return + } + next, _ := current[part].(map[string]any) + if next == nil { + next = map[string]any{} + current[part] = next + } + current = next + } + last := strings.TrimSpace(parts[len(parts)-1]) + if last != "" { + current[last] = value + } +} + +func deleteOpenAIRequestMapPath(reqBody map[string]any, path string) { + path = strings.TrimSpace(path) + if reqBody == nil || path == "" { + return + } + parts := strings.Split(path, ".") + current := reqBody + for _, part := range parts[:len(parts)-1] { + part = strings.TrimSpace(part) + if part == "" { + return + } + next, _ := current[part].(map[string]any) + if next == nil { + return + } + current = next + } + last := strings.TrimSpace(parts[len(parts)-1]) + if last != "" { + delete(current, last) + } +} + +func extractOpenAIRequestMetaFromBody(body []byte) (model string, stream bool, promptCacheKey string) { + view := newOpenAIRequestView(body) + return view.Model, view.Stream, view.PromptCacheKey } // normalizeOpenAIPassthroughOAuthBody 将透传 OAuth 请求体收敛为旧链路关键行为： @@ -6553,8 +6861,84 @@ func buildOpenAIFastPolicyBlockedWSEvent(err *OpenAIFastBlockedError) []byte { return payload } +func openAIRequestBodyMayContainImageInput(body []byte) bool { + if len(body) == 0 { + return false + } + input := gjson.GetBytes(body, "input") + messages := gjson.GetBytes(body, "messages.#-1") + return openAIJSONValueMayContainImageInput(input) || openAIJSONValueMayContainImageInput(messages) +} + +func openAIJSONValueMayContainImageInput(value gjson.Result) bool { + if !value.Exists() { + return false + } + if value.IsArray() { + found := false + value.ForEach(func(_, item gjson.Result) bool { + if openAIJSONValueMayContainImageInput(item) { + found = true + return false + } + return true + }) + return found + } + if value.IsObject() { + if strings.TrimSpace(value.Get("type").String()) == "input_image" || value.Get("image_url").Exists() { + return true + } + return openAIJSONValueMayContainImageInput(value.Get("content")) + } + return false +} + +func openAIRequestBodyMayContainEmptyBase64InputImage(body []byte) bool { + if len(body) == 0 || !openAIRequestBodyMayContainInputImageToken(body) { + return false + } + input := gjson.GetBytes(body, "input") + if !input.Exists() { + return false + } + return openAIJSONValueMayContainEmptyBase64InputImage(input) +} + +func openAIRequestBodyMayContainInputImageToken(body []byte) bool { + if bytes.Contains(body, []byte("input_image")) { + return true + } + // JSON 字符串任意字符都可能被 unicode escape，遇到 \u 时交给 gjson 解码后的结构扫描兜底。 + return bytes.Contains(body, []byte("\\u")) +} + +func openAIJSONValueMayContainEmptyBase64InputImage(value gjson.Result) bool { + if !value.Exists() { + return false + } + if value.IsArray() { + found := false + value.ForEach(func(_, item gjson.Result) bool { + if openAIJSONValueMayContainEmptyBase64InputImage(item) { + found = true + return false + } + return true + }) + return found + } + if value.IsObject() { + if strings.TrimSpace(value.Get("type").String()) == "input_image" && isEmptyBase64DataURI(value.Get("image_url").String()) { + return true + } + return openAIJSONValueMayContainEmptyBase64InputImage(value.Get("content")) + } + return false +} + func sanitizeEmptyBase64InputImagesInOpenAIBody(body []byte) ([]byte, bool, error) { - if len(body) == 0 || !bytes.Contains(body, []byte(`"image_url"`)) || !bytes.Contains(body, []byte(`base64,`)) { + if !openAIRequestBodyMayContainEmptyBase64InputImage(body) { return body, false, nil } @@ -6565,7 +6949,7 @@ func sanitizeEmptyBase64InputImagesInOpenAIBody(body []byte) ([]byte, bool, erro if !sanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(reqBody) { return body, false, nil } - normalized, err := json.Marshal(reqBody) + normalized, err := marshalOpenAIUpstreamJSON(reqBody) if err != nil { return body, false, fmt.Errorf("serialize sanitized request body: %w", err) } @@ -6670,32 +7054,14 @@ func isEmptyBase64DataURI(raw string) bool { return strings.TrimSpace(strings.TrimPrefix(rest, "base64,")) == "" } -func getOpenAIRequestBodyMap(c *gin.Context, body []byte) (map[string]any, error) { - if c != nil { - if cached, ok := c.Get(OpenAIParsedRequestBodyKey); ok { - if reqBody, ok := cached.(map[string]any); ok && reqBody != nil { - return reqBody, nil - } - } - } - +func getOpenAIRequestBodyMap(_ *gin.Context, body []byte) (map[string]any, error) { var reqBody map[string]any if err := json.Unmarshal(body, &reqBody); err != nil { return nil, fmt.Errorf("parse request: %w", err) } - if c != nil { - c.Set(OpenAIParsedRequestBodyKey, reqBody) - } return reqBody, nil } -func releaseOpenAIParsedRequestBody(c *gin.Context) { - if c == nil { - return - } - delete(c.Keys, OpenAIParsedRequestBodyKey) -} - func extractOpenAIReasoningEffort(reqBody map[string]any, requestedModel string) *string { if value, present := getOpenAIReasoningEffortFromReqBody(reqBody); present { if value == "" { diff --git a/backend/internal/service/openai_gateway_service_codex_snapshot_test.go b/backend/internal/service/openai_gateway_service_codex_snapshot_test.go index 654dd4cabe8c..22f5fa749248 100644 --- a/backend/internal/service/openai_gateway_service_codex_snapshot_test.go +++ b/backend/internal/service/openai_gateway_service_codex_snapshot_test.go @@ -104,6 +104,40 @@ func TestBuildCodexUsageExtraUpdates_UsesSnapshotUpdatedAt(t *testing.T) { } } +func TestBuildCodexUsageExtraUpdates_NormalizesFiveHourRemainingToUsedPercent(t *testing.T) { + primaryUsed := 93.0 + primaryReset := 86400 + primaryWindow := 10080 + secondaryRemaining := 6.0 + secondaryReset := 3600 + secondaryWindow := 300 + + snapshot := &OpenAICodexUsageSnapshot{ + PrimaryUsedPercent: &primaryUsed, + PrimaryResetAfterSeconds: &primaryReset, + PrimaryWindowMinutes: &primaryWindow, + SecondaryUsedPercent: &secondaryRemaining, + SecondaryResetAfterSeconds: &secondaryReset, + SecondaryWindowMinutes: &secondaryWindow, + UpdatedAt: "2026-05-30T07:04:09Z", + } + + updates := buildCodexUsageExtraUpdates(snapshot, time.Time{}) + if updates == nil { + t.Fatal("expected non-nil updates") + } + + if got := updates["codex_secondary_used_percent"]; got != 6.0 { + t.Fatalf("codex_secondary_used_percent = %v, want raw upstream value 6", got) + } + if got := updates["codex_5h_used_percent"]; got != 94.0 { + t.Fatalf("codex_5h_used_percent = %v, want 94", got) + } + if got := updates["codex_7d_used_percent"]; got != 93.0 { + t.Fatalf("codex_7d_used_percent = %v, want 93", got) + } +} + func TestBuildCodexUsageExtraUpdates_FallbackToNowWhenUpdatedAtInvalid(t *testing.T) { primaryUsed := 15.0 primaryReset := 30 diff --git a/backend/internal/service/openai_gateway_service_hotpath_test.go b/backend/internal/service/openai_gateway_service_hotpath_test.go index 234dee00cf86..92a0d1ac03dc 100644 --- a/backend/internal/service/openai_gateway_service_hotpath_test.go +++ b/backend/internal/service/openai_gateway_service_hotpath_test.go @@ -1,14 +1,533 @@ package service import ( + "context" "encoding/json" + "io" + "net/http" "net/http/httptest" + "strings" "testing" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" ) +func TestOpenAIRequestView_ExtractsRawScalars(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":" gpt-5 ","stream":true,"prompt_cache_key":" ses-1 ","previous_response_id":" resp-1 ","service_tier":" fast ","reasoning":{"effort":" medium "}}`)) + + require.Equal(t, "gpt-5", view.Model) + require.True(t, view.Stream) + require.Equal(t, "ses-1", view.PromptCacheKey) + require.Equal(t, "resp-1", view.PreviousResponseID) + require.Equal(t, "fast", view.ServiceTier) + require.Equal(t, "medium", view.ReasoningEffort) +} + +func TestOpenAIRequestView_DecodeKeepsFullMapBehavior(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5","stream":true,"input":[{"type":"message","content":"hi"}]}`)) + + reqBody, err := view.Decode(nil) + require.NoError(t, err) + require.Equal(t, "gpt-5", reqBody["model"]) + require.IsType(t, []any{}, reqBody["input"]) +} + +func TestOpenAIRequestView_ApplyPatches(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5","previous_response_id":"resp_1","reasoning":{"effort":"minimal"},"input":[{"type":"message","content":"hi"}]}`)) + view.MarkPatchSet("model", "gpt-5.1") + view.MarkPatchDelete("previous_response_id") + view.MarkPatchSet("reasoning.effort", "none") + + patched, err := view.ApplyPatches() + require.NoError(t, err) + require.JSONEq(t, `{"model":"gpt-5.1","reasoning":{"effort":"none"},"input":[{"type":"message","content":"hi"}]}`, string(patched)) +} + +func TestOpenAIRequestView_RejectsEscapedPatchPath(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"metadata":{"user.id":"old"}}`)) + view.MarkPatchSet(`metadata.user\.id`, "new") + + require.False(t, view.HasPatches()) + _, err := view.ApplyPatches() + require.Error(t, err) +} + +func TestOpenAIRequestView_ApplyPatchesDisabled(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`)) + view.MarkPatchSet("model", "gpt-5.1") + view.DisablePatches() + + _, err := view.ApplyPatches() + require.Error(t, err) +} + +func TestOpenAIRequestView_HasPatches(t *testing.T) { + view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`)) + require.False(t, view.HasPatches()) + + view.MarkPatchSet("model", "gpt-5.1") + require.True(t, view.HasPatches()) + + view.DisablePatches() + require.False(t, view.HasPatches()) +} + +func TestOpenAIGatewayService_Forward_HTTPPatchPathKeepsLargeInputRaw(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( + `{"usage":{"input_tokens":1,"output_tokens":2,"input_tokens_details":{"cached_tokens":0}}}`, + )), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 1, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"reasoning":{"effort":"minimal"},"input":[{"type":"message","content":[{"type":"input_text","text":"hi","nonce":9007199254740993}]}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, upstream.lastReq) + require.JSONEq(t, `{"model":"gpt-5","stream":false,"reasoning":{"effort":"none"},"instructions":"You are a helpful coding assistant.","input":[{"type":"message","content":[{"type":"input_text","text":"hi","nonce":9007199254740993}]}]}`, string(upstream.lastBody)) + require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.lastBody, "input.0.content.0.nonce").Raw) +} + +func TestOpenAIGatewayService_Forward_DecodedMutationKeepsLaterFieldDeletes(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(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 2, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5.4","stream":false,"max_completion_tokens":12,"tools":[{"type":"image_generation","format":"png"}],"input":[{"type":"message","content":"draw"}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, gjson.GetBytes(upstream.lastBody, "max_completion_tokens").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "tools.0.format").Exists()) + require.Equal(t, "png", gjson.GetBytes(upstream.lastBody, "tools.0.output_format").String()) +} + +func TestOpenAIGatewayService_Forward_MappedImageModelUsesImageGate(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(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 3, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + "model_mapping": map[string]any{"draw-alias": "gpt-image-2"}, + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Set("api_key", &APIKey{Group: &Group{AllowImageGeneration: false}}) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"draw-alias","stream":false,"input":"draw"}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.Error(t, err) + require.Nil(t, result) + require.Nil(t, upstream.lastReq) + require.Equal(t, http.StatusForbidden, rec.Code) +} + +func TestOpenAIGatewayService_Forward_TextDataImageDoesNotForceMapMarshal(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(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 4, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"input":[{"type":"message","content":[{"type":"input_text","text":"literal data:image/png;base64, only","nonce":1e1000000}]}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "1e1000000", gjson.GetBytes(upstream.lastBody, "input.0.content.0.nonce").Raw) +} + +func TestOpenAIGatewayService_Forward_ImageToolBillingDoesNotForceFullDecode(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( + `{"output":[{"id":"ig_1","type":"image_generation_call","result":"final-image"}],"usage":{"input_tokens":1,"output_tokens":2}}`, + )), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 9, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[{"type":"message","content":[{"type":"input_text","text":"draw","nonce":1e1000000}]}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "1e1000000", gjson.GetBytes(upstream.lastBody, "input.0.content.0.nonce").Raw) + require.Equal(t, 1, result.ImageCount) + require.Equal(t, "2K", result.ImageSize) + require.Equal(t, "gpt-image-2", result.BillingModel) +} + +func TestOpenAIGatewayService_Forward_ImageToolWithImageOnlyModelIsNormalized(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(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 11, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-image-2","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2"}],"input":"draw"}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String()) +} + +func TestOpenAIGatewayService_Forward_HTTPRetryRecoveryDoesNotDecodeBeforeError(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + responses: []*http.Response{ + { + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_encrypted_content","type":"invalid_request_error","message":"bad encrypted content"}}`)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 10, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"input":[{"type":"reasoning","encrypted_content":"gAAA","summary":[{"type":"summary_text","text":"keep me"}]},{"type":"message","content":[{"type":"input_text","text":"hi","nonce":9007199254740993}]}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.bodies, 2) + require.Equal(t, "gAAA", gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").String()) + require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[0], "input.1.content.0.nonce").Raw) + require.False(t, gjson.GetBytes(upstream.bodies[1], "input.0.encrypted_content").Exists()) + require.Equal(t, "summary_text", gjson.GetBytes(upstream.bodies[1], "input.0.summary.0.type").String()) +} + +func TestOpenAIGatewayService_Forward_CodexSparkRejectsEscapedInputImage(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(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 5, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5.3-codex-spark","stream":false,"input":[{"type":"input_` + "\\u0069" + `mage","file_id":"file_1"}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.Error(t, err) + require.Nil(t, result) + require.Nil(t, upstream.lastReq) + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestOpenAIGatewayService_Forward_CodexBridgeInjectionSetsImageBilling(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( + `{"output":[{"id":"ig_1","type":"image_generation_call","result":"final-image","size":"1024x1024"}],"usage":{"input_tokens":1,"output_tokens":2}}`, + )), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Gateway.ForceCodexCLI = true + cfg.Gateway.CodexImageGenerationBridgeEnabled = true + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 7, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Set("api_key", &APIKey{Group: &Group{AllowImageGeneration: true}}) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5","stream":false,"input":"draw if needed"}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, result.ImageCount) + require.Equal(t, "2K", result.ImageSize) + require.Equal(t, "gpt-image-2", result.BillingModel) +} + +func TestOpenAIGatewayService_Forward_HTTPDeletesPreviousResponseIDWhenPresent(t *testing.T) { + gin.SetMode(gin.TestMode) + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + account := &Account{ + ID: 8, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + + for _, body := range [][]byte{ + []byte(`{"model":"gpt-5","stream":false,"previous_response_id":"","input":"hi"}`), + []byte(`{"model":"gpt-5","stream":false,"previous_response_id":null,"input":"hi"}`), + } { + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, gjson.GetBytes(upstream.lastBody, "previous_response_id").Exists()) + } +} + +func TestOpenAIRequestBodyMayContainEmptyBase64InputImageSeesEscapedJSON(t *testing.T) { + body := []byte(`{"input":[{"type":"message","content":[{"type":"input_image","image_` + "\\u0075" + `rl":"data:image/png;base64` + "\\u002c" + ` "}]}]}`) + + require.True(t, openAIRequestBodyMayContainEmptyBase64InputImage(body)) +} + +func TestOpenAIRequestBodyMayContainEmptyBase64InputImageSeesEscapedImageType(t *testing.T) { + body := []byte(`{"input":[{"type":"message","content":[{"type":"input_` + "\\u0069" + `mage","image_url":"data:image/png;base64, "}]}]}`) + + require.True(t, openAIRequestBodyMayContainEmptyBase64InputImage(body)) +} + +func TestOpenAIRequestBodyMayContainEmptyBase64InputImageSeesEscapedInputPrefix(t *testing.T) { + body := []byte(`{"input":[{"type":"message","content":[{"type":"inp` + "\\u0075" + `t_image","image_url":"data:image/png;base64, "}]}]}`) + + require.True(t, openAIRequestBodyMayContainEmptyBase64InputImage(body)) +} + +func TestOpenAIGatewayService_Forward_ImageOnlyModelKeepsSupportedVerbosity(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(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 6, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com", + }, + Extra: map[string]any{"use_responses_api": true}, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-image-2","stream":false,"text":{"verbosity":"low"},"input":"draw"}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "low", gjson.GetBytes(upstream.lastBody, "text.verbosity").String()) + require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String()) +} + func TestExtractOpenAIRequestMetaFromBody(t *testing.T) { tests := []struct { name string @@ -106,26 +625,13 @@ func TestExtractOpenAIReasoningEffortFromBody(t *testing.T) { } } -func TestGetOpenAIRequestBodyMap_UsesContextCache(t *testing.T) { - gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - - cached := map[string]any{"model": "cached-model", "stream": true} - c.Set(OpenAIParsedRequestBodyKey, cached) - - got, err := getOpenAIRequestBodyMap(c, []byte(`{invalid-json`)) - require.NoError(t, err) - require.Equal(t, cached, got) -} - -func TestGetOpenAIRequestBodyMap_ParseErrorWithoutCache(t *testing.T) { +func TestGetOpenAIRequestBodyMap_ParseError(t *testing.T) { _, err := getOpenAIRequestBodyMap(nil, []byte(`{invalid-json`)) require.Error(t, err) require.Contains(t, err.Error(), "parse request") } -func TestGetOpenAIRequestBodyMap_WriteBackContextCache(t *testing.T) { +func TestGetOpenAIRequestBodyMap_DoesNotWriteContextCache(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) @@ -133,12 +639,7 @@ func TestGetOpenAIRequestBodyMap_WriteBackContextCache(t *testing.T) { got, err := getOpenAIRequestBodyMap(c, []byte(`{"model":"gpt-5","stream":true}`)) require.NoError(t, err) require.Equal(t, "gpt-5", got["model"]) - - cached, ok := c.Get(OpenAIParsedRequestBodyKey) - require.True(t, ok) - cachedMap, ok := cached.(map[string]any) - require.True(t, ok) - require.Equal(t, got, cachedMap) + require.Empty(t, c.Keys) } func TestSanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(t *testing.T) { diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index b2bfa6d0c8a1..c4fc7d317a83 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -1774,7 +1774,7 @@ func TestOpenAIUpdateCodexUsageSnapshotFromHeaders(t *testing.T) { select { case updates := <-repo.updateExtraCalls: - require.Equal(t, 12.0, updates["codex_5h_used_percent"]) + require.Equal(t, 88.0, updates["codex_5h_used_percent"]) require.Equal(t, 34.0, updates["codex_7d_used_percent"]) require.Equal(t, 600, updates["codex_5h_reset_after_seconds"]) require.Equal(t, 86400, updates["codex_7d_reset_after_seconds"]) diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index 1bcd947c4495..beb34780f3bf 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -622,7 +622,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( return nil, fmt.Errorf("upstream request failed: %s", safeErr) } if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) @@ -1276,21 +1276,22 @@ func resolveOpenAIImageBytes( headers http.Header, conversationID string, pointer openAIImagePointerInfo, + errorBodyReadLimit int64, ) ([]byte, error) { if normalized := normalizeOpenAIImageBase64(pointer.B64JSON); normalized != "" { return base64.StdEncoding.DecodeString(normalized) } if downloadURL := strings.TrimSpace(pointer.DownloadURL); downloadURL != "" { - return downloadOpenAIImageBytes(ctx, client, headers, downloadURL) + return downloadOpenAIImageBytes(ctx, client, headers, downloadURL, errorBodyReadLimit) } if strings.TrimSpace(pointer.Pointer) == "" { return nil, fmt.Errorf("image asset is missing pointer, url, and base64 data") } - downloadURL, err := fetchOpenAIImageDownloadURL(ctx, client, headers, conversationID, pointer.Pointer) + downloadURL, err := fetchOpenAIImageDownloadURL(ctx, client, headers, conversationID, pointer.Pointer, errorBodyReadLimit) if err != nil { return nil, err } - return downloadOpenAIImageBytes(ctx, client, headers, downloadURL) + return downloadOpenAIImageBytes(ctx, client, headers, downloadURL, errorBodyReadLimit) } func normalizeOpenAIImageBase64(raw string) string { @@ -1395,6 +1396,7 @@ func fetchOpenAIImageDownloadURL( headers http.Header, conversationID string, pointer string, + errorBodyReadLimit int64, ) (string, error) { url := "" allowConversationRetry := false @@ -1425,7 +1427,7 @@ func fetchOpenAIImageDownloadURL( } else if resp.IsSuccessState() && strings.TrimSpace(result.DownloadURL) != "" { return strings.TrimSpace(result.DownloadURL), nil } else { - statusErr := newOpenAIImageStatusError(resp, "fetch image download url failed") + statusErr := newOpenAIImageStatusError(resp, "fetch image download url failed", errorBodyReadLimit) if !allowConversationRetry || !isOpenAIImageTransientConversationNotFoundError(statusErr) { return "", statusErr } @@ -1450,7 +1452,7 @@ func fetchOpenAIImageDownloadURL( return "", lastErr } -func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers http.Header, downloadURL string) ([]byte, error) { +func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers http.Header, downloadURL string, errorBodyReadLimit int64) ([]byte, error) { request := client.R(). SetContext(ctx). DisableAutoReadResponse() @@ -1478,7 +1480,7 @@ func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers h } }() if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return nil, newOpenAIImageStatusError(resp, "download image bytes failed") + return nil, newOpenAIImageStatusError(resp, "download image bytes failed", errorBodyReadLimit) } return io.ReadAll(io.LimitReader(resp.Body, openAIImageMaxDownloadBytes)) } @@ -1505,7 +1507,7 @@ func (e *openAIImageStatusError) Error() string { return "openai image backend request failed" } -func newOpenAIImageStatusError(resp *req.Response, fallback string) error { +func newOpenAIImageStatusError(resp *req.Response, fallback string, errorBodyReadLimit int64) error { if resp == nil { if strings.TrimSpace(fallback) == "" { fallback = "openai image backend request failed" @@ -1526,7 +1528,10 @@ func newOpenAIImageStatusError(resp *req.Response, fallback string) error { requestURL = resp.Request.URL.String() } if resp.Body != nil { - body, _ = io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + if errorBodyReadLimit <= 0 { + errorBodyReadLimit = openAIUpstreamErrorBodyReadLimit + } + body, _ = io.ReadAll(io.LimitReader(resp.Body, errorBodyReadLimit)) _ = resp.Body.Close() } } diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 849ad7920c13..db9c7b1627f5 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -1172,7 +1172,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( return nil, fmt.Errorf("upstream request failed: %s", safeErr) } if resp.StatusCode >= 400 { - respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() resp.Body = io.NopCloser(bytes.NewReader(respBody)) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) diff --git a/backend/internal/service/openai_images_test.go b/backend/internal/service/openai_images_test.go index a87e96c10778..c3efdc931aa8 100644 --- a/backend/internal/service/openai_images_test.go +++ b/backend/internal/service/openai_images_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "errors" + "fmt" "io" "mime/multipart" "net/http" @@ -14,6 +15,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/gin-gonic/gin" + "github.com/imroc/req/v3" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" ) @@ -398,11 +400,38 @@ func TestCollectOpenAIImagePointers_RecognizesDirectAssets(t *testing.T) { func TestResolveOpenAIImageBytes_PrefersInlineBase64(t *testing.T) { data, err := resolveOpenAIImageBytes(context.Background(), nil, nil, "", openAIImagePointerInfo{ B64JSON: "data:image/png;base64,QUJD", - }) + }, openAIUpstreamErrorBodyReadLimit) require.NoError(t, err) require.Equal(t, []byte("ABC"), data) } +func TestNewOpenAIImageStatusError_UsesProvidedReadLimit(t *testing.T) { + padding := strings.Repeat("x", int(openAIUpstreamErrorBodyReadLimit)+1024) + body := fmt.Sprintf(`{"error":{"padding":"%s","message":"diagnostic-marker"}}`, padding) + resp := &req.Response{Response: &http.Response{ + StatusCode: http.StatusBadGateway, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(body)), + }} + + err := newOpenAIImageStatusError(resp, "download image bytes failed", int64(len(body))) + require.Error(t, err) + require.Equal(t, "diagnostic-marker", err.Error()) + + var statusErr *openAIImageStatusError + require.ErrorAs(t, err, &statusErr) + require.Len(t, statusErr.ResponseBody, len(body)) +} + +func TestOpenAIUpstreamErrorBodyReadLimitForConfig_RespectsDiagnosticLimit(t *testing.T) { + cfg := &config.Config{Gateway: config.GatewayConfig{ + LogUpstreamErrorBody: true, + LogUpstreamErrorBodyMaxBytes: int(openAIUpstreamErrorBodyReadLimit) + 1024, + }} + + require.Equal(t, int64(cfg.Gateway.LogUpstreamErrorBodyMaxBytes), openAIUpstreamErrorBodyReadLimitForConfig(cfg)) +} + func TestAccountSupportsOpenAIImageCapability_OAuthSupportsNative(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, diff --git a/backend/internal/service/openai_oauth_passthrough_test.go b/backend/internal/service/openai_oauth_passthrough_test.go index 398cbb850b48..2710c696e69a 100644 --- a/backend/internal/service/openai_oauth_passthrough_test.go +++ b/backend/internal/service/openai_oauth_passthrough_test.go @@ -14,6 +14,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -101,6 +102,57 @@ func TestOpenAIGatewayService_ResponsesUnknownModelDoesNotFallbackToGPT54(t *tes require.True(t, rec.Code >= http.StatusBadRequest) } +func TestOpenAIGatewayService_NativeResponsesBodyModificationPreservesHTMLChars(t *testing.T) { + gin.SetMode(gin.TestMode) + + payloadText := strings.Repeat(`&value`, 128) + originalBody := []byte(fmt.Sprintf(`{"model":"gpt-5.5","stream":false,"max_output_tokens":100,"previous_response_id":"resp_prev","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":%q}]}]}`, payloadText)) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(originalBody)) + c.Request.Header.Set("Content-Type", "application/json") + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid_native_reencode"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop after capture"}}`)), + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{ + Enabled: false, + AllowInsecureHTTP: true, + }}}, + httpUpstream: upstream, + } + account := &Account{ + ID: 456, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "http://upstream.example", + }, + Extra: map[string]any{ + openai_compat.ExtraKeyResponsesMode: string(openai_compat.ResponsesSupportModeAuto), + openai_compat.ExtraKeyResponsesSupported: true, + }, + Status: StatusActive, + Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, originalBody) + require.Error(t, err) + require.Nil(t, result) + require.NotNil(t, upstream.lastReq) + require.Equal(t, "http://upstream.example/v1/responses", upstream.lastReq.URL.String()) + require.Contains(t, string(upstream.lastBody), payloadText) + require.NotContains(t, string(upstream.lastBody), `\\u003c`) + require.NotContains(t, string(upstream.lastBody), `\\u003e`) + require.NotContains(t, string(upstream.lastBody), `\\u0026`) +} + func TestOpenAIGatewayService_OAuthMessagesBridgeDoesNotInjectDefaultInstructions(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_oauth_service.go b/backend/internal/service/openai_oauth_service.go index dc094d43ce2c..0ee357a97c05 100644 --- a/backend/internal/service/openai_oauth_service.go +++ b/backend/internal/service/openai_oauth_service.go @@ -278,11 +278,29 @@ func (s *OpenAIOAuthService) enrichTokenInfo(ctx context.Context, tokenInfo *Ope tokenInfo.Email = info.Email } } + if strings.TrimSpace(tokenInfo.SubscriptionExpiresAt) == "" { + if expiresAt := fetchChatGPTSubscriptionExpiresAt(ctx, s.privacyClientFactory, tokenInfo.AccessToken, proxyURL, resolveChatGPTSubscriptionAccountID(tokenInfo, orgID)); expiresAt != "" { + tokenInfo.SubscriptionExpiresAt = expiresAt + } + } // 尝试设置隐私（关闭训练数据共享），best-effort tokenInfo.PrivacyMode = disableOpenAITraining(ctx, s.privacyClientFactory, tokenInfo.AccessToken, proxyURL) } +func resolveChatGPTSubscriptionAccountID(tokenInfo *OpenAITokenInfo, orgID string) string { + for _, candidate := range []string{ + tokenInfo.ChatGPTAccountID, + tokenInfo.OrganizationID, + orgID, + } { + if trimmed := strings.TrimSpace(candidate); trimmed != "" { + return trimmed + } + } + return "" +} + // RefreshAccountToken refreshes token for an OpenAI OAuth account func (s *OpenAIOAuthService) RefreshAccountToken(ctx context.Context, account *Account) (*OpenAITokenInfo, error) { if account.Platform != PlatformOpenAI { @@ -292,38 +310,40 @@ func (s *OpenAIOAuthService) RefreshAccountToken(ctx context.Context, account *A return nil, infraerrors.New(http.StatusBadRequest, "OPENAI_OAUTH_INVALID_ACCOUNT_TYPE", "account is not an OAuth account") } + var proxyURL string + if account.ProxyID != nil { + proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID) + if err == nil && proxy != nil { + proxyURL = proxy.URL() + } + } + refreshToken := account.GetCredential("refresh_token") if refreshToken == "" { accessToken := account.GetCredential("access_token") if accessToken != "" { tokenInfo := &OpenAITokenInfo{ - AccessToken: accessToken, - RefreshToken: "", - IDToken: account.GetCredential("id_token"), - ClientID: account.GetCredential("client_id"), - Email: account.GetCredential("email"), - ChatGPTAccountID: account.GetCredential("chatgpt_account_id"), - ChatGPTUserID: account.GetCredential("chatgpt_user_id"), - OrganizationID: account.GetCredential("organization_id"), - PlanType: account.GetCredential("plan_type"), + AccessToken: accessToken, + RefreshToken: "", + IDToken: account.GetCredential("id_token"), + ClientID: account.GetCredential("client_id"), + Email: account.GetCredential("email"), + ChatGPTAccountID: account.GetCredential("chatgpt_account_id"), + ChatGPTUserID: account.GetCredential("chatgpt_user_id"), + OrganizationID: account.GetCredential("organization_id"), + PlanType: account.GetCredential("plan_type"), + SubscriptionExpiresAt: account.GetCredential("subscription_expires_at"), } if expiresAt := account.GetCredentialAsTime("expires_at"); expiresAt != nil { tokenInfo.ExpiresAt = expiresAt.Unix() tokenInfo.ExpiresIn = int64(time.Until(*expiresAt).Seconds()) } + s.enrichTokenInfo(ctx, tokenInfo, proxyURL) return tokenInfo, nil } return nil, infraerrors.New(http.StatusBadRequest, "OPENAI_OAUTH_NO_REFRESH_TOKEN", "no refresh token available") } - var proxyURL string - if account.ProxyID != nil { - proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID) - if err == nil && proxy != nil { - proxyURL = proxy.URL() - } - } - clientID := account.GetCredential("client_id") return s.RefreshTokenWithClientID(ctx, refreshToken, proxyURL, clientID) } diff --git a/backend/internal/service/openai_oauth_service_refresh_test.go b/backend/internal/service/openai_oauth_service_refresh_test.go index 84b68ea643a6..75588c8db4d0 100644 --- a/backend/internal/service/openai_oauth_service_refresh_test.go +++ b/backend/internal/service/openai_oauth_service_refresh_test.go @@ -8,6 +8,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/openai" + "github.com/imroc/req/v3" "github.com/stretchr/testify/require" ) @@ -32,6 +33,11 @@ func (s *openaiOAuthClientRefreshStub) RefreshTokenWithClientID(ctx context.Cont func TestOpenAIOAuthService_RefreshAccountToken_NoRefreshTokenUsesExistingAccessToken(t *testing.T) { client := &openaiOAuthClientRefreshStub{} svc := NewOpenAIOAuthService(nil, client) + var privacyClientCalls int32 + svc.SetPrivacyClientFactory(func(proxyURL string) (*req.Client, error) { + atomic.AddInt32(&privacyClientCalls, 1) + return nil, errors.New("stop before request") + }) expiresAt := time.Now().Add(30 * time.Minute).UTC().Format(time.RFC3339) account := &Account{ @@ -51,6 +57,7 @@ func TestOpenAIOAuthService_RefreshAccountToken_NoRefreshTokenUsesExistingAccess require.Equal(t, "existing-access-token", info.AccessToken) require.Equal(t, "client-id-1", info.ClientID) require.Zero(t, atomic.LoadInt32(&client.refreshCalls), "existing access token should be reused without calling refresh") + require.Positive(t, atomic.LoadInt32(&privacyClientCalls), "existing access token should still run enrichment") } func TestOpenAITokenRefresher_NeedsRefresh_SkipsAccountWithoutRefreshToken(t *testing.T) { diff --git a/backend/internal/service/openai_privacy_service.go b/backend/internal/service/openai_privacy_service.go index da6dbefc93f6..99cbb7267e46 100644 --- a/backend/internal/service/openai_privacy_service.go +++ b/backend/internal/service/openai_privacy_service.go @@ -95,6 +95,8 @@ type ChatGPTAccountInfo struct { const chatGPTAccountsCheckURL = "https://chatgpt.com/backend-api/accounts/check/v4-2023-04-27" +var chatGPTSubscriptionsURL = "https://chatgpt.com/backend-api/subscriptions" + // fetchChatGPTAccountInfo calls ChatGPT backend-api to get account info (plan_type, etc.). // Used as fallback when id_token doesn't contain these fields (e.g., Mobile RT). // orgID is used to match the correct account when multiple accounts exist (e.g., personal + team). @@ -199,6 +201,62 @@ func fetchChatGPTAccountInfo(ctx context.Context, clientFactory PrivacyClientFac return info } +// fetchChatGPTSubscriptionExpiresAt reads the lightweight subscription endpoint used by +// ChatGPT/Codex clients. Some Plus accounts no longer expose entitlement.expires_at in +// accounts/check, but this endpoint still returns active_until. +func fetchChatGPTSubscriptionExpiresAt(ctx context.Context, clientFactory PrivacyClientFactory, accessToken, proxyURL, accountID string) string { + accountID = strings.TrimSpace(accountID) + if accessToken == "" || accountID == "" || clientFactory == nil { + return "" + } + + ctx, cancel := context.WithTimeout(ctx, 15*time.Second) + defer cancel() + + client, err := clientFactory(proxyURL) + if err != nil { + slog.Debug("chatgpt_subscription_client_error", "error", err.Error()) + return "" + } + + var result struct { + PlanType string `json:"plan_type"` + ActiveUntil string `json:"active_until"` + WillRenew bool `json:"will_renew"` + ID string `json:"id"` + } + resp, err := client.R(). + SetContext(ctx). + SetHeader("Authorization", "Bearer "+accessToken). + SetHeader("Origin", "https://chatgpt.com"). + SetHeader("Referer", "https://chatgpt.com/"). + SetHeader("Accept", "application/json"). + SetSuccessResult(&result). + SetQueryParam("account_id", accountID). + Get(chatGPTSubscriptionsURL) + if err != nil { + slog.Debug("chatgpt_subscription_request_error", "error", err.Error()) + return "" + } + if !resp.IsSuccessState() { + slog.Debug("chatgpt_subscription_failed", "status", resp.StatusCode, "body", truncate(resp.String(), 200)) + return "" + } + + activeUntil := strings.TrimSpace(result.ActiveUntil) + if activeUntil == "" { + slog.Debug("chatgpt_subscription_no_active_until", "plan_type", result.PlanType, "has_subscription_id", strings.TrimSpace(result.ID) != "", "will_renew", result.WillRenew) + return "" + } + if _, err := time.Parse(time.RFC3339, activeUntil); err != nil { + slog.Debug("chatgpt_subscription_bad_active_until", "active_until", activeUntil, "error", err.Error()) + return "" + } + + slog.Info("chatgpt_subscription_success", "plan_type", result.PlanType, "subscription_expires_at", activeUntil, "account_id", accountID) + return activeUntil +} + // fillAccountInfo 从单个 account 对象中提取 plan_type 和 subscription_expires_at func fillAccountInfo(info *ChatGPTAccountInfo, acct map[string]any) { info.PlanType = extractPlanType(acct) diff --git a/backend/internal/service/openai_subscription_test.go b/backend/internal/service/openai_subscription_test.go new file mode 100644 index 000000000000..89df54dbacb3 --- /dev/null +++ b/backend/internal/service/openai_subscription_test.go @@ -0,0 +1,42 @@ +package service + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/imroc/req/v3" + "github.com/stretchr/testify/require" +) + +func TestFetchChatGPTSubscriptionExpiresAt(t *testing.T) { + const wantExpiresAt = "2026-06-10T02:52:15Z" + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/backend-api/subscriptions", r.URL.Path) + require.Equal(t, "acc_123", r.URL.Query().Get("account_id")) + require.Equal(t, "Bearer access-token", r.Header.Get("Authorization")) + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "plan_type": "plus", + "active_until": wantExpiresAt, + "will_renew": true, + "id": "sub_123", + }) + })) + defer server.Close() + + oldURL := chatGPTSubscriptionsURL + chatGPTSubscriptionsURL = server.URL + "/backend-api/subscriptions" + t.Cleanup(func() { chatGPTSubscriptionsURL = oldURL }) + + got := fetchChatGPTSubscriptionExpiresAt(context.Background(), func(proxyURL string) (*req.Client, error) { + return req.C().SetTimeout(5 * time.Second), nil + }, "access-token", "", "acc_123") + + require.Equal(t, wantExpiresAt, got) +} diff --git a/backend/internal/service/openai_tool_continuation.go b/backend/internal/service/openai_tool_continuation.go index 7d503f5a457b..6515c0c4e567 100644 --- a/backend/internal/service/openai_tool_continuation.go +++ b/backend/internal/service/openai_tool_continuation.go @@ -1,6 +1,10 @@ package service -import "strings" +import ( + "strings" + + "github.com/tidwall/gjson" +) // ToolContinuationSignals 聚合工具续链相关信号，避免重复遍历 input。 type ToolContinuationSignals struct { @@ -150,6 +154,67 @@ func AnalyzeToolContinuationSignals(reqBody map[string]any) ToolContinuationSign return signals } +// ValidateFunctionCallOutputContextBytes 基于 raw JSON 校验工具输出续链，避免 handler 预校验阶段全量解码大 input。 +func ValidateFunctionCallOutputContextBytes(body []byte) FunctionCallOutputValidation { + result := FunctionCallOutputValidation{} + if len(body) == 0 { + return result + } + // handler 热路径只读扫描 input，避免 GetBytes 为大 Responses body 复制整段 JSON。 + input := parseRawJSONView(body).Get("input") + if !input.IsArray() { + return result + } + + var callIDs map[string]struct{} + var referenceIDs map[string]struct{} + input.ForEach(func(_, item gjson.Result) bool { + if !item.IsObject() { + return true + } + itemType := item.Get("type").String() + switch { + case isCodexToolCallOutputItemType(itemType): + result.HasFunctionCallOutput = true + callID := strings.TrimSpace(item.Get("call_id").String()) + if callID == "" { + result.HasFunctionCallOutputMissingCallID = true + return true + } + if callIDs == nil { + callIDs = make(map[string]struct{}) + } + callIDs[callID] = struct{}{} + case isCodexToolCallContextItemType(itemType): + if strings.TrimSpace(item.Get("call_id").String()) != "" { + result.HasToolCallContext = true + } + case itemType == "item_reference": + idValue := strings.TrimSpace(item.Get("id").String()) + if idValue == "" { + return true + } + if referenceIDs == nil { + referenceIDs = make(map[string]struct{}) + } + referenceIDs[idValue] = struct{}{} + } + return !result.HasFunctionCallOutput || !result.HasToolCallContext + }) + if !result.HasFunctionCallOutput || result.HasToolCallContext || len(callIDs) == 0 || len(referenceIDs) == 0 { + return result + } + allReferenced := true + for callID := range callIDs { + if _, ok := referenceIDs[callID]; !ok { + allReferenced = false + break + } + } + result.HasItemReferenceForAllCallIDs = allReferenced + return result +} + // ValidateFunctionCallOutputContext 为 handler 提供低开销校验结果： // 1) 无工具输出直接返回 // 2) 若已存在工具调用上下文则提前返回 diff --git a/backend/internal/service/openai_tool_continuation_test.go b/backend/internal/service/openai_tool_continuation_test.go index 0e0552f654b5..4610652b6c2c 100644 --- a/backend/internal/service/openai_tool_continuation_test.go +++ b/backend/internal/service/openai_tool_continuation_test.go @@ -1,6 +1,7 @@ package service import ( + "encoding/json" "testing" "github.com/stretchr/testify/require" @@ -118,3 +119,68 @@ func TestHasItemReferenceForCallIDs(t *testing.T) { require.True(t, HasItemReferenceForCallIDs(req, []string{"call_1", "call_2"})) require.False(t, HasItemReferenceForCallIDs(req, []string{"call_1", "call_3"})) } + +func TestValidateFunctionCallOutputContextBytesMatchesMapValidation(t *testing.T) { + // handler 预校验走 raw JSON 扫描，语义必须与 service 内部 map 校验保持一致。 + cases := []struct { + name string + body map[string]any + }{ + { + name: "no_input", + body: map[string]any{"model": "gpt-5.4"}, + }, + { + name: "missing_call_id", + body: map[string]any{"input": []any{map[string]any{"type": "function_call_output"}}}, + }, + { + name: "call_id_without_reference", + body: map[string]any{"input": []any{map[string]any{"type": "function_call_output", "call_id": "call_1"}}}, + }, + { + name: "matching_reference", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_1"}, + map[string]any{"type": "item_reference", "id": "call_1"}, + }}, + }, + { + name: "partial_reference", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_1"}, + map[string]any{"type": "tool_search_output", "call_id": "call_2"}, + map[string]any{"type": "item_reference", "id": "call_1"}, + }}, + }, + { + name: "tool_context", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_1"}, + map[string]any{"type": "function_call", "call_id": "call_1"}, + }}, + }, + { + name: "all_codex_tool_outputs", + body: map[string]any{"input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_function"}, + map[string]any{"type": "tool_search_output", "call_id": "call_search"}, + map[string]any{"type": "custom_tool_call_output", "call_id": "call_custom"}, + map[string]any{"type": "mcp_tool_call_output", "call_id": "call_mcp"}, + map[string]any{"type": "item_reference", "id": "call_function"}, + map[string]any{"type": "item_reference", "id": "call_search"}, + map[string]any{"type": "item_reference", "id": "call_custom"}, + map[string]any{"type": "item_reference", "id": "call_mcp"}, + }}, + }, + } + + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + bodyBytes, err := json.Marshal(tt.body) + require.NoError(t, err) + + require.Equal(t, ValidateFunctionCallOutputContext(tt.body), ValidateFunctionCallOutputContextBytes(bodyBytes)) + }) + } +} diff --git a/backend/internal/service/openai_ws_account_sticky_test.go b/backend/internal/service/openai_ws_account_sticky_test.go index c8b28a465b43..6fc44298a400 100644 --- a/backend/internal/service/openai_ws_account_sticky_test.go +++ b/backend/internal/service/openai_ws_account_sticky_test.go @@ -48,6 +48,46 @@ func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_Hit(t *testing.T } } +func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_QuotaAutoPausedMiss(t *testing.T) { + ctx := context.Background() + groupID := int64(23) + account := Account{ + ID: 77, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 2, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + "codex_5h_used_percent": 96.0, + "auto_pause_5h_threshold": 0.95, + }, + } + cache := &stubGatewayCache{} + store := NewOpenAIWSStateStore(cache) + cfg := newOpenAIWSV2TestConfig() + svc := &OpenAIGatewayService{ + accountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}, + cache: cache, + cfg: cfg, + concurrencyService: NewConcurrencyService(stubConcurrencyCache{}), + openaiWSStateStore: store, + } + + require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_quota", account.ID, time.Hour)) + + selection, err := svc.SelectAccountByPreviousResponseID(ctx, &groupID, "resp_prev_quota", "gpt-5.1", nil, false) + require.NoError(t, err) + require.Nil(t, selection, "超过 5h 配额阈值的账号不应继续命中 previous_response_id 粘连") + + // Auto-pause is transient, so the binding is preserved: the chain can resume on the + // same account once the quota window resets. + boundAccountID, getErr := store.GetResponseAccount(ctx, groupID, "resp_prev_quota") + require.NoError(t, getErr) + require.Equal(t, account.ID, boundAccountID) +} + func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_RateLimitedMiss(t *testing.T) { ctx := context.Background() groupID := int64(23) diff --git a/backend/internal/service/openai_ws_forwarder.go b/backend/internal/service/openai_ws_forwarder.go index 4745e5bca88c..ddeb1a83b66d 100644 --- a/backend/internal/service/openai_ws_forwarder.go +++ b/backend/internal/service/openai_ws_forwarder.go @@ -1614,6 +1614,38 @@ func openAIWSRawItemsHasFunctionCallOutput(items []json.RawMessage) bool { return false } +func openAIWSRawItemsHaveToolCallContextForOutputs(items []json.RawMessage) bool { + if len(items) == 0 { + return false + } + contextCallIDs := make(map[string]struct{}) + outputCallIDs := make(map[string]struct{}) + for _, item := range items { + itemType := gjson.GetBytes(item, "type").String() + callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String()) + switch { + case isCodexToolCallContextItemType(itemType): + if callID != "" { + contextCallIDs[callID] = struct{}{} + } + case isCodexToolCallOutputItemType(itemType): + if callID == "" { + return false + } + outputCallIDs[callID] = struct{}{} + } + } + if len(outputCallIDs) == 0 || len(contextCallIDs) == 0 { + return false + } + for callID := range outputCallIDs { + if _, ok := contextCallIDs[callID]; !ok { + return false + } + } + return true +} + func openAIWSRawPayloadHasToolCallOutput(payload []byte) bool { if len(payload) == 0 { return false @@ -2604,6 +2636,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( wsPath = normalizeOpenAIWSLogValue(parsedURL.Path) } debugEnabled := isOpenAIWSModeDebugEnabled() + isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator")) || (s.cfg != nil && s.cfg.Gateway.ForceCodexCLI) type openAIWSClientPayload struct { payloadRaw []byte @@ -2713,6 +2746,34 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( } normalized = next } + apiKey := getAPIKeyFromContext(c) + imageGenerationAllowed := GroupAllowsImageGeneration(apiKeyGroup(apiKey)) + codexBridgeEnabled := isCodexCLI && imageGenerationAllowed && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) + if codexBridgeEnabled { + payloadMap := make(map[string]any) + if err := json.Unmarshal(normalized, &payloadMap); err != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", err) + } + bridgeModified := false + if ensureOpenAIResponsesImageGenerationTool(payloadMap) { + bridgeModified = true + logOpenAIWSModeInfo("ingress_ws_codex_image_tool_injected account_id=%d", account.ID) + } + if normalizeOpenAIResponsesImageGenerationTools(payloadMap) { + bridgeModified = true + } + if applyCodexImageGenerationBridgeInstructions(payloadMap) { + bridgeModified = true + logOpenAIWSModeInfo("ingress_ws_codex_image_bridge_instructions_added account_id=%d", account.ID) + } + if bridgeModified { + rebuilt, marshalErr := json.Marshal(payloadMap) + if marshalErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", marshalErr) + } + normalized = rebuilt + } + } upstreamModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)) if modelMissing || upstreamModel != originalModel { next, setErr := applyPayloadMutation(normalized, "model", upstreamModel) @@ -2722,7 +2783,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( normalized = next } imageIntent := IsImageGenerationIntent(openAIResponsesEndpoint, originalModel, normalized) - if imageIntent && !GroupAllowsImageGeneration(apiKeyGroup(getAPIKeyFromContext(c))) { + if imageIntent && !imageGenerationAllowed { return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, ImageGenerationPermissionMessage(), nil) } imageBillingModel := "" @@ -2791,6 +2852,27 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( }, nil } + writeClientMessage := func(message []byte) error { + writeCtx, cancel := context.WithTimeout(ctx, s.openAIWSWriteTimeout()) + defer cancel() + return clientConn.Write(writeCtx, coderws.MessageText, message) + } + + readClientMessage := func() ([]byte, error) { + msgType, payload, readErr := clientConn.Read(ctx) + if readErr != nil { + return nil, readErr + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + return nil, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + fmt.Sprintf("unsupported websocket client message type: %s", msgType.String()), + nil, + ) + } + return payload, nil + } + firstPayload, err := parseClientPayload(firstClientMessage) if err != nil { return err @@ -2799,29 +2881,155 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( turnState := strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader)) stateStore := s.getOpenAIWSStateStore() groupID := getOpenAIGroupIDFromContext(c) - sessionHash := s.GenerateSessionHash(c, firstPayload.rawForHash) - if turnState == "" && stateStore != nil && sessionHash != "" { - if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { - turnState = savedTurnState + storeDisabledConnMode := s.openAIWSStoreDisabledConnMode() + sessionHash := "" + preferredConnID := "" + storeDisabled := false + refreshIngressRouteState := func(payload openAIWSClientPayload) { + sessionHash = s.GenerateSessionHash(c, payload.rawForHash) + if turnState == "" && stateStore != nil && sessionHash != "" { + if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { + turnState = savedTurnState + } } - } - preferredConnID := "" - if stateStore != nil && firstPayload.previousResponseID != "" { - if connID, ok := stateStore.GetResponseConn(firstPayload.previousResponseID); ok { - preferredConnID = connID + preferredConnID = "" + if stateStore != nil && payload.previousResponseID != "" { + if connID, ok := stateStore.GetResponseConn(payload.previousResponseID); ok { + preferredConnID = connID + } + } + + storeDisabled = s.isOpenAIWSStoreDisabledInRequestRaw(payload.payloadRaw, account) + if stateStore != nil && storeDisabled && payload.previousResponseID == "" && sessionHash != "" { + if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok { + preferredConnID = connID + } } } + refreshIngressRouteState(firstPayload) - storeDisabled := s.isOpenAIWSStoreDisabledInRequestRaw(firstPayload.payloadRaw, account) - storeDisabledConnMode := s.openAIWSStoreDisabledConnMode() - if stateStore != nil && storeDisabled && firstPayload.previousResponseID == "" && sessionHash != "" { - if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok { - preferredConnID = connID + if s.shouldBridgeOpenAIWSHTTP(firstPayload.payloadBytes, firstPayload.previousResponseID) { + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_start account_id=%d account_type=%s payload_bytes=%d threshold_bytes=%d has_session_hash=%v store_disabled=%v", + account.ID, + account.Type, + firstPayload.payloadBytes, + s.openAIWSHTTPBridgeThresholdBytes(), + sessionHash != "", + storeDisabled, + ) + currentBridgePayload := firstPayload + var bridgeReplayInput []json.RawMessage + bridgeReplayInputExists := false + for turn := 1; ; turn++ { + if turn > 1 && hooks != nil && hooks.BeforeRequest != nil { + if err := hooks.BeforeRequest(turn, currentBridgePayload.payloadRaw, currentBridgePayload.originalModel); err != nil { + return err + } + } + if hooks != nil && hooks.BeforeTurn != nil { + if err := hooks.BeforeTurn(turn); err != nil { + return err + } + } + if turnState != "" && c != nil && c.Request != nil { + c.Request.Header.Set(openAIWSTurnStateHeader, turnState) + } + bridgePayloadRaw := currentBridgePayload.payloadRaw + bridgePayloadBytes := currentBridgePayload.payloadBytes + needsBridgeReplay := currentBridgePayload.previousResponseID != "" || openAIWSRawPayloadHasToolCallOutput(currentBridgePayload.payloadRaw) + turnReplayInput, turnReplayInputExists, replayInputErr := buildOpenAIWSReplayInputSequence( + bridgeReplayInput, + bridgeReplayInputExists, + currentBridgePayload.payloadRaw, + needsBridgeReplay, + ) + if replayInputErr != nil { + return fmt.Errorf("build websocket http bridge replay input: %w", replayInputErr) + } + if needsBridgeReplay && turnReplayInputExists { + updatedPayload, setInputErr := setOpenAIWSPayloadInputSequence( + currentBridgePayload.payloadRaw, + turnReplayInput, + true, + ) + if setInputErr != nil { + return fmt.Errorf("set websocket http bridge replay input: %w", setInputErr) + } + bridgePayloadRaw = updatedPayload + bridgePayloadBytes = len(updatedPayload) + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_replay_input account_id=%d turn=%d input_items=%d previous_response_id_present=%v has_tool_output=%v", + account.ID, + turn, + len(turnReplayInput), + currentBridgePayload.previousResponseID != "", + openAIWSRawPayloadHasToolCallOutput(currentBridgePayload.payloadRaw), + ) + } + result, bridgeErr := s.proxyOpenAIWSHTTPBridgeTurn( + ctx, + c, + account, + token, + bridgePayloadRaw, + bridgePayloadBytes, + currentBridgePayload.originalModel, + currentBridgePayload.imageBillingModel, + currentBridgePayload.imageSizeTier, + currentBridgePayload.imageInputSize, + turn, + writeClientMessage, + ) + if hooks != nil && hooks.AfterTurn != nil { + hooks.AfterTurn(turn, result, bridgeErr) + } + if bridgeErr != nil { + return bridgeErr + } + if result == nil { + return errors.New("websocket http bridge turn result is nil") + } + bridgeReplayInput = cloneOpenAIWSRawMessages(turnReplayInput) + bridgeReplayInputExists = turnReplayInputExists + if result.wsReplayInputExists { + bridgeReplayInput = append(bridgeReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...) + bridgeReplayInputExists = true + } + if bridgeTurnState := strings.TrimSpace(result.ResponseHeaders.Get(openAIWSTurnStateHeader)); bridgeTurnState != "" { + turnState = bridgeTurnState + if stateStore != nil && sessionHash != "" { + stateStore.BindSessionTurnState(groupID, sessionHash, bridgeTurnState, s.openAIWSSessionStickyTTL()) + } + } + responseID := strings.TrimSpace(result.RequestID) + if responseID != "" && stateStore != nil { + ttl := s.openAIWSResponseStickyTTL() + logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) + } + nextClientMessage, readErr := readClientMessage() + if readErr != nil { + if isOpenAIWSClientDisconnectError(readErr) { + closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_client_closed account_id=%d close_status=%s close_reason=%s", + account.ID, + closeStatus, + truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), + ) + return nil + } + return fmt.Errorf("read client websocket request: %w", readErr) + } + nextPayload, parseErr := parseClientPayload(nextClientMessage) + if parseErr != nil { + return parseErr + } + currentBridgePayload = nextPayload } } - isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator")) || (s.cfg != nil && s.cfg.Gateway.ForceCodexCLI) wsHeaders, _ := s.buildOpenAIWSHeaders(c, account, token, wsDecision, isCodexCLI, turnState, strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)), firstPayload.promptCacheKey) baseAcquireReq := openAIWSAcquireRequest{ Account: account, @@ -2971,27 +3179,6 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( return lease, nil } - writeClientMessage := func(message []byte) error { - writeCtx, cancel := context.WithTimeout(ctx, s.openAIWSWriteTimeout()) - defer cancel() - return clientConn.Write(writeCtx, coderws.MessageText, message) - } - - readClientMessage := func() ([]byte, error) { - msgType, payload, readErr := clientConn.Read(ctx) - if readErr != nil { - return nil, readErr - } - if msgType != coderws.MessageText && msgType != coderws.MessageBinary { - return nil, NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - fmt.Sprintf("unsupported websocket client message type: %s", msgType.String()), - nil, - ) - } - return payload, nil - } - sendAndRelay := func(turn int, lease *openAIWSConnLease, payload []byte, payloadBytes int, originalModel string, imageBillingModel string, imageSizeTier string, imageInputSize string) (*OpenAIForwardResult, error) { if lease == nil { return nil, errors.New("upstream websocket lease is nil") @@ -3028,6 +3215,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( eventCount := 0 tokenEventCount := 0 terminalEventCount := 0 + replayCollector := &openAIWSToolCallReplayCollector{} firstEventType := "" lastEventType := "" needModelReplace := false @@ -3158,6 +3346,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( upstreamMessage = corrected } } + replayCollector.AddEvent(eventType, upstreamMessage) if err := writeClientMessage(upstreamMessage); err != nil { if isOpenAIWSClientDisconnectError(err) { clientDisconnected = true @@ -3221,6 +3410,10 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( Duration: time.Since(turnStart), FirstTokenMs: firstTokenMs, } + if replayInput := replayCollector.Items(); len(replayInput) > 0 { + result.wsReplayInput = replayInput + result.wsReplayInputExists = true + } if imageCount > 0 { result.ImageCount = imageCount result.ImageSize = imageSizeTier @@ -3614,9 +3807,12 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( if forcePreferredConn { // 携带 function_call_output 的请求不能丢弃 previous_response_id： // 上游 API 需要 response chain 来匹配 tool_result 与之前的 tool_use， - // 丢弃后会导致 "No tool call found for function call output" 400 错误。 + // 除非 replay input 已经包含与每个 tool_result 匹配的 tool_use 上下文。 hasFCOutput := hasFunctionCallOutput - if !turnPrevRecoveryTried && currentPreviousResponseID != "" && !hasFCOutput { + hasReplayToolContext := hasFCOutput && + currentTurnReplayInputExists && + openAIWSRawItemsHaveToolCallContextForOutputs(currentTurnReplayInput) + if !turnPrevRecoveryTried && currentPreviousResponseID != "" && (!hasFCOutput || hasReplayToolContext) { updatedPayload, removed, dropErr := dropPreviousResponseIDFromRawPayload(currentPayload) if dropErr != nil || !removed { reason := "not_removed" @@ -3648,11 +3844,13 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( ) } else { logOpenAIWSModeInfo( - "ingress_ws_preflight_ping_recovery account_id=%d turn=%d conn_id=%s action=drop_previous_response_id_retry previous_response_id=%s", + "ingress_ws_preflight_ping_recovery account_id=%d turn=%d conn_id=%s action=drop_previous_response_id_retry previous_response_id=%s has_function_call_output=%v has_replay_tool_context=%v", account.ID, turn, truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + hasFCOutput, + hasReplayToolContext, ) turnPrevRecoveryTried = true currentPayload = updatedWithInput @@ -3664,12 +3862,18 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( } } if hasFCOutput && currentPreviousResponseID != "" { + reason := "function_call_output_missing_replay_context" + if hasReplayToolContext { + reason = "function_call_output_replay_not_applied" + } logOpenAIWSModeInfo( - "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=function_call_output action=fail_close previous_response_id=%s", + "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=%s action=fail_close previous_response_id=%s has_replay_tool_context=%v", account.ID, turn, truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + reason, truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + hasReplayToolContext, ) } resetSessionLease(true) @@ -3749,6 +3953,10 @@ func (s *OpenAIGatewayService) proxyOpenAIWebSocketFromClient( lastTurnPayload = cloneOpenAIWSPayloadBytes(currentPayload) lastTurnReplayInput = cloneOpenAIWSRawMessages(currentTurnReplayInput) lastTurnReplayInputExists = currentTurnReplayInputExists + if result.wsReplayInputExists { + lastTurnReplayInput = append(lastTurnReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...) + lastTurnReplayInputExists = true + } nextStrictState, strictStateErr := buildOpenAIWSIngressPreviousTurnStrictState(currentPayload) if strictStateErr != nil { lastTurnStrictState = nil @@ -4195,6 +4403,13 @@ func (s *OpenAIGatewayService) selectAccountByPreviousResponseIDForCapability( if !account.SupportsOpenAIEndpointCapability(requiredCapability) { return nil, nil } + // Quota auto-pause must also gate the previous_response_id sticky path; otherwise an + // account over its 5h/7d threshold keeps serving the same response chain even though + // normal scheduling skips it. Pause is transient, so fall through to normal scheduling + // without deleting the binding (the window may reset before the next turn). + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { + return nil, nil + } if s.schedulerSnapshot != nil && s.accountRepo != nil { latest, latestErr := s.accountRepo.GetByID(ctx, account.ID) if latestErr != nil || latest == nil { @@ -4211,6 +4426,9 @@ func (s *OpenAIGatewayService) selectAccountByPreviousResponseIDForCapability( if !latest.SupportsOpenAIEndpointCapability(requiredCapability) { return nil, nil } + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, latest); paused { + return nil, nil + } if s.isOpenAIAccountRuntimeBlocked(latest) { _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) return nil, nil 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 b7f1bc4f78da..dc48b990ea25 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -298,6 +298,142 @@ 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) { + 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.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + 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_codex_image_bridge","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + captureDialer := &openAIWSCaptureDialer{conn: captureConn} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(captureDialer) + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + + groupID := int64(3) + apiKey := &APIKey{ + ID: 1, + UserID: 1, + GroupID: &groupID, + Group: &Group{ + ID: groupID, + AllowImageGeneration: true, + }, + } + account := &Account{ + ID: 31, + Name: "openai-codex-image-ws", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "test-token", + }, + Extra: map[string]any{ + "openai_oauth_responses_websockets_v2_enabled": true, + "codex_image_generation_bridge": 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() + }() + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "codex_cli_rs/0.98.0") + ginCtx.Request = req + ginCtx.Set("api_key", apiKey) + + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + serverErrCh <- errors.New("unsupported websocket client message type") + return + } + + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "test-token", 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.5","stream":false,"input":"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_bridge", gjson.GetBytes(message, "response.id").String()) + + _ = clientConn.Close(coderws.StatusNormalClosure, "done") + + select { + case serverErr := <-serverErrCh: + require.NoError(t, serverErr) + case <-time.After(5 * time.Second): + 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.Contains(t, gjson.Get(upstreamPayload, "instructions").String(), "image_generation") +} + func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_DedicatedModeDoesNotReuseConnAcrossSessions(t *testing.T) { gin.SetMode(gin.TestMode) @@ -2305,6 +2441,161 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledStr require.Equal(t, "world", gjson.Get(secondWrite, "input.1.text").String()) } +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPreflightPingFailReplaysFunctionCallOutputWithContext(t *testing.T) { + gin.SetMode(gin.TestMode) + prevPreflightPingIdle := openAIWSIngressPreflightPingIdle + openAIWSIngressPreflightPingIdle = 0 + defer func() { + openAIWSIngressPreflightPingIdle = prevPreflightPingIdle + }() + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + firstConn := &openAIWSPreflightFailConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_replay_ctx_1","model":"gpt-5.1","output":[{"type":"function_call","id":"fc_replay_1","call_id":"call_replay_1","name":"shell","arguments":"{}"}],"usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + secondConn := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_replay_ctx_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + dialer := &openAIWSQueueDialer{ + conns: []openAIWSClientConn{firstConn, secondConn}, + } + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(dialer) + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + + account := &Account{ + ID: 128, + Name: "openai-ingress-preflight-replay-function-output-with-context", + 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() + }() + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "unit-test-agent/1.0") + ginCtx.Request = req + + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + serverErrCh <- errors.New("unsupported websocket client message type") + return + } + + 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() + }() + + writeMessage := func(payload string) { + writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload))) + } + readMessage := func() []byte { + readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + msgType, message, readErr := clientConn.Read(readCtx) + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + return message + } + + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"call tool"}]}]}`) + firstTurn := readMessage() + require.Equal(t, "resp_turn_ping_replay_ctx_1", gjson.GetBytes(firstTurn, "response.id").String()) + + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"previous_response_id":"resp_turn_ping_replay_ctx_1","input":[{"type":"function_call_output","call_id":"call_replay_1","output":"ok"}]}`) + secondTurn := readMessage() + require.Equal(t, "resp_turn_ping_replay_ctx_2", gjson.GetBytes(secondTurn, "response.id").String()) + + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + select { + case serverErr := <-serverErrCh: + require.NoError(t, serverErr) + case <-time.After(5 * time.Second): + t.Fatal("等待 ingress websocket function_call_output 自包含重放后结束超时") + } + + require.Equal(t, 2, dialer.DialCount(), "带完整 tool 上下文的 function_call_output 应在 ping 失败后换新连接重放") + require.Equal(t, 1, firstConn.WriteCount()) + require.GreaterOrEqual(t, firstConn.PingCount(), 1) + secondConn.mu.Lock() + secondWrites := append([]map[string]any(nil), secondConn.writes...) + secondConn.mu.Unlock() + require.Len(t, secondWrites, 1) + secondWrite := requestToJSONString(secondWrites[0]) + require.False(t, gjson.Get(secondWrite, "previous_response_id").Exists()) + require.Equal(t, 3, len(gjson.Get(secondWrite, "input").Array())) + require.Equal(t, "message", gjson.Get(secondWrite, "input.0.type").String()) + require.Equal(t, "function_call", gjson.Get(secondWrite, "input.1.type").String()) + require.Equal(t, "call_replay_1", gjson.Get(secondWrite, "input.1.call_id").String()) + require.Equal(t, "function_call_output", gjson.Get(secondWrite, "input.2.type").String()) + require.Equal(t, "call_replay_1", gjson.Get(secondWrite, "input.2.call_id").String()) +} + func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPreflightPingFailClosesWhenFunctionCallOutputNeedsPreviousResponseID(t *testing.T) { gin.SetMode(gin.TestMode) prevPreflightPingIdle := openAIWSIngressPreflightPingIdle diff --git a/backend/internal/service/openai_ws_forwarder_success_test.go b/backend/internal/service/openai_ws_forwarder_success_test.go index ed5759e710c2..c11b13a472d9 100644 --- a/backend/internal/service/openai_ws_forwarder_success_test.go +++ b/backend/internal/service/openai_ws_forwarder_success_test.go @@ -450,6 +450,90 @@ func TestOpenAIGatewayService_ProxyRealtimeWebSocketFromClient_RewritesSessionUp require.Equal(t, "next", gjson.GetBytes(hookPayloads[1].Payload, "session.instructions").String()) } +func TestOpenAIGatewayService_Forward_WSv2_UsesPatchedBodyAfterValidationDecode(t *testing.T) { + gin.SetMode(gin.TestMode) + + type receivedPayload struct { + MaxCompletionTokensExists bool + } + receivedCh := make(chan receivedPayload, 1) + + upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket failed: %v", err) + return + } + defer func() { _ = conn.Close() }() + + var request map[string]any + if err := conn.ReadJSON(&request); err != nil { + t.Errorf("read ws request failed: %v", err) + return + } + requestJSON := requestToJSONString(request) + receivedCh <- receivedPayload{MaxCompletionTokensExists: gjson.Get(requestJSON, "max_completion_tokens").Exists()} + + if err := conn.WriteJSON(map[string]any{ + "type": "response.completed", + "response": map[string]any{ + "id": "resp_patched_ws_1", + "model": "gpt-5.3-codex-spark", + "usage": map[string]any{"input_tokens": 1, "output_tokens": 1}, + }, + }); err != nil { + t.Errorf("write response.completed failed: %v", err) + return + } + })) + defer wsServer.Close() + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Request.Header.Set("User-Agent", "unit-test-agent/1.0") + + 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.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 30 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 10 + + svc := &OpenAIGatewayService{ + cfg: cfg, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + } + + account := &Account{ + ID: 10, + Name: "openai-ws", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": wsServer.URL, + }, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + + body := []byte(`{"model":"gpt-5.4","stream":false,"max_completion_tokens":12,"tools":[{"type":"image_generation"}],"input":[{"type":"input_text","text":"hello"}]}`) + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, result.OpenAIWSMode) + + received := <-receivedCh + require.False(t, received.MaxCompletionTokensExists) +} func TestOpenAIGatewayService_Forward_WSv2_ImageGenerationCountsOutputs(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go new file mode 100644 index 000000000000..1f0f32a05754 --- /dev/null +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -0,0 +1,387 @@ +package service + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +const ( + openAIWSClientReadLimitBytesDefault int64 = 64 * 1024 * 1024 + openAIWSHTTPBridgeThresholdBytesDefault int64 = 15 * 1024 * 1024 + openAIWSHTTPBridgeErrorBodyLimitBytes = 64 * 1024 +) + +func ResolveOpenAIWSClientReadLimitBytes(cfg *config.Config) int64 { + if cfg == nil || cfg.Gateway.OpenAIWS.ClientReadLimitBytes <= 0 { + return openAIWSClientReadLimitBytesDefault + } + return cfg.Gateway.OpenAIWS.ClientReadLimitBytes +} + +func (s *OpenAIGatewayService) openAIWSHTTPBridgeEnabled() bool { + return s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.HTTPBridgeEnabled +} + +func (s *OpenAIGatewayService) openAIWSHTTPBridgeThresholdBytes() int64 { + if s == nil || s.cfg == nil || s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes <= 0 { + return openAIWSHTTPBridgeThresholdBytesDefault + } + return s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes +} + +func (s *OpenAIGatewayService) shouldBridgeOpenAIWSHTTP(payloadBytes int, previousResponseID string) bool { + if !s.openAIWSHTTPBridgeEnabled() { + return false + } + if strings.TrimSpace(previousResponseID) != "" { + return false + } + threshold := s.openAIWSHTTPBridgeThresholdBytes() + return threshold > 0 && int64(payloadBytes) >= threshold +} + +func prepareOpenAIWSHTTPBridgeBody(payload []byte) ([]byte, error) { + var body map[string]any + if err := json.Unmarshal(payload, &body); err != nil { + return nil, err + } + if body == nil { + return nil, errors.New("response.create payload must be a JSON object") + } + delete(body, "type") + delete(body, "generate") + delete(body, "previous_response_id") + body["stream"] = true + return json.Marshal(body) +} + +type openAIWSToolCallReplayCollector struct { + items []json.RawMessage + seen map[string]struct{} +} + +func (c *openAIWSToolCallReplayCollector) AddEvent(eventType string, message []byte) { + switch strings.TrimSpace(eventType) { + case "response.output_item.done": + c.addItem(gjson.GetBytes(message, "item")) + case "response.completed", "response.done": + output := gjson.GetBytes(message, "response.output") + if !output.IsArray() { + return + } + for _, item := range output.Array() { + c.addItem(item) + } + } +} + +func (c *openAIWSToolCallReplayCollector) Items() []json.RawMessage { + return cloneOpenAIWSRawMessages(c.items) +} + +func (c *openAIWSToolCallReplayCollector) addItem(item gjson.Result) { + if !item.Exists() || item.Type != gjson.JSON { + return + } + raw := strings.TrimSpace(item.Raw) + if raw == "" || !strings.HasPrefix(raw, "{") { + return + } + if !isCodexToolCallContextItemType(item.Get("type").String()) { + return + } + key := strings.TrimSpace(item.Get("id").String()) + if key == "" { + key = strings.TrimSpace(item.Get("call_id").String()) + } + if key == "" { + key = raw + } + if c.seen == nil { + c.seen = make(map[string]struct{}) + } + if _, ok := c.seen[key]; ok { + return + } + c.seen[key] = struct{}{} + c.items = append(c.items, json.RawMessage(raw)) +} + +func buildOpenAIWSHTTPBridgeErrorEvent(statusCode int, message string) []byte { + message = strings.TrimSpace(message) + if message == "" { + message = http.StatusText(statusCode) + } + if message == "" { + message = "upstream request failed" + } + event := map[string]any{ + "type": "error", + "status": statusCode, + "error": map[string]any{ + "type": "upstream_error", + "message": message, + }, + } + body, err := json.Marshal(event) + if err != nil { + return []byte(`{"type":"error","error":{"type":"upstream_error","message":"upstream request failed"}}`) + } + return body +} + +func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( + ctx context.Context, + c *gin.Context, + account *Account, + token string, + payload []byte, + payloadBytes int, + originalModel string, + imageBillingModel string, + imageSizeTier string, + imageInputSize string, + turn int, + writeClientMessage func([]byte) error, +) (*OpenAIForwardResult, error) { + if s == nil { + return nil, errors.New("service is nil") + } + if s.httpUpstream == nil { + return nil, errors.New("openai http upstream is nil") + } + if account == nil { + return nil, errors.New("account is nil") + } + if writeClientMessage == nil { + return nil, errors.New("client websocket writer is nil") + } + + body, err := prepareOpenAIWSHTTPBridgeBody(payload) + if err != nil { + return nil, fmt.Errorf("prepare http bridge body: %w", err) + } + + upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) + upstreamReq, err := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token) + releaseUpstreamCtx() + if err != nil { + return nil, err + } + + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + if c != nil { + c.Set("openai_passthrough", true) + c.Set("openai_ws_http_bridge", true) + } + + turnStart := time.Now() + resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) + if err != nil { + safeErr := sanitizeUpstreamErrorMessage(err.Error()) + _ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(http.StatusBadGateway, "Upstream request failed")) + return nil, fmt.Errorf("upstream http bridge request failed: %s", safeErr) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode >= 400 { + respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes)) + upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) + if upstreamMsg == "" { + upstreamMsg = http.StatusText(resp.StatusCode) + } + _ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(resp.StatusCode, upstreamMsg)) + return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg) + } + + responseID := "" + usage := OpenAIUsage{} + imageCounter := newOpenAIImageOutputCounter() + var firstTokenMs *int + reqStream := openAIWSPayloadBoolFromRaw(body, "stream", true) + eventCount := 0 + tokenEventCount := 0 + terminalEventCount := 0 + replayCollector := &openAIWSToolCallReplayCollector{} + firstEventType := "" + lastEventType := "" + sawDone := false + wroteDownstream := false + clientDisconnected := false + mappedModel := "" + needModelReplace := false + var mappedModelBytes []byte + if originalModel != "" { + mappedModel = normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)) + needModelReplace = mappedModel != "" && mappedModel != originalModel + if needModelReplace { + mappedModelBytes = []byte(mappedModel) + } + } + + resultWithUsage := func() *OpenAIForwardResult { + imageCount := imageCounter.Count() + result := &OpenAIForwardResult{ + RequestID: responseID, + Usage: usage, + Model: originalModel, + UpstreamModel: mappedModel, + ServiceTier: extractOpenAIServiceTierFromBody(body), + ReasoningEffort: extractOpenAIReasoningEffortFromBody(body, originalModel), + Stream: reqStream, + OpenAIWSMode: true, + ResponseHeaders: cloneHeader(resp.Header), + Duration: time.Since(turnStart), + FirstTokenMs: firstTokenMs, + } + if replayInput := replayCollector.Items(); len(replayInput) > 0 { + result.wsReplayInput = replayInput + result.wsReplayInputExists = true + } + if imageCount > 0 { + result.ImageCount = imageCount + result.ImageSize = imageSizeTier + result.ImageInputSize = imageInputSize + result.ImageOutputSizes = imageCounter.Sizes() + result.BillingModel = imageBillingModel + } + return result + } + + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + defer putSSEScannerBuf64K(scanBuf) + + for scanner.Scan() { + line := scanner.Text() + data, ok := extractOpenAISSEDataLine(line) + if !ok { + continue + } + trimmedData := strings.TrimSpace(data) + if trimmedData == "" { + continue + } + if trimmedData == "[DONE]" { + sawDone = true + continue + } + + upstreamMessage := []byte(trimmedData) + eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage) + if responseID == "" && eventResponseID != "" { + responseID = eventResponseID + } + if eventType != "" { + eventCount++ + if firstEventType == "" { + firstEventType = eventType + } + lastEventType = eventType + } + if isOpenAIWSTokenEvent(eventType) { + tokenEventCount++ + if firstTokenMs == nil { + ms := int(time.Since(turnStart).Milliseconds()) + firstTokenMs = &ms + } + } + if openAIWSEventShouldParseUsage(eventType) { + parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage) + } + imageCounter.AddSSEData(upstreamMessage) + + if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && strings.Contains(trimmedData, mappedModel) { + upstreamMessage = replaceOpenAIWSMessageModel(upstreamMessage, mappedModel, originalModel) + } + if s.toolCorrector != nil && openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(upstreamMessage) { + if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(upstreamMessage); changed { + upstreamMessage = corrected + } + } + replayCollector.AddEvent(eventType, upstreamMessage) + + if !clientDisconnected { + if err := writeClientMessage(upstreamMessage); err != nil { + if isOpenAIWSClientDisconnectError(err) { + clientDisconnected = true + closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err) + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_client_disconnected_drain account_id=%d turn=%d close_status=%s close_reason=%s", + account.ID, + turn, + closeStatus, + truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), + ) + } else { + return nil, wrapOpenAIWSIngressTurnError( + "write_client", + fmt.Errorf("write client websocket event: %w", err), + wroteDownstream, + ) + } + } else { + wroteDownstream = true + } + } + + if eventType == "error" { + errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) + s.persistOpenAIWSRateLimitSignal(ctx, account, resp.Header, upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw) + errMessage := strings.TrimSpace(errMsgRaw) + if errMessage == "" { + errMessage = "upstream error event" + } + return resultWithUsage(), errors.New(errMessage) + } + if isOpenAIWSTerminalEvent(eventType) { + terminalEventCount++ + firstTokenMsValue := -1 + if firstTokenMs != nil { + firstTokenMsValue = *firstTokenMs + } + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_turn_completed account_id=%d turn=%d response_id=%s payload_bytes=%d duration_ms=%d events=%d token_events=%d terminal_events=%d first_event=%s last_event=%s first_token_ms=%d client_disconnected=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), + payloadBytes, + time.Since(turnStart).Milliseconds(), + eventCount, + tokenEventCount, + terminalEventCount, + truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen), + firstTokenMsValue, + clientDisconnected, + ) + return resultWithUsage(), nil + } + } + if err := scanner.Err(); err != nil { + return resultWithUsage(), fmt.Errorf("read upstream http bridge stream: %w", err) + } + if sawDone && eventCount > 0 { + return resultWithUsage(), nil + } + return resultWithUsage(), errors.New("upstream http bridge stream ended before terminal event") +} diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go new file mode 100644 index 000000000000..0a1d6b564bb3 --- /dev/null +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -0,0 +1,461 @@ +package service + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestPrepareOpenAIWSHTTPBridgeBodyStripsWSFields(t *testing.T) { + body, err := prepareOpenAIWSHTTPBridgeBody([]byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":false,"previous_response_id":"resp_prev","input":"hi"}`)) + require.NoError(t, err) + require.False(t, gjson.GetBytes(body, "type").Exists()) + require.False(t, gjson.GetBytes(body, "generate").Exists()) + require.False(t, gjson.GetBytes(body, "previous_response_id").Exists()) + require.Equal(t, "gpt-5", gjson.GetBytes(body, "model").String()) + require.True(t, gjson.GetBytes(body, "stream").Bool()) + require.Equal(t, "hi", gjson.GetBytes(body, "input").String()) +} + +func TestOpenAIWSHTTPBridgeDecisionKeepsSmallFramesOnWS(t *testing.T) { + svc := &OpenAIGatewayService{ + cfg: &config.Config{ + Gateway: config.GatewayConfig{ + OpenAIWS: config.GatewayOpenAIWSConfig{ + HTTPBridgeEnabled: true, + HTTPBridgeThresholdBytes: 100, + }, + }, + }, + } + + require.False(t, svc.shouldBridgeOpenAIWSHTTP(99, "")) + require.True(t, svc.shouldBridgeOpenAIWSHTTP(100, "")) + require.False(t, svc.shouldBridgeOpenAIWSHTTP(1000, "resp_existing")) + + svc.cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = false + require.False(t, svc.shouldBridgeOpenAIWSHTTP(1000, "")) +} + +func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_bridge","model":"gpt-5"}}`, + "", + `data: {"type":"response.output_text.delta","response":{"id":"resp_bridge"},"delta":"ok"}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_bridge","model":"gpt-5","usage":{"input_tokens":3,"output_tokens":2}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "x-request-id": []string{"rid_bridge"}, + }, + Body: io.NopCloser(strings.NewReader(sseBody)), + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{ + Gateway: config.GatewayConfig{ + MaxLineSize: defaultMaxLineSize, + OpenAIWS: config.GatewayOpenAIWSConfig{ + HTTPBridgeEnabled: true, + HTTPBridgeThresholdBytes: 1, + }, + }, + }, + httpUpstream: upstream, + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 7, + Name: "api-key", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Status: StatusActive, + } + payload := []byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"input":"hi"}`) + + type bridgeResult struct { + result *OpenAIForwardResult + err error + } + resultCh := make(chan bridgeResult, 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 { + resultCh <- bridgeResult{err: err} + return + } + defer func() { _ = conn.CloseNow() }() + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + ginCtx.Request = req + + writeClient := func(message []byte) error { + writeCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + defer cancel() + return conn.Write(writeCtx, coderws.MessageText, message) + } + result, bridgeErr := svc.proxyOpenAIWSHTTPBridgeTurn( + r.Context(), + ginCtx, + account, + "sk-test", + payload, + len(payload), + "gpt-5", + "", + "", + "", + 1, + writeClient, + ) + resultCh <- bridgeResult{result: result, err: bridgeErr} + })) + 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() }() + + readEvent := func() []byte { + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + msgType, event, readErr := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + return event + } + + created := readEvent() + delta := readEvent() + completed := readEvent() + + require.Equal(t, "response.created", gjson.GetBytes(created, "type").String()) + require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String()) + require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String()) + + select { + case bridge := <-resultCh: + require.NoError(t, bridge.err) + require.NotNil(t, bridge.result) + require.Equal(t, "resp_bridge", bridge.result.RequestID) + require.Equal(t, 3, bridge.result.Usage.InputTokens) + require.Equal(t, 2, bridge.result.Usage.OutputTokens) + require.True(t, bridge.result.OpenAIWSMode) + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for bridge result") + } + + require.NotNil(t, upstream.lastReq) + require.Equal(t, http.MethodPost, upstream.lastReq.Method) + 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 TestOpenAIWSHTTPBridgeAcceptsFirstFrameAboveLegacy16MiB(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_large_bridge","model":"gpt-5"}}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_large_bridge","model":"gpt-5","usage":{"input_tokens":9,"output_tokens":1}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "x-request-id": []string{"rid_large_bridge"}, + }, + Body: io.NopCloser(strings.NewReader(sseBody)), + }} + cfg := &config.Config{ + Gateway: config.GatewayConfig{ + MaxLineSize: defaultMaxLineSize, + OpenAIWS: config.GatewayOpenAIWSConfig{ + Enabled: true, + APIKeyEnabled: true, + ResponsesWebsocketsV2: true, + ClientReadLimitBytes: 64 * 1024 * 1024, + HTTPBridgeEnabled: true, + HTTPBridgeThresholdBytes: 15 * 1024 * 1024, + }, + }, + } + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: upstream, + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 9, + Name: "api-key", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "sk-upstream"}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + }, + Concurrency: 1, + Status: StatusActive, + } + + payload := []byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"input":"` + strings.Repeat("x", 17*1024*1024) + `"}`) + require.Greater(t, len(payload), 16*1024*1024) + require.Less(t, int64(len(payload)), ResolveOpenAIWSClientReadLimitBytes(cfg)) + + 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() }() + conn.SetReadLimit(ResolveOpenAIWSClientReadLimitBytes(cfg)) + + readCtx, cancelRead := context.WithTimeout(r.Context(), 10*time.Second) + msgType, firstMessage, err := conn.Read(readCtx) + cancelRead() + if err != nil { + errCh <- err + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + errCh <- NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "unexpected client websocket message type", nil) + return + } + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "codex_cli_rs/0.135.0") + ginCtx.Request = req + + proxyCtx, cancelProxy := context.WithTimeout(r.Context(), 20*time.Second) + defer cancelProxy() + errCh <- svc.ProxyResponsesWebSocketFromClient(proxyCtx, ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 5*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(), 20*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, payload) + cancelWrite() + require.NoError(t, err) + + var eventTypes []string + for { + readCtx, cancelRead := context.WithTimeout(context.Background(), 10*time.Second) + msgType, event, readErr := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + + eventType := gjson.GetBytes(event, "type").String() + eventTypes = append(eventTypes, eventType) + if eventType == "response.completed" { + break + } + } + require.Contains(t, eventTypes, "response.created") + require.Contains(t, eventTypes, "response.completed") + + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + select { + case proxyErr := <-errCh: + require.NoError(t, proxyErr) + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for websocket bridge proxy to finish") + } + + require.NotNil(t, upstream.lastReq) + require.Equal(t, http.MethodPost, upstream.lastReq.Method) + require.Greater(t, len(upstream.lastBody), 16*1024*1024) + 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()) + require.Equal(t, "gpt-5", gjson.GetBytes(upstream.lastBody, "model").String()) +} + +func TestOpenAIWSHTTPBridgeKeepsContinuationFramesOnHTTPWithoutPreviousResponseID(t *testing.T) { + gin.SetMode(gin.TestMode) + + firstSSEBody := strings.Join([]string{ + `data: {"type":"response.completed","response":{"id":"resp_bridge_first","model":"gpt-5.1","output":[{"type":"function_call","id":"fc_bridge_1","call_id":"call_bridge_1","name":"shell","arguments":"{}"}],"usage":{"input_tokens":9,"output_tokens":1}}}`, + "", + }, "\n") + secondSSEBody := strings.Join([]string{ + `data: {"type":"response.completed","response":{"id":"resp_bridge_second","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + { + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + }, + Body: io.NopCloser(strings.NewReader(firstSSEBody)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + }, + Body: io.NopCloser(strings.NewReader(secondSSEBody)), + }, + }} + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true + cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = 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{} + captureDialer := &openAIWSCaptureDialer{conn: captureConn} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(captureDialer) + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: upstream, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + account := &Account{ + ID: 19, + Name: "api-key-bridge-handoff", + 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) + msgType, firstMessage, err := conn.Read(readCtx) + cancelRead() + if err != nil { + errCh <- err + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + errCh <- NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "unexpected client websocket message type", nil) + return + } + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "codex_cli_rs/0.135.0") + ginCtx.Request = req + + 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() }() + + writeMessage := func(payload string) { + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + defer cancelWrite() + require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload))) + } + readMessage := func() []byte { + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + defer cancelRead() + msgType, event, readErr := clientConn.Read(readCtx) + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + return event + } + + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":true,"input":"first"}`) + firstTurnEvent := readMessage() + require.Equal(t, "response.completed", gjson.GetBytes(firstTurnEvent, "type").String()) + require.Equal(t, "resp_bridge_first", gjson.GetBytes(firstTurnEvent, "response.id").String()) + + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_bridge_first","input":[{"type":"function_call_output","call_id":"call_bridge_1","output":"ok"}]}`) + secondTurnEvent := readMessage() + require.Equal(t, "response.completed", gjson.GetBytes(secondTurnEvent, "type").String()) + require.Equal(t, "resp_bridge_second", gjson.GetBytes(secondTurnEvent, "response.id").String()) + + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + select { + case proxyErr := <-errCh: + require.NoError(t, proxyErr) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for websocket bridge proxy to finish") + } + + require.Len(t, upstream.bodies, 2, "进入 HTTP bridge 后同一客户端 WS 连接内应保持 HTTP/SSE bridge") + require.False(t, gjson.GetBytes(upstream.bodies[0], "previous_response_id").Exists()) + require.False(t, gjson.GetBytes(upstream.bodies[1], "previous_response_id").Exists()) + secondInput := gjson.GetBytes(upstream.bodies[1], "input").Array() + require.Len(t, secondInput, 3) + require.Equal(t, "first", secondInput[0].String()) + require.Equal(t, "function_call", secondInput[1].Get("type").String()) + require.Equal(t, "call_bridge_1", secondInput[1].Get("call_id").String()) + require.Equal(t, "function_call_output", secondInput[2].Get("type").String()) + require.Equal(t, "call_bridge_1", secondInput[2].Get("call_id").String()) + require.Equal(t, 0, captureDialer.DialCount()) + require.Empty(t, captureConn.writes) +} diff --git a/backend/internal/service/ops_service.go b/backend/internal/service/ops_service.go index 1cea72fae2ce..2d7c5bd42d7c 100644 --- a/backend/internal/service/ops_service.go +++ b/backend/internal/service/ops_service.go @@ -41,6 +41,11 @@ type OpsService struct { // cleanupReloader 由 wire 在 OpsCleanupService 构造完成后通过 SetCleanupReloader 注入。 // 解耦避免 OpsService -> OpsCleanupService 的硬依赖（cleanup 也读 settings，会循环）。 cleanupReloader CleanupReloader + + // quotaAutoPauseSink 由 wire 注入（通常是 SettingService.SetOpenAIQuotaAutoPauseSettings）。 + // UpdateOpsAdvancedSettings 写入新配置后调用，把最新的 quota auto-pause 全局默认阈值 + // 立即同步到调度热路径读取的内存缓存，避免下次请求才能感知新值。 + quotaAutoPauseSink func(OpsOpenAIAccountQuotaAutoPauseSettings) } // CleanupReloader 由 OpsCleanupService 实现。 @@ -57,6 +62,16 @@ func (s *OpsService) SetCleanupReloader(r CleanupReloader) { s.cleanupReloader = r } +// SetOpenAIQuotaAutoPauseSettingsSink 由 wire 注入，把最新的 quota auto-pause 全局默认 +// 阈值 push 到调度热路径读取的内存缓存。同 SetCleanupReloader 的解耦目的：避免 OpsService +// 持有 *SettingService 引入循环依赖。 +func (s *OpsService) SetOpenAIQuotaAutoPauseSettingsSink(sink func(OpsOpenAIAccountQuotaAutoPauseSettings)) { + if s == nil { + return + } + s.quotaAutoPauseSink = sink +} + func NewOpsService( opsRepo OpsRepository, settingRepo SettingRepository, diff --git a/backend/internal/service/ops_settings.go b/backend/internal/service/ops_settings.go index 68c1d9ddea49..472f4e3207aa 100644 --- a/backend/internal/service/ops_settings.go +++ b/backend/internal/service/ops_settings.go @@ -369,6 +369,7 @@ func defaultOpsAdvancedSettings() *OpsAdvancedSettings { Aggregation: OpsAggregationSettings{ AggregationEnabled: false, }, + OpenAIAccountQuotaAutoPause: OpsOpenAIAccountQuotaAutoPauseSettings{}, IgnoreCountTokensErrors: true, // count_tokens 404 是预期行为，默认忽略 IgnoreContextCanceled: true, // Default to true - client disconnects are not errors IgnoreNoAvailableAccounts: false, // Default to false - this is a real routing issue @@ -384,6 +385,8 @@ func normalizeOpsAdvancedSettings(cfg *OpsAdvancedSettings) { if cfg == nil { return } + cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold5h = clampOpsQuotaAutoPauseThreshold(cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold5h) + cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold7d = clampOpsQuotaAutoPauseThreshold(cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold7d) cfg.DataRetention.CleanupSchedule = strings.TrimSpace(cfg.DataRetention.CleanupSchedule) if cfg.DataRetention.CleanupSchedule == "" { cfg.DataRetention.CleanupSchedule = opsCleanupDefaultSchedule @@ -405,6 +408,16 @@ func normalizeOpsAdvancedSettings(cfg *OpsAdvancedSettings) { } } +func clampOpsQuotaAutoPauseThreshold(value float64) float64 { + if value <= 0 { + return 0 + } + if value > 1 { + return 1 + } + return value +} + func validateOpsAdvancedSettings(cfg *OpsAdvancedSettings) error { if cfg == nil { return errors.New("invalid config") @@ -477,6 +490,12 @@ func (s *OpsService) UpdateOpsAdvancedSettings(ctx context.Context, cfg *OpsAdva if err := s.settingRepo.Set(ctx, SettingKeyOpsAdvancedSettings, string(raw)); err != nil { return nil, err } + // Push the new quota auto-pause settings straight into the in-memory cache that + // the OpenAI scheduling hot path reads, so the next request observes the new value + // without waiting for the background refresher's TTL. + if s.quotaAutoPauseSink != nil { + s.quotaAutoPauseSink(cfg.OpenAIAccountQuotaAutoPause) + } // notify cleanup service to reload schedule/enabled. if s.cleanupReloader != nil { diff --git a/backend/internal/service/ops_settings_advanced_test.go b/backend/internal/service/ops_settings_advanced_test.go index 06cc545bb218..62803f948b9b 100644 --- a/backend/internal/service/ops_settings_advanced_test.go +++ b/backend/internal/service/ops_settings_advanced_test.go @@ -4,6 +4,9 @@ import ( "context" "encoding/json" "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" ) func TestGetOpsAdvancedSettings_DefaultHidesOpenAITokenStats(t *testing.T) { @@ -95,3 +98,64 @@ func TestGetOpsAdvancedSettings_BackfillsNewDisplayFlagsFromDefaults(t *testing. t.Fatalf("DisplayAlertEvents = false, want true default backfill") } } + +func TestGetOpenAIQuotaAutoPauseSettings_ReadsDefaultsFromOpsAdvancedSettings(t *testing.T) { + repo := newRuntimeSettingRepoStub() + repo.values[SettingKeyOpsAdvancedSettings] = `{"openai_account_quota_auto_pause":{"default_threshold_5h":0.95,"default_threshold_7d":0.9}}` + svc := NewSettingService(repo, &config.Config{}) + + // Warm the in-memory cache synchronously so the assertion below is deterministic. + // GetOpenAIQuotaAutoPauseSettings is non-blocking on the hot path (returns the + // cached value, refreshes asynchronously); for tests and startup, Warm is the + // synchronous entry point that guarantees a populated cache. + settings := svc.WarmOpenAIQuotaAutoPauseSettings(context.Background()) + if settings.DefaultThreshold5h != 0.95 { + t.Fatalf("DefaultThreshold5h = %v, want 0.95", settings.DefaultThreshold5h) + } + if settings.DefaultThreshold7d != 0.9 { + t.Fatalf("DefaultThreshold7d = %v, want 0.9", settings.DefaultThreshold7d) + } + + // Subsequent Get must hit the warm cache and return the same value without any DB + // access — that's the hot-path invariant. + cached := svc.GetOpenAIQuotaAutoPauseSettings(context.Background()) + if cached.DefaultThreshold5h != 0.95 || cached.DefaultThreshold7d != 0.9 { + t.Fatalf("cached read = %+v, want {0.95, 0.9}", cached) + } +} + +// Hot-path invariant: a Get with cold cache must return immediately (zero defaults) +// rather than blocking on the DB. The async refresher will populate the cache for +// subsequent calls. +func TestGetOpenAIQuotaAutoPauseSettings_ColdCacheNonBlocking(t *testing.T) { + repo := newRuntimeSettingRepoStub() + repo.values[SettingKeyOpsAdvancedSettings] = `{"openai_account_quota_auto_pause":{"default_threshold_5h":0.7}}` + svc := NewSettingService(repo, &config.Config{}) + + start := time.Now() + settings := svc.GetOpenAIQuotaAutoPauseSettings(context.Background()) + elapsed := time.Since(start) + if elapsed > 50*time.Millisecond { + t.Fatalf("cold-cache Get must be non-blocking, took %v", elapsed) + } + // Cold cache means we get zero defaults (the async refresh hasn't completed yet). + if settings.DefaultThreshold5h != 0 || settings.DefaultThreshold7d != 0 { + t.Fatalf("cold-cache Get = %+v, want zeroes", settings) + } +} + +// Explicit cache write (e.g. from UpdateOpsAdvancedSettings) must be visible on the +// very next read without any DB roundtrip. +func TestSetOpenAIQuotaAutoPauseSettings_VisibleImmediately(t *testing.T) { + svc := NewSettingService(newRuntimeSettingRepoStub(), &config.Config{}) + + svc.SetOpenAIQuotaAutoPauseSettings(OpsOpenAIAccountQuotaAutoPauseSettings{ + DefaultThreshold5h: 0.88, + DefaultThreshold7d: 0.77, + }) + + got := svc.GetOpenAIQuotaAutoPauseSettings(context.Background()) + if got.DefaultThreshold5h != 0.88 || got.DefaultThreshold7d != 0.77 { + t.Fatalf("after Set, Get = %+v, want {0.88, 0.77}", got) + } +} diff --git a/backend/internal/service/ops_settings_models.go b/backend/internal/service/ops_settings_models.go index fa18b05fb7cb..4d459e2b9496 100644 --- a/backend/internal/service/ops_settings_models.go +++ b/backend/internal/service/ops_settings_models.go @@ -92,17 +92,23 @@ type OpsAlertRuntimeSettings struct { // OpsAdvancedSettings stores advanced ops configuration (data retention, aggregation). type OpsAdvancedSettings struct { - DataRetention OpsDataRetentionSettings `json:"data_retention"` - Aggregation OpsAggregationSettings `json:"aggregation"` - IgnoreCountTokensErrors bool `json:"ignore_count_tokens_errors"` - IgnoreContextCanceled bool `json:"ignore_context_canceled"` - IgnoreNoAvailableAccounts bool `json:"ignore_no_available_accounts"` - IgnoreInvalidApiKeyErrors bool `json:"ignore_invalid_api_key_errors"` - IgnoreInsufficientBalanceErrors bool `json:"ignore_insufficient_balance_errors"` - DisplayOpenAITokenStats bool `json:"display_openai_token_stats"` - DisplayAlertEvents bool `json:"display_alert_events"` - AutoRefreshEnabled bool `json:"auto_refresh_enabled"` - AutoRefreshIntervalSec int `json:"auto_refresh_interval_seconds"` + DataRetention OpsDataRetentionSettings `json:"data_retention"` + Aggregation OpsAggregationSettings `json:"aggregation"` + OpenAIAccountQuotaAutoPause OpsOpenAIAccountQuotaAutoPauseSettings `json:"openai_account_quota_auto_pause"` + IgnoreCountTokensErrors bool `json:"ignore_count_tokens_errors"` + IgnoreContextCanceled bool `json:"ignore_context_canceled"` + IgnoreNoAvailableAccounts bool `json:"ignore_no_available_accounts"` + IgnoreInvalidApiKeyErrors bool `json:"ignore_invalid_api_key_errors"` + IgnoreInsufficientBalanceErrors bool `json:"ignore_insufficient_balance_errors"` + DisplayOpenAITokenStats bool `json:"display_openai_token_stats"` + DisplayAlertEvents bool `json:"display_alert_events"` + AutoRefreshEnabled bool `json:"auto_refresh_enabled"` + AutoRefreshIntervalSec int `json:"auto_refresh_interval_seconds"` +} + +type OpsOpenAIAccountQuotaAutoPauseSettings struct { + DefaultThreshold5h float64 `json:"default_threshold_5h"` + DefaultThreshold7d float64 `json:"default_threshold_7d"` } type OpsDataRetentionSettings struct { diff --git a/backend/internal/service/pricing_service_test.go b/backend/internal/service/pricing_service_test.go index 24e64803f64e..30c56668e252 100644 --- a/backend/internal/service/pricing_service_test.go +++ b/backend/internal/service/pricing_service_test.go @@ -149,9 +149,9 @@ func TestDefaultPricingIncludesCodexAutoReview(t *testing.T) { got := svc.GetModelPricing("codex-auto-review") require.NotNil(t, got) - require.InDelta(t, 2.5e-6, got.InputCostPerToken, 1e-12) - require.InDelta(t, 1.5e-5, got.OutputCostPerToken, 1e-12) - require.InDelta(t, 2.5e-7, got.CacheReadInputTokenCost, 1e-12) + require.InDelta(t, 5e-6, got.InputCostPerToken, 1e-12) + require.InDelta(t, 3e-5, got.OutputCostPerToken, 1e-12) + require.InDelta(t, 5e-7, got.CacheReadInputTokenCost, 1e-12) } func TestGetModelPricing_Gpt54MiniUsesDedicatedStaticFallbackWhenRemoteMissing(t *testing.T) { diff --git a/backend/internal/service/ratelimit_service_openai_test.go b/backend/internal/service/ratelimit_service_openai_test.go index aa5a070c74fc..107ac27ef414 100644 --- a/backend/internal/service/ratelimit_service_openai_test.go +++ b/backend/internal/service/ratelimit_service_openai_test.go @@ -51,7 +51,7 @@ func TestCalculateOpenAI429ResetTime_5hExhausted(t *testing.T) { headers.Set("x-codex-primary-used-percent", "50") headers.Set("x-codex-primary-reset-after-seconds", "500000") headers.Set("x-codex-primary-window-minutes", "10080") // 7 days - headers.Set("x-codex-secondary-used-percent", "100") + headers.Set("x-codex-secondary-used-percent", "0") headers.Set("x-codex-secondary-reset-after-seconds", "3600") // 1 hour headers.Set("x-codex-secondary-window-minutes", "300") // 5 hours @@ -122,7 +122,7 @@ func TestCalculateOpenAI429ResetTime_ReversedWindowOrder(t *testing.T) { // Test when OpenAI sends primary as 5h and secondary as 7d (reversed) headers := http.Header{} - headers.Set("x-codex-primary-used-percent", "100") // This is 5h + headers.Set("x-codex-primary-used-percent", "0") // This is 5h remaining% headers.Set("x-codex-primary-reset-after-seconds", "3600") // 1 hour headers.Set("x-codex-primary-window-minutes", "300") // 5 hours - smaller! headers.Set("x-codex-secondary-used-percent", "50") @@ -180,7 +180,7 @@ func TestHandle429_OpenAIPersistsCodexSnapshotImmediately(t *testing.T) { headers.Set("x-codex-primary-used-percent", "100") headers.Set("x-codex-primary-reset-after-seconds", "604800") headers.Set("x-codex-primary-window-minutes", "10080") - headers.Set("x-codex-secondary-used-percent", "100") + headers.Set("x-codex-secondary-used-percent", "0") headers.Set("x-codex-secondary-reset-after-seconds", "18000") headers.Set("x-codex-secondary-window-minutes", "300") @@ -224,7 +224,7 @@ func TestNormalizedCodexLimits(t *testing.T) { pUsed := 100.0 pReset := 384607 pWindow := 10080 - sUsed := 3.0 + sRemaining := 3.0 sReset := 17369 sWindow := 300 @@ -232,7 +232,7 @@ func TestNormalizedCodexLimits(t *testing.T) { PrimaryUsedPercent: &pUsed, PrimaryResetAfterSeconds: &pReset, PrimaryWindowMinutes: &pWindow, - SecondaryUsedPercent: &sUsed, + SecondaryUsedPercent: &sRemaining, SecondaryResetAfterSeconds: &sReset, SecondaryWindowMinutes: &sWindow, } @@ -249,8 +249,8 @@ func TestNormalizedCodexLimits(t *testing.T) { if normalized.Reset7dSeconds == nil || *normalized.Reset7dSeconds != 384607 { t.Errorf("expected Reset7dSeconds=384607, got %v", normalized.Reset7dSeconds) } - if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 3.0 { - t.Errorf("expected Used5hPercent=3, got %v", normalized.Used5hPercent) + if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 97.0 { + t.Errorf("expected Used5hPercent=97, got %v", normalized.Used5hPercent) } if normalized.Reset5hSeconds == nil || *normalized.Reset5hSeconds != 17369 { t.Errorf("expected Reset5hSeconds=17369, got %v", normalized.Reset5hSeconds) @@ -338,11 +338,11 @@ func TestRateLimitService_HandleUpstreamError_403FallsBackToRawBody(t *testing.T func TestNormalizedCodexLimits_OnlySecondaryData(t *testing.T) { // Test when only secondary has data, no window_minutes - sUsed := 60.0 + sRemaining := 60.0 sReset := 3000 snapshot := &OpenAICodexUsageSnapshot{ - SecondaryUsedPercent: &sUsed, + SecondaryUsedPercent: &sRemaining, SecondaryResetAfterSeconds: &sReset, // No window_minutes, no primary data } @@ -354,8 +354,8 @@ func TestNormalizedCodexLimits_OnlySecondaryData(t *testing.T) { // Legacy assumption: primary=7d, secondary=5h // So secondary goes to 5h - if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 60.0 { - t.Errorf("expected Used5hPercent=60, got %v", normalized.Used5hPercent) + if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 40.0 { + t.Errorf("expected Used5hPercent=40, got %v", normalized.Used5hPercent) } if normalized.Reset5hSeconds == nil || *normalized.Reset5hSeconds != 3000 { t.Errorf("expected Reset5hSeconds=3000, got %v", normalized.Reset5hSeconds) @@ -370,13 +370,13 @@ func TestNormalizedCodexLimits_BothDataNoWindowMinutes(t *testing.T) { // Test when both have data but no window_minutes pUsed := 100.0 pReset := 400000 - sUsed := 50.0 + sRemaining := 30.0 sReset := 10000 snapshot := &OpenAICodexUsageSnapshot{ PrimaryUsedPercent: &pUsed, PrimaryResetAfterSeconds: &pReset, - SecondaryUsedPercent: &sUsed, + SecondaryUsedPercent: &sRemaining, SecondaryResetAfterSeconds: &sReset, // No window_minutes } @@ -393,8 +393,8 @@ func TestNormalizedCodexLimits_BothDataNoWindowMinutes(t *testing.T) { if normalized.Reset7dSeconds == nil || *normalized.Reset7dSeconds != 400000 { t.Errorf("expected Reset7dSeconds=400000, got %v", normalized.Reset7dSeconds) } - if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 50.0 { - t.Errorf("expected Used5hPercent=50, got %v", normalized.Used5hPercent) + if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 70.0 { + t.Errorf("expected Used5hPercent=70, got %v", normalized.Used5hPercent) } if normalized.Reset5hSeconds == nil || *normalized.Reset5hSeconds != 10000 { t.Errorf("expected Reset5hSeconds=10000, got %v", normalized.Reset5hSeconds) @@ -425,7 +425,7 @@ func TestCalculateOpenAI429ResetTime_UserProvidedScenario(t *testing.T) { // This is the exact scenario from the user: // codex_7d_used_percent: 100 // codex_7d_reset_after_seconds: 384607 (约4.5天后重置) - // codex_5h_used_percent: 3 + // codex_5h_used_percent: 97 (from upstream 3% remaining) // codex_5h_reset_after_seconds: 17369 (约4.8小时后重置) svc := &RateLimitService{} diff --git a/backend/internal/service/scheduler_snapshot_hydration_test.go b/backend/internal/service/scheduler_snapshot_hydration_test.go index 778cab235165..0a1d0a0a52e7 100644 --- a/backend/internal/service/scheduler_snapshot_hydration_test.go +++ b/backend/internal/service/scheduler_snapshot_hydration_test.go @@ -6,6 +6,8 @@ import ( "context" "testing" "time" + + "github.com/Wei-Shaw/sub2api/internal/config" ) type snapshotHydrationCache struct { @@ -186,3 +188,89 @@ func TestGatewaySelectAccountWithLoadAwareness_HydratesSelectedAccountFromSchedu t.Fatalf("expected hydrated api key, got %q", got) } } + +func TestGatewaySelectAccountWithLoadAwareness_SkipsAntigravityGeminiFamilyRateLimitedSnapshot(t *testing.T) { + resetAt := time.Now().Add(10 * time.Minute).Format(time.RFC3339) + cache := &snapshotHydrationCache{ + snapshot: []*Account{ + { + ID: 1, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 1, + AccountGroups: []AccountGroup{ + {AccountID: 1, GroupID: 22}, + }, + GroupIDs: []int64{22}, + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + antigravityGeminiModelRateLimitKey: map[string]any{ + "rate_limit_reset_at": resetAt, + }, + }, + }, + }, + { + ID: 2, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 2, + AccountGroups: []AccountGroup{ + {AccountID: 2, GroupID: 22}, + }, + GroupIDs: []int64{22}, + Extra: map[string]any{ + "mixed_scheduling": true, + }, + }, + }, + accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformAntigravity, Type: AccountTypeOAuth}, + 2: {ID: 2, Platform: PlatformAntigravity, Type: AccountTypeOAuth}, + }, + } + groupID := int64(22) + svc := &GatewayService{ + schedulerSnapshot: NewSchedulerSnapshotService(cache, nil, nil, nil, nil), + groupRepo: &mockGroupRepoForGateway{ + groups: map[int64]*Group{ + groupID: { + ID: groupID, + Platform: PlatformGemini, + Status: StatusActive, + Hydrated: true, + }, + }, + }, + concurrencyService: NewConcurrencyService(&mockConcurrencyCache{}), + cfg: &config.Config{ + Gateway: config.GatewayConfig{ + Scheduling: config.GatewaySchedulingConfig{ + LoadBatchEnabled: true, + StickySessionMaxWaiting: 3, + StickySessionWaitTimeout: time.Second, + FallbackWaitTimeout: time.Second, + FallbackMaxWaiting: 10, + }, + }, + }, + } + + result, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gemini-3-flash-preview", nil, "", 0) + if err != nil { + t.Fatalf("SelectAccountWithLoadAwareness error: %v", err) + } + if result == nil || result.Account == nil { + t.Fatalf("expected selected account") + } + if result.Account.ID != 2 { + t.Fatalf("expected scheduler to skip Gemini-family limited antigravity account 1, got %d", result.Account.ID) + } +} diff --git a/backend/internal/service/setting_service.go b/backend/internal/service/setting_service.go index 08c0d045c6e0..98acdb80f56d 100644 --- a/backend/internal/service/setting_service.go +++ b/backend/internal/service/setting_service.go @@ -137,6 +137,11 @@ type cachedOpenAICodexUserAgent struct { expiresAt int64 // unix nano } +type cachedOpenAIQuotaAutoPauseSettings struct { + settings OpsOpenAIAccountQuotaAutoPauseSettings + expiresAt int64 +} + const openAICodexUserAgentCacheTTL = 60 * time.Second const openAICodexUserAgentErrorTTL = 5 * time.Second const openAICodexUserAgentDBTimeout = 5 * time.Second @@ -152,6 +157,12 @@ const openAIAllowCodexPluginCacheTTL = 60 * time.Second const openAIAllowCodexPluginErrorTTL = 5 * time.Second const openAIAllowCodexPluginDBTimeout = 5 * time.Second +const openAIQuotaAutoPauseSettingsCacheTTL = 60 * time.Second +const openAIQuotaAutoPauseSettingsErrorTTL = 5 * time.Second +const openAIQuotaAutoPauseSettingsDBTimeout = 5 * time.Second + +const openAIQuotaAutoPauseSettingsRefreshKey = "openai_quota_auto_pause_settings" + // DefaultSubscriptionGroupReader validates group references used by default subscriptions. type DefaultSubscriptionGroupReader interface { GetByID(ctx context.Context, id int64) (*Group, error) @@ -176,6 +187,15 @@ type SettingService struct { openAICodexUASF singleflight.Group openAIAllowCodexPluginCache atomic.Value // *cachedOpenAIAllowCodexPlugin openAIAllowCodexPluginSF singleflight.Group + + // openAIQuotaAutoPauseSettingsCache holds the most recently observed quota auto-pause + // settings. GetOpenAIQuotaAutoPauseSettings reads this atomic.Value on the request hot + // path without ever blocking on the DB; when the cached entry expires, a background + // goroutine refreshes it via openAIQuotaAutoPauseSettingsSF (stale-while-revalidate). + // This per-service field also gives tests natural isolation — each SettingService + // instance owns its own cache, no shared package-level state. + openAIQuotaAutoPauseSettingsCache atomic.Value // *cachedOpenAIQuotaAutoPauseSettings + openAIQuotaAutoPauseSettingsSF singleflight.Group } // DefaultPlatformQuotaSetting 单 platform 三档限额（nil = 沿用上层；0 = 显式禁用；>0 = 上限） @@ -2027,6 +2047,17 @@ func (s *SettingService) refreshCachedSettings(settings *SystemSettings) { enabled: settings.OpenAIAdvancedSchedulerEnabled, expiresAt: time.Now().Add(openAIAdvancedSchedulerSettingCacheTTL).UnixNano(), }) + // Invalidate the quota auto-pause cache and let the next read trigger a fresh load. + // We can't know from here whether ops_advanced_settings was also touched, so be + // defensive: store an expired entry — GetOpenAIQuotaAutoPauseSettings will serve + // stale and kick off an async refresh, never blocking the request that follows. + s.openAIQuotaAutoPauseSettingsSF.Forget(openAIQuotaAutoPauseSettingsRefreshKey) + if cached, _ := s.openAIQuotaAutoPauseSettingsCache.Load().(*cachedOpenAIQuotaAutoPauseSettings); cached != nil { + s.openAIQuotaAutoPauseSettingsCache.Store(&cachedOpenAIQuotaAutoPauseSettings{ + settings: cached.settings, + expiresAt: 0, + }) + } if s.cfg != nil { s.cfg.SetTrustForwardedIPForAPIKeyACL(settings.APIKeyACLTrustForwardedIP) } @@ -4448,6 +4479,106 @@ func (s *SettingService) GetClaudeCodeVersionBounds(ctx context.Context) (min, m return b.min, b.max } +// GetOpenAIQuotaAutoPauseSettings returns the current global default quota auto-pause +// settings. It is invoked on the OpenAI scheduling hot path (once per request) and is +// therefore designed to never block on the DB: +// +// - Fresh cached value → returned immediately. +// - Stale or empty cache → the last known value is returned, and a background +// goroutine refreshes the cache via singleflight (stale-while-revalidate). +// - First call with no cache yet → zero defaults are returned and the same async +// refresh is kicked off; the next call gets the freshly populated value. +// +// Callers that need the freshly persisted value synchronously (tests, post-update +// confirmation, optional startup warm-up) should call WarmOpenAIQuotaAutoPauseSettings. +func (s *SettingService) GetOpenAIQuotaAutoPauseSettings(ctx context.Context) OpsOpenAIAccountQuotaAutoPauseSettings { + if s == nil { + return OpsOpenAIAccountQuotaAutoPauseSettings{} + } + cached, _ := s.openAIQuotaAutoPauseSettingsCache.Load().(*cachedOpenAIQuotaAutoPauseSettings) + now := time.Now().UnixNano() + if cached != nil && now < cached.expiresAt { + return cached.settings + } + // Stale or unset: trigger background refresh without blocking this request. + // singleflight.DoChan dedupes concurrent refreshes; we deliberately ignore the + // returned channel — the result is observable via the atomic cache. + s.openAIQuotaAutoPauseSettingsSF.DoChan(openAIQuotaAutoPauseSettingsRefreshKey, func() (any, error) { + s.refreshOpenAIQuotaAutoPauseSettings(context.Background()) + return nil, nil + }) + if cached != nil { + return cached.settings // serve stale value while revalidating + } + return OpsOpenAIAccountQuotaAutoPauseSettings{} +} + +// WarmOpenAIQuotaAutoPauseSettings synchronously loads the quota auto-pause settings +// into the in-memory cache. Useful for application startup (so the first request hits +// a warm cache) and for tests that need deterministic reads immediately after +// constructing the service. +func (s *SettingService) WarmOpenAIQuotaAutoPauseSettings(ctx context.Context) OpsOpenAIAccountQuotaAutoPauseSettings { + if s == nil { + return OpsOpenAIAccountQuotaAutoPauseSettings{} + } + s.refreshOpenAIQuotaAutoPauseSettings(ctx) + cached, _ := s.openAIQuotaAutoPauseSettingsCache.Load().(*cachedOpenAIQuotaAutoPauseSettings) + if cached == nil { + return OpsOpenAIAccountQuotaAutoPauseSettings{} + } + return cached.settings +} + +// refreshOpenAIQuotaAutoPauseSettings reads the latest settings from the DB and stores +// them into the in-memory cache. On error it stores the prior value (or zero defaults +// if nothing is cached yet) with the shorter error TTL so the next refresh comes +// sooner. Always uses its own timeout-bounded context to keep refresh latency +// predictable regardless of the caller. +func (s *SettingService) refreshOpenAIQuotaAutoPauseSettings(ctx context.Context) { + if s == nil || s.settingRepo == nil { + return + } + dbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), openAIQuotaAutoPauseSettingsDBTimeout) + defer cancel() + + settings := OpsOpenAIAccountQuotaAutoPauseSettings{} + ttl := openAIQuotaAutoPauseSettingsCacheTTL + raw, err := s.settingRepo.GetValue(dbCtx, SettingKeyOpsAdvancedSettings) + if err == nil { + cfg := defaultOpsAdvancedSettings() + if strings.TrimSpace(raw) != "" { + if jsonErr := json.Unmarshal([]byte(raw), cfg); jsonErr == nil { + normalizeOpsAdvancedSettings(cfg) + } + } + settings = cfg.OpenAIAccountQuotaAutoPause + } else if !errors.Is(err, ErrSettingNotFound) { + // Real error: keep serving prior value but refresh sooner. + if prior, _ := s.openAIQuotaAutoPauseSettingsCache.Load().(*cachedOpenAIQuotaAutoPauseSettings); prior != nil { + settings = prior.settings + } + ttl = openAIQuotaAutoPauseSettingsErrorTTL + } + + s.openAIQuotaAutoPauseSettingsCache.Store(&cachedOpenAIQuotaAutoPauseSettings{ + settings: settings, + expiresAt: time.Now().Add(ttl).UnixNano(), + }) +} + +// SetOpenAIQuotaAutoPauseSettings writes the given settings directly into the in-memory +// cache. Called from settings-write code paths so that the next read reflects the new +// value immediately, without waiting for the background refresh. +func (s *SettingService) SetOpenAIQuotaAutoPauseSettings(settings OpsOpenAIAccountQuotaAutoPauseSettings) { + if s == nil { + return + } + s.openAIQuotaAutoPauseSettingsCache.Store(&cachedOpenAIQuotaAutoPauseSettings{ + settings: settings, + expiresAt: time.Now().Add(openAIQuotaAutoPauseSettingsCacheTTL).UnixNano(), + }) +} + // GetRectifierSettings 获取请求整流器配置 func (s *SettingService) GetRectifierSettings(ctx context.Context) (*RectifierSettings, error) { value, err := s.settingRepo.GetValue(ctx, SettingKeyRectifierSettings) diff --git a/backend/internal/service/user.go b/backend/internal/service/user.go index f98336111bac..edb944ee05cd 100644 --- a/backend/internal/service/user.go +++ b/backend/internal/service/user.go @@ -32,6 +32,7 @@ type User struct { LastUsedAt *time.Time CreatedAt time.Time UpdatedAt time.Time + DeletedAt *time.Time // 非 nil 表示用户已软删除 // GroupRates 用户专属分组倍率配置 // map[groupID]rateMultiplier diff --git a/backend/internal/service/user_msg_queue_service.go b/backend/internal/service/user_msg_queue_service.go index a0ce95a8638e..f3f105ac9386 100644 --- a/backend/internal/service/user_msg_queue_service.go +++ b/backend/internal/service/user_msg_queue_service.go @@ -12,6 +12,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/tidwall/gjson" ) // UserMsgQueueCache 用户消息串行队列 Redis 缓存接口 @@ -62,43 +63,48 @@ func NewUserMessageQueueService(cache UserMsgQueueCache, rpmCache RPMCache, cfg // 2. 最后一条消息 role == "user" // 3. 最后一条消息 content（如果是数组）中不含 type:"tool_result" / "tool_use_result" func IsRealUserMessage(parsed *ParsedRequest) bool { - if parsed == nil || len(parsed.Messages) == 0 { + if parsed == nil { return false } - - lastMsg := parsed.Messages[len(parsed.Messages)-1] - msgMap, ok := lastMsg.(map[string]any) - if !ok { + messagesRaw := parsed.MessagesRaw() + if len(messagesRaw) == 0 { return false } - role, _ := msgMap["role"].(string) - if role != "user" { + messages := gjson.ParseBytes(messagesRaw) + if !messages.IsArray() { return false } - - // 检查 content 是否包含 tool_result 类型 - content, ok := msgMap["content"] - if !ok { - return true // 没有 content 字段，视为普通用户消息 + lastMsg := gjson.Result{} + messages.ForEach(func(_, msg gjson.Result) bool { + lastMsg = msg + return true + }) + if !lastMsg.Exists() || !lastMsg.IsObject() { + return false + } + if lastMsg.Get("role").String() != "user" { + return false } - contentArr, ok := content.([]any) - if !ok { - return true // content 不是数组（可能是 string），视为普通用户消息 + content := lastMsg.Get("content") + if !content.Exists() { + return true + } + if !content.IsArray() { + return true } - for _, item := range contentArr { - itemMap, ok := item.(map[string]any) - if !ok { - continue - } - itemType, _ := itemMap["type"].(string) + isReal := true + content.ForEach(func(_, item gjson.Result) bool { + itemType := item.Get("type").String() if itemType == "tool_result" || itemType == "tool_use_result" { + isReal = false return false } - } - return true + return true + }) + return isReal } // TryAcquire 尝试立即获取串行锁 diff --git a/backend/internal/service/user_platform_quota_flusher.go b/backend/internal/service/user_platform_quota_flusher.go new file mode 100644 index 000000000000..3ee23d2c9372 --- /dev/null +++ b/backend/internal/service/user_platform_quota_flusher.go @@ -0,0 +1,267 @@ +package service + +import ( + "context" + "errors" + "sync/atomic" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" +) + +// quotaDirtyCache 是 flusher 依赖的窄接口（来自 BillingCache）。 +type quotaDirtyCache interface { + PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) + ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error + BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) +} + +// quotaSnapshotWriter 是 flusher 依赖的 DB 写入窄接口。 +// 使用 service 层的 UserPlatformQuotaSnapshot，避免与 repository 包形成循环依赖； +// 实际实现由 repository adapter 在 B7 注入。 +type quotaSnapshotWriter interface { + BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error +} + +// FlusherMetrics 记录 flusher 运行时指标（原子量，零值可用）。 +type FlusherMetrics struct { + FlushSuccessTotal atomic.Int64 + FlushErrorTotal atomic.Int64 + FlushBatchSizeTotal atomic.Int64 + FlushLatencyMsMax atomic.Int64 + DirtyReaddTotal atomic.Int64 + // DirtyLostTotal：Readd 失败导致脏 key 丢失——已 SPOP+主操作失败+Readd 也失败； + // Redis 仍权威，活跃 key 下次 SADD 自愈。 + DirtyLostTotal atomic.Int64 + FlushFKViolationTotal atomic.Int64 +} + +// flusherMaxBatchesPerTick 单次 tick 最多消费的批数，防止 tick 执行时间过长。 +const flusherMaxBatchesPerTick = 16 + +// maxFlushBatchSize 限制单批行数,必须 ≤ repository.BatchSnapshotUsage 的 batchRows(6000), +// 以保证单次 flush 的 snapshots 仅生成一条 UPSERT(单事务原子)。两处需手动保持一致。 +const maxFlushBatchSize = 6000 + +// defaultFlushBatchSize 是配置 flush_batch_size 非法(≤0)时的回退值。 +const defaultFlushBatchSize = 1000 + +// UserPlatformQuotaUsageFlusher 将 Redis 脏集快照定期批量写入 DB。 +// 不维护任何 delta/in-process 状态；每批读取 Redis 当前绝对值覆盖写入。 +type UserPlatformQuotaUsageFlusher struct { + cache quotaDirtyCache + quotaRepo quotaSnapshotWriter + timingWheel *TimingWheelService + // enabled 对应 flusher_enabled 配置；false 时 Start() 不注册定时器。 + enabled bool + interval time.Duration + batchSize int + flushTimeout time.Duration + metrics *FlusherMetrics + stopped atomic.Bool +} + +// NewUserPlatformQuotaUsageFlusher 创建 UserPlatformQuotaUsageFlusher。 +// cache(BillingCache) 隐式满足 quotaDirtyCache；quotaRepo(UserPlatformQuotaRepository) 隐式满足 quotaSnapshotWriter。 +func NewUserPlatformQuotaUsageFlusher(cfg *config.Config, cache BillingCache, quotaRepo UserPlatformQuotaRepository, tw *TimingWheelService) *UserPlatformQuotaUsageFlusher { + batchSize := cfg.Database.UserPlatformQuotaFlushBatchSize + if batchSize <= 0 { + batchSize = defaultFlushBatchSize + } + if batchSize > maxFlushBatchSize { + logger.LegacyPrintf("quota_flusher", + "[QuotaFlusher] flush_batch_size %d 超过上限 %d,已 clamp(避免 BatchSnapshotUsage 多子批非原子)", + cfg.Database.UserPlatformQuotaFlushBatchSize, maxFlushBatchSize) + batchSize = maxFlushBatchSize + } + interval := time.Duration(cfg.Database.UserPlatformQuotaFlushIntervalMs) * time.Millisecond + if interval <= 0 { + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] flush_interval_ms %d 非法,回退 2000ms", cfg.Database.UserPlatformQuotaFlushIntervalMs) + interval = 2 * time.Second + } + return &UserPlatformQuotaUsageFlusher{ + cache: cache, + quotaRepo: quotaRepo, + timingWheel: tw, + enabled: cfg.Database.UserPlatformQuotaFlusherEnabled, + interval: interval, + batchSize: batchSize, + flushTimeout: 3 * time.Second, + metrics: &FlusherMetrics{}, + } +} + +// updateLatencyMax 用 CAS 单调更新最大延迟。 +func (s *UserPlatformQuotaUsageFlusher) updateLatencyMax(ms int64) { + for { + old := s.metrics.FlushLatencyMsMax.Load() + if ms <= old { + return + } + if s.metrics.FlushLatencyMsMax.CompareAndSwap(old, ms) { + return + } + } +} + +// readdOrCountLost 尝试把 keys 回填脏集：成功计 DirtyReaddTotal，失败计 DirtyLostTotal 并 ALERT。 +func (s *UserPlatformQuotaUsageFlusher) readdOrCountLost(ctx context.Context, keys []UserPlatformQuotaKey, stage string) { + if err := s.cache.ReaddDirtyUserPlatformQuotaKeys(ctx, keys); err != nil { + s.metrics.DirtyLostTotal.Add(int64(len(keys))) + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] ALERT: Readd after %s failed, %d keys 丢出脏集(DB 镜像缺这批,Redis 仍权威,活跃 key 下次 SADD 自愈): %v", stage, len(keys), err) + return + } + s.metrics.DirtyReaddTotal.Add(int64(len(keys))) +} + +// flushOneBatch 处理单批：Pop → BatchGet → 组装 snaps → BatchSnapshotUsage。 +// 返回 (shouldContinue bool)：false 表示本轮循环应停止（空集/错误/最后一批）。 +// 每次调用独立创建带 timeout 的 ctx 并 defer cancel，不会在循环中累积泄漏。 +func (s *UserPlatformQuotaUsageFlusher) flushOneBatch(parentCtx context.Context) bool { + ctx, cancel := context.WithTimeout(parentCtx, s.flushTimeout) + defer cancel() + + // 1. Pop 脏集 + keys, err := s.cache.PopDirtyUserPlatformQuotaKeys(ctx, s.batchSize) + if err != nil { + s.metrics.FlushErrorTotal.Add(1) + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] PopDirty error: %v", err) + return false + } + if len(keys) == 0 { + // 脏集已空 + return false + } + + // 2. 批量读 Redis 快照 + entries, err := s.cache.BatchGetUserPlatformQuotaCache(ctx, keys) + if err != nil { + s.metrics.FlushErrorTotal.Add(1) + s.readdOrCountLost(ctx, keys, "BatchGet") + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] BatchGet error: %v", err) + return false + } + + // 3. 组装 snapshots（MISS 或任一 WindowStart==nil → 跳过） + snaps := make([]UserPlatformQuotaSnapshot, 0, len(keys)) + for i, key := range keys { + e := entries[i] + if e == nil { + continue + } + if e.DailyWindowStart == nil || e.WeeklyWindowStart == nil || e.MonthlyWindowStart == nil { + continue + } + snaps = append(snaps, UserPlatformQuotaSnapshot{ + UserID: key.UserID, + Platform: key.Platform, + DailyUsageUSD: e.DailyUsageUSD, + WeeklyUsageUSD: e.WeeklyUsageUSD, + MonthlyUsageUSD: e.MonthlyUsageUSD, + DailyWindowStart: *e.DailyWindowStart, + WeeklyWindowStart: *e.WeeklyWindowStart, + MonthlyWindowStart: *e.MonthlyWindowStart, + }) + } + + // 4. 全部 MISS/异常跳过时 + if len(snaps) == 0 { + // 若 Pop 数量已不满一批，表示脏集将空，停止 + if len(keys) < s.batchSize { + return false + } + // 否则继续下一批（可能还有更多脏 key） + return true + } + + // 已知竞态(admin 写 × flusher 刷,仅 flusher_enabled=true 时存在): + // admin ResetExpiredWindow/UpsertForUser 是"先写 DB 再 DeleteCache"。若本批已 SPOP + BatchGet + // 读到旧 usage 快照(此刻 member 已离开脏集),而 admin 随后写 DB、本行 UPSERT 又在 admin 写之后落库, + // 则旧快照会覆盖 admin 刚写入的值;DeleteCache 后 Redis MISS,下次 preflight 从 DB 重载被覆盖的旧值。 + // 因 member 已被 SPOP,admin 侧 SREM/清脏标记无法拦截本批(故未做)。影响有限,暂列为已知取舍: + // - UpsertForUser 改 limit,而本 UPSERT 不写 limit 列 → limit 配置不受影响; + // - ResetExpiredWindow 改 usage,但 preflight windowExpired 会在窗口真正过期时自愈重置, + // 仅"强制重置未过期窗口"且与本批精确交错时短暂失效; + // - 低频 admin 操作 + 默认 flusher_enabled=false。彻底消除需 version OCC(DB 加 version 列条件 UPSERT), + // 成本高;启用 flusher 后如需强一致再评估。 + + // 5. 写入 DB + start := time.Now() + writeErr := s.quotaRepo.BatchSnapshotUsage(ctx, snaps, time.Now().UTC()) + s.updateLatencyMax(time.Since(start).Milliseconds()) + + if writeErr != nil { + if errors.Is(writeErr, ErrUserPlatformQuotaFKViolation) { + // 注意:PG FK violation 是整条 INSERT 回滚 → 整批(含同批正常用户)均未写入 DB, + // 且这些 key 已被 SPOP 出脏集、此处不 Readd。活跃 key 会在下次请求重新 SADD, + // flusher 读 Redis 当前累计绝对值刷库即自愈;低活跃 key 这轮 DB usage 偏低 + // (Redis 仍是 enforcement 权威,不受影响;DB 仅展示)。已删用户边角的接受取舍,不做逐行重试。 + // FK 违反：用户已被删除，直接丢弃不 Readd + s.metrics.FlushFKViolationTotal.Add(1) + s.metrics.FlushErrorTotal.Add(1) + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] FK violation (dropped %d snaps): %v", len(snaps), writeErr) + } else { + // 其他错误：回填脏集，保留下次重试 + s.metrics.FlushErrorTotal.Add(1) + s.readdOrCountLost(ctx, keys, "BatchSnapshotUsage") + logger.LegacyPrintf("quota_flusher", "[QuotaFlusher] BatchSnapshotUsage error: %v", writeErr) + } + return false + } + + // 6. 成功 + s.metrics.FlushSuccessTotal.Add(1) + s.metrics.FlushBatchSizeTotal.Add(int64(len(snaps))) + + // 若 Pop 数量不满一批，脏集已空，停止 + if len(keys) < s.batchSize { + return false + } + return true +} + +// flush 执行一次完整的 flush，循环消费至脏集空或达到 maxBatchesPerTick。 +func (s *UserPlatformQuotaUsageFlusher) flush() { + if s == nil { + return + } + parentCtx := context.Background() + for b := 0; b < flusherMaxBatchesPerTick; b++ { + if !s.flushOneBatch(parentCtx) { + return + } + } + // 连续消费满 flusherMaxBatchesPerTick 批仍未取空脏集:本 tick 主动让出,剩余积压留待下一 tick。 + // 记一条 log 便于 oncall 发现 distinct 活跃 key 远超 maxBatchesPerTick×batchSize(DB 镜像延迟上升); + // 可配合 Redis SCARD billing:upq:dirty 观察脏集存量。 + logger.LegacyPrintf("quota_flusher", + "[QuotaFlusher] 单 tick 达到 max batches 上限(%d × batchSize=%d),脏集仍非空,积压顺延至下一 tick", + flusherMaxBatchesPerTick, s.batchSize) +} + +// tick 是 TimingWheel 回调。若 flusher 已停止则直接返回。 +func (s *UserPlatformQuotaUsageFlusher) tick() { + if s == nil || s.stopped.Load() { + return + } + s.flush() +} + +// Start 注册定时 tick。flusher_enabled=false 时直接返回，不注册定时器。 +func (s *UserPlatformQuotaUsageFlusher) Start() { + if s == nil || !s.enabled { + return + } + s.timingWheel.ScheduleRecurring("deferred:platform_quota", s.interval, s.tick) +} + +// Stop 停止 flusher：标记 stopped → Cancel 定时器 → 执行最后一次 flush。 +func (s *UserPlatformQuotaUsageFlusher) Stop() { + if s == nil { + return + } + s.stopped.Store(true) + s.timingWheel.Cancel("deferred:platform_quota") + s.flush() +} diff --git a/backend/internal/service/user_platform_quota_flusher_test.go b/backend/internal/service/user_platform_quota_flusher_test.go new file mode 100644 index 000000000000..4f7344819ae0 --- /dev/null +++ b/backend/internal/service/user_platform_quota_flusher_test.go @@ -0,0 +1,511 @@ +package service + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" +) + +// --------------------------------------------------------------------------- +// Mock: quotaDirtyCache +// --------------------------------------------------------------------------- + +type mockQuotaDirtyCache struct { + // popSequence: 第 0 次 Pop 返回 popSequence[0]，之后返回 nil（空集） + popSequence [][]UserPlatformQuotaKey + popCallIdx int + + // getEntries: BatchGetUserPlatformQuotaCache 返回的 entries（与 keys 对齐） + getEntries []*UserPlatformQuotaCacheEntry + getErr error + + // readdCalled: 记录 Readd 收到的 keys（累积所有次调用） + readdCalled [][]UserPlatformQuotaKey + readdErr error +} + +func (m *mockQuotaDirtyCache) PopDirtyUserPlatformQuotaKeys(_ context.Context, _ int) ([]UserPlatformQuotaKey, error) { + if m.popCallIdx < len(m.popSequence) { + keys := m.popSequence[m.popCallIdx] + m.popCallIdx++ + return keys, nil + } + // 超出序列 → 空集（模拟脏集已清空） + return nil, nil +} + +func (m *mockQuotaDirtyCache) ReaddDirtyUserPlatformQuotaKeys(_ context.Context, keys []UserPlatformQuotaKey) error { + m.readdCalled = append(m.readdCalled, keys) + return m.readdErr +} + +func (m *mockQuotaDirtyCache) BatchGetUserPlatformQuotaCache(_ context.Context, _ []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + if m.getErr != nil { + return nil, m.getErr + } + return m.getEntries, nil +} + +// --------------------------------------------------------------------------- +// Mock: quotaSnapshotWriter +// --------------------------------------------------------------------------- + +type mockQuotaSnapshotWriter struct { + receivedSnaps []UserPlatformQuotaSnapshot + returnErr error +} + +func (m *mockQuotaSnapshotWriter) BatchSnapshotUsage(_ context.Context, snaps []UserPlatformQuotaSnapshot, _ time.Time) error { + m.receivedSnaps = append(m.receivedSnaps, snaps...) + return m.returnErr +} + +// --------------------------------------------------------------------------- +// Helper: 构造窗口起始时间（非 nil） +// --------------------------------------------------------------------------- + +func flusherPtrTime(t time.Time) *time.Time { return &t } + +func makeEntry(daily, weekly, monthly float64) *UserPlatformQuotaCacheEntry { + now := time.Now().UTC() + return &UserPlatformQuotaCacheEntry{ + DailyUsageUSD: daily, + WeeklyUsageUSD: weekly, + MonthlyUsageUSD: monthly, + DailyWindowStart: flusherPtrTime(now), + WeeklyWindowStart: flusherPtrTime(now), + MonthlyWindowStart: flusherPtrTime(now), + } +} + +// --------------------------------------------------------------------------- +// newTestFlusher: 直接构造 struct（跳过构造函数，B7 才注入） +// --------------------------------------------------------------------------- + +func newTestFlusher(cache quotaDirtyCache, writer quotaSnapshotWriter) *UserPlatformQuotaUsageFlusher { + return &UserPlatformQuotaUsageFlusher{ + cache: cache, + quotaRepo: writer, + timingWheel: nil, // 单测不启动 TimingWheel + interval: 5 * time.Second, + batchSize: 100, + flushTimeout: 5 * time.Second, + metrics: &FlusherMetrics{}, + } +} + +// --------------------------------------------------------------------------- +// 场景 1: PopSnapshotUpsert — 2 key + 2 个含 window 的 entry → writer 收 2 行 +// --------------------------------------------------------------------------- + +func TestFlusher_PopSnapshotUpsert(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 1, Platform: "anthropic"}, + {UserID: 2, Platform: "openai"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, // 第 1 次返回 keys，之后空 + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + makeEntry(4.0, 5.0, 6.0), + }, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + f.flush() + + if len(writer.receivedSnaps) != 2 { + t.Fatalf("expected 2 snaps, got %d", len(writer.receivedSnaps)) + } + if f.metrics.FlushBatchSizeTotal.Load() != 2 { + t.Errorf("FlushBatchSizeTotal = %d, want 2", f.metrics.FlushBatchSizeTotal.Load()) + } + if f.metrics.FlushSuccessTotal.Load() != 1 { + t.Errorf("FlushSuccessTotal = %d, want 1", f.metrics.FlushSuccessTotal.Load()) + } + if f.metrics.FlushErrorTotal.Load() != 0 { + t.Errorf("FlushErrorTotal = %d, want 0", f.metrics.FlushErrorTotal.Load()) + } +} + +// --------------------------------------------------------------------------- +// 场景 2: MissKeySkipped — 2 key，BatchGet 返回 [entry, nil] → 只刷 1 行，nil 跳过，不 Readd +// --------------------------------------------------------------------------- + +func TestFlusher_MissKeySkipped(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 1, Platform: "anthropic"}, + {UserID: 2, Platform: "openai"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + nil, // MISS + }, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + f.flush() + + if len(writer.receivedSnaps) != 1 { + t.Fatalf("expected 1 snap, got %d", len(writer.receivedSnaps)) + } + if writer.receivedSnaps[0].UserID != 1 { + t.Errorf("expected snap for UserID=1, got %d", writer.receivedSnaps[0].UserID) + } + if len(cache.readdCalled) != 0 { + t.Errorf("Readd should NOT be called on MISS, got %d calls", len(cache.readdCalled)) + } + if f.metrics.FlushSuccessTotal.Load() != 1 { + t.Errorf("FlushSuccessTotal = %d, want 1", f.metrics.FlushSuccessTotal.Load()) + } +} + +// --------------------------------------------------------------------------- +// 场景 3: UpsertFailReadds — writer 返普通 error → keys 被 Readd，FlushErrorTotal=1，DirtyReaddTotal=len +// --------------------------------------------------------------------------- + +func TestFlusher_UpsertFailReadds(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 1, Platform: "anthropic"}, + {UserID: 2, Platform: "openai"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + makeEntry(4.0, 5.0, 6.0), + }, + } + writeErr := errors.New("db connection timeout") + writer := &mockQuotaSnapshotWriter{returnErr: writeErr} + f := newTestFlusher(cache, writer) + + f.flush() + + if f.metrics.FlushErrorTotal.Load() != 1 { + t.Errorf("FlushErrorTotal = %d, want 1", f.metrics.FlushErrorTotal.Load()) + } + if len(cache.readdCalled) == 0 { + t.Fatal("Readd should be called after write error") + } + totalReadd := 0 + for _, rk := range cache.readdCalled { + totalReadd += len(rk) + } + if totalReadd != len(keys) { + t.Errorf("DirtyReaddTotal (from Readd calls) = %d, want %d", totalReadd, len(keys)) + } + if f.metrics.DirtyReaddTotal.Load() != int64(len(keys)) { + t.Errorf("DirtyReaddTotal metric = %d, want %d", f.metrics.DirtyReaddTotal.Load(), len(keys)) + } + if f.metrics.FlushSuccessTotal.Load() != 0 { + t.Errorf("FlushSuccessTotal = %d, want 0", f.metrics.FlushSuccessTotal.Load()) + } +} + +// --------------------------------------------------------------------------- +// 场景 4: FKViolationDropsNoReadd — writer 返 ErrUserPlatformQuotaFKViolation → 不 Readd，FlushFKViolationTotal=1 +// --------------------------------------------------------------------------- + +func TestFlusher_FKViolationDropsNoReadd(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 999, Platform: "anthropic"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + }, + } + writer := &mockQuotaSnapshotWriter{returnErr: ErrUserPlatformQuotaFKViolation} + f := newTestFlusher(cache, writer) + + f.flush() + + if f.metrics.FlushFKViolationTotal.Load() != 1 { + t.Errorf("FlushFKViolationTotal = %d, want 1", f.metrics.FlushFKViolationTotal.Load()) + } + if f.metrics.FlushErrorTotal.Load() != 1 { + t.Errorf("FlushErrorTotal = %d, want 1", f.metrics.FlushErrorTotal.Load()) + } + if len(cache.readdCalled) != 0 { + t.Errorf("Readd should NOT be called for FK violation (drop), got %d calls", len(cache.readdCalled)) + } + if f.metrics.DirtyReaddTotal.Load() != 0 { + t.Errorf("DirtyReaddTotal = %d, want 0 (FK violation drops)", f.metrics.DirtyReaddTotal.Load()) + } +} + +// --------------------------------------------------------------------------- +// 场景 5: NilSafe — var f *UserPlatformQuotaUsageFlusher; f.flush(); f.Stop() 不 panic +// --------------------------------------------------------------------------- + +func TestFlusher_NilSafe(t *testing.T) { + var f *UserPlatformQuotaUsageFlusher + // 下面两行不应 panic + f.flush() + f.Stop() +} + +// --------------------------------------------------------------------------- +// 场景 6: StopPreventsFlush — stopped=true 后 tick() 不调 flush（writer 没收到 snaps） +// --------------------------------------------------------------------------- + +func TestFlusher_StopPreventsFlush(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 1, Platform: "anthropic"}, + } + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: []*UserPlatformQuotaCacheEntry{ + makeEntry(1.0, 2.0, 3.0), + }, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + // 标记为已停止 + f.stopped.Store(true) + + // tick 应该直接返回，不触发 flush + f.tick() + + if len(writer.receivedSnaps) != 0 { + t.Errorf("expected 0 snaps after stop, got %d", len(writer.receivedSnaps)) + } + if cache.popCallIdx != 0 { + t.Errorf("Pop should not be called after stop, popCallIdx = %d", cache.popCallIdx) + } +} + +// --------------------------------------------------------------------------- +// 场景 B13-1: ZeroPercentCompany — 0% 公司脏集恒空，flusher 空跑无 DB 写 +// +// 模拟几乎没有用户配置 quota limit 的公司：脏集始终为空（popSequence 为空切片）， +// Pop 每次返回空集。flush() 应早退，不写 DB、不计成功、不 Readd。 +// --------------------------------------------------------------------------- + +func TestScenario_ZeroPercentCompany(t *testing.T) { + cache := &mockQuotaDirtyCache{ + // popSequence 为空 → Pop 超出序列 → 始终返回 nil（空集） + popSequence: [][]UserPlatformQuotaKey{}, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + f.flush() + + if len(writer.receivedSnaps) != 0 { + t.Errorf("0%% company: expected 0 snaps, got %d", len(writer.receivedSnaps)) + } + if f.metrics.FlushBatchSizeTotal.Load() != 0 { + t.Errorf("0%% company: FlushBatchSizeTotal = %d, want 0", f.metrics.FlushBatchSizeTotal.Load()) + } + if f.metrics.FlushSuccessTotal.Load() != 0 { + t.Errorf("0%% company: FlushSuccessTotal = %d, want 0 (empty-set early return)", f.metrics.FlushSuccessTotal.Load()) + } + if f.metrics.FlushErrorTotal.Load() != 0 { + t.Errorf("0%% company: FlushErrorTotal = %d, want 0", f.metrics.FlushErrorTotal.Load()) + } + if len(cache.readdCalled) != 0 { + t.Errorf("0%% company: Readd should never be called, got %d calls", len(cache.readdCalled)) + } +} + +// --------------------------------------------------------------------------- +// P1: IntervalFallback — flush_interval_ms ≤0 时回退 2s；正常值保留 +// --------------------------------------------------------------------------- + +func TestNewUserPlatformQuotaUsageFlusher_IntervalFallback(t *testing.T) { + cases := []struct { + name string + inMs int + wantDu time.Duration + }{ + {"零值回退 2s", 0, 2 * time.Second}, + {"负数回退 2s", -100, 2 * time.Second}, + {"正常 2000ms 保留", 2000, 2 * time.Second}, + {"正常 500ms 保留", 500, 500 * time.Millisecond}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + cfg := &config.Config{} + cfg.Database.UserPlatformQuotaFlushIntervalMs = tc.inMs + f := NewUserPlatformQuotaUsageFlusher(cfg, nil, nil, nil) + if f.interval != tc.wantDu { + t.Fatalf("interval = %v, want %v", f.interval, tc.wantDu) + } + }) + } +} + +// --------------------------------------------------------------------------- +// P1: EnabledField — flusher_enabled 配置正确写入 f.enabled +// --------------------------------------------------------------------------- + +func TestNewUserPlatformQuotaUsageFlusher_EnabledField(t *testing.T) { + for _, enabled := range []bool{true, false} { + cfg := &config.Config{} + cfg.Database.UserPlatformQuotaFlusherEnabled = enabled + cfg.Database.UserPlatformQuotaFlushIntervalMs = 500 + f := NewUserPlatformQuotaUsageFlusher(cfg, nil, nil, nil) + if f.enabled != enabled { + t.Errorf("enabled = %v, want %v", f.enabled, enabled) + } + } +} + +// --------------------------------------------------------------------------- +// P2: ReaddFailCounts — BatchGet 失败 + Readd 失败 → DirtyLostTotal 增、DirtyReaddTotal 不变 +// BatchGet 失败 + Readd 成功 → DirtyReaddTotal 增、DirtyLostTotal 不变 +// --------------------------------------------------------------------------- + +func TestFlusher_ReaddFailCounts(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 10, Platform: "anthropic"}, + {UserID: 11, Platform: "openai"}, + } + + t.Run("Readd 失败计 DirtyLostTotal", func(t *testing.T) { + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getErr: errors.New("redis timeout"), // 触发 BatchGet 失败路径 + readdErr: errors.New("redis connection refused"), // Readd 也失败 + } + f := newTestFlusher(cache, &mockQuotaSnapshotWriter{}) + + f.flush() + + if f.metrics.DirtyLostTotal.Load() != int64(len(keys)) { + t.Errorf("DirtyLostTotal = %d, want %d", f.metrics.DirtyLostTotal.Load(), len(keys)) + } + if f.metrics.DirtyReaddTotal.Load() != 0 { + t.Errorf("DirtyReaddTotal = %d, want 0 (Readd 失败不应计入)", f.metrics.DirtyReaddTotal.Load()) + } + }) + + t.Run("Readd 成功计 DirtyReaddTotal", func(t *testing.T) { + cache := &mockQuotaDirtyCache{ + popSequence: [][]UserPlatformQuotaKey{keys}, + getErr: errors.New("redis timeout"), // 触发 BatchGet 失败路径 + readdErr: nil, // Readd 成功 + } + f := newTestFlusher(cache, &mockQuotaSnapshotWriter{}) + + f.flush() + + if f.metrics.DirtyReaddTotal.Load() != int64(len(keys)) { + t.Errorf("DirtyReaddTotal = %d, want %d", f.metrics.DirtyReaddTotal.Load(), len(keys)) + } + if f.metrics.DirtyLostTotal.Load() != 0 { + t.Errorf("DirtyLostTotal = %d, want 0 (Readd 成功不应计 lost)", f.metrics.DirtyLostTotal.Load()) + } + }) +} + +// --------------------------------------------------------------------------- +// ClampsBatchSize — NewUserPlatformQuotaUsageFlusher 构造时按 +// [defaultFlushBatchSize, maxFlushBatchSize] 区间 clamp batchSize +// --------------------------------------------------------------------------- + +func TestNewUserPlatformQuotaUsageFlusher_ClampsBatchSize(t *testing.T) { + cases := []struct { + name string + in int + want int + }{ + {"超上限被 clamp", 7000, maxFlushBatchSize}, + {"恰好上限保留", maxFlushBatchSize, maxFlushBatchSize}, + {"零回退默认", 0, defaultFlushBatchSize}, + {"负数回退默认", -5, defaultFlushBatchSize}, + {"正常值保留", 500, 500}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + cfg := &config.Config{} + cfg.Database.UserPlatformQuotaFlushBatchSize = tc.in + f := NewUserPlatformQuotaUsageFlusher(cfg, nil, nil, nil) + if f.batchSize != tc.want { + t.Fatalf("batchSize = %d, want %d", f.batchSize, tc.want) + } + }) + } +} + +// --------------------------------------------------------------------------- +// 场景 B13-2: NinetyPercentCompany — 90% 公司大量用户配 limit，一批 5 key 批量刷库 +// +// 模拟大量用户配置了 quota limit 的公司：脏集第一次 Pop 返回 5 个不同用户的 key， +// 之后返回空集（避免 flush 循环）。flush() 应构造 5 条 snapshot 写入 DB， +// 断言绝对值语义（snap 的 DailyUsageUSD 等于 entry 的值）、metrics 正确、不 Readd。 +// --------------------------------------------------------------------------- + +func TestScenario_NinetyPercentCompany(t *testing.T) { + keys := []UserPlatformQuotaKey{ + {UserID: 101, Platform: "anthropic"}, + {UserID: 102, Platform: "anthropic"}, + {UserID: 103, Platform: "openai"}, + {UserID: 104, Platform: "openai"}, + {UserID: 105, Platform: "anthropic"}, + } + entries := []*UserPlatformQuotaCacheEntry{ + makeEntry(1.1, 2.2, 3.3), + makeEntry(4.4, 5.5, 6.6), + makeEntry(7.7, 8.8, 9.9), + makeEntry(0.5, 1.0, 1.5), + makeEntry(10.0, 20.0, 30.0), + } + cache := &mockQuotaDirtyCache{ + // 第 1 次 Pop 返回 5 keys，之后返回空集（防止 flush 无限循环） + popSequence: [][]UserPlatformQuotaKey{keys}, + getEntries: entries, + } + writer := &mockQuotaSnapshotWriter{} + f := newTestFlusher(cache, writer) + + f.flush() + + // 应收到 5 条 snapshot + if len(writer.receivedSnaps) != 5 { + t.Fatalf("90%% company: expected 5 snaps, got %d", len(writer.receivedSnaps)) + } + + // 验证绝对值语义：第 1 条 snap 的各窗口 usage 应等于 entries[0] 的值 + snap0 := writer.receivedSnaps[0] + entry0 := entries[0] + if snap0.DailyUsageUSD != entry0.DailyUsageUSD { + t.Errorf("snap[0].DailyUsageUSD = %v, want %v", snap0.DailyUsageUSD, entry0.DailyUsageUSD) + } + if snap0.WeeklyUsageUSD != entry0.WeeklyUsageUSD { + t.Errorf("snap[0].WeeklyUsageUSD = %v, want %v", snap0.WeeklyUsageUSD, entry0.WeeklyUsageUSD) + } + if snap0.MonthlyUsageUSD != entry0.MonthlyUsageUSD { + t.Errorf("snap[0].MonthlyUsageUSD = %v, want %v", snap0.MonthlyUsageUSD, entry0.MonthlyUsageUSD) + } + + // FlushBatchSizeTotal 应为 5（本批 keys 数量） + if f.metrics.FlushBatchSizeTotal.Load() != 5 { + t.Errorf("90%% company: FlushBatchSizeTotal = %d, want 5", f.metrics.FlushBatchSizeTotal.Load()) + } + // FlushSuccessTotal 应为 1（1 个批次写成功） + if f.metrics.FlushSuccessTotal.Load() != 1 { + t.Errorf("90%% company: FlushSuccessTotal = %d, want 1", f.metrics.FlushSuccessTotal.Load()) + } + // 无错误、无 Readd + if f.metrics.FlushErrorTotal.Load() != 0 { + t.Errorf("90%% company: FlushErrorTotal = %d, want 0", f.metrics.FlushErrorTotal.Load()) + } + if f.metrics.DirtyReaddTotal.Load() != 0 { + t.Errorf("90%% company: DirtyReaddTotal = %d, want 0", f.metrics.DirtyReaddTotal.Load()) + } + if len(cache.readdCalled) != 0 { + t.Errorf("90%% company: Readd should not be called, got %d calls", len(cache.readdCalled)) + } +} diff --git a/backend/internal/service/user_platform_quota_port.go b/backend/internal/service/user_platform_quota_port.go index cb09542a421f..0f88eda4c3b7 100644 --- a/backend/internal/service/user_platform_quota_port.go +++ b/backend/internal/service/user_platform_quota_port.go @@ -11,6 +11,23 @@ import ( // handler 只需引用 service 包，无需直接依赖 repository 包。 var ErrUserPlatformQuotaNotFound = errors.New("user platform quota not found") +// ErrUserPlatformQuotaFKViolation service 层 sentinel：批量 snapshot UPSERT 时存在 +// user_id 不在 users 表的记录（外键违反）。adapter 负责将 repository 层同名 sentinel 包装为此错误。 +var ErrUserPlatformQuotaFKViolation = errors.New("user platform quota snapshot FK violation") + +// UserPlatformQuotaSnapshot 是 service 层 flusher 向 DB 写入快照时使用的传输结构。 +// 字段语义与 repository.UserPlatformQuotaSnapshot 完全对应，由 adapter 负责转换。 +type UserPlatformQuotaSnapshot struct { + UserID int64 + Platform string + DailyUsageUSD float64 + WeeklyUsageUSD float64 + MonthlyUsageUSD float64 + DailyWindowStart time.Time + WeeklyWindowStart time.Time + MonthlyWindowStart time.Time +} + // UserPlatformQuotaRecord service 层传输结构体（与 repository 层解耦）。 type UserPlatformQuotaRecord struct { UserID int64 @@ -47,4 +64,6 @@ type UserPlatformQuotaRepository interface { // ResetExpiredWindow 重置指定窗口（"daily"|"weekly"|"monthly"）的用量与起始时间。 // 未命中活跃记录时返回（service-side wrapper of repository.ErrUserPlatformQuotaNotFound）。 ResetExpiredWindow(ctx context.Context, userID int64, platform string, window string, newStart time.Time) error + // BatchSnapshotUsage 绝对值覆盖写入整批 usage 快照。FK 违反返回 ErrUserPlatformQuotaFKViolation。 + BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error } diff --git a/backend/internal/service/user_service.go b/backend/internal/service/user_service.go index 36bcf1c881d5..f801e2c4cb63 100644 --- a/backend/internal/service/user_service.go +++ b/backend/internal/service/user_service.go @@ -74,11 +74,16 @@ type UserListFilters struct { // For large datasets this can be expensive; admin list pages should enable it on demand. // nil means not specified (default: load subscriptions for backward compatibility). IncludeSubscriptions *bool + // IncludeDeleted 为 true 时绕过软删除过滤，返回含已删除（deleted_at 非空）的用户。 + // 仅供 /admin/usage 的 SearchUsers 端点使用，其他列表调用方不要设置。 + IncludeDeleted bool } type UserRepository interface { Create(ctx context.Context, user *User) error GetByID(ctx context.Context, id int64) (*User, error) + // GetByIDIncludeDeleted 绕过软删除过滤按 ID 取用户（含已删）。仅供管理员审计/usage 点击使用。 + GetByIDIncludeDeleted(ctx context.Context, id int64) (*User, error) GetByEmail(ctx context.Context, email string) (*User, error) GetFirstAdmin(ctx context.Context) (*User, error) Update(ctx context.Context, user *User) error diff --git a/backend/internal/service/user_service_test.go b/backend/internal/service/user_service_test.go index 19aec5d312be..417140adcc84 100644 --- a/backend/internal/service/user_service_test.go +++ b/backend/internal/service/user_service_test.go @@ -236,6 +236,10 @@ func (m *mockUserRepo) UnbindUserAuthProvider(_ context.Context, _ int64, provid return nil } +func (m *mockUserRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*User, error) { + return m.GetByID(ctx, id) +} + func (m *mockUserRepo) WithUserProfileIdentityTx(ctx context.Context, fn func(txCtx context.Context) error) error { m.txCalls++ txState := &mockUserRepoTxState{ @@ -327,10 +331,22 @@ func (m *mockBillingCache) DeleteUserPlatformQuotaCache(context.Context, int64, return nil } -func (m *mockBillingCache) IncrUserPlatformQuotaUsageCache(context.Context, int64, string, float64, time.Duration) error { +func (m *mockBillingCache) IncrUserPlatformQuotaUsageCache(context.Context, int64, string, float64, time.Duration, bool) error { return nil } +func (m *mockBillingCache) PopDirtyUserPlatformQuotaKeys(context.Context, int) ([]UserPlatformQuotaKey, error) { + return nil, nil +} + +func (m *mockBillingCache) ReaddDirtyUserPlatformQuotaKeys(context.Context, []UserPlatformQuotaKey) error { + return nil +} + +func (m *mockBillingCache) BatchGetUserPlatformQuotaCache(context.Context, []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) { + return nil, nil +} + // --- 测试 --- func TestUpdateBalance_Success(t *testing.T) { diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index b22e10ae382b..fbee1b056caa 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -45,6 +45,17 @@ func ProvideOAuthRefreshAPI(accountRepo AccountRepository, tokenCache GeminiToke return NewOAuthRefreshAPI(accountRepo, tokenCache) } +// ProvideOpenAIOAuthService creates OpenAIOAuthService with privacy/account enrichment support. +func ProvideOpenAIOAuthService( + proxyRepo ProxyRepository, + oauthClient OpenAIOAuthClient, + privacyClientFactory PrivacyClientFactory, +) *OpenAIOAuthService { + svc := NewOpenAIOAuthService(proxyRepo, oauthClient) + svc.SetPrivacyClientFactory(privacyClientFactory) + return svc +} + // ProvideTokenRefreshService creates and starts TokenRefreshService func ProvideTokenRefreshService( accountRepo AccountRepository, @@ -396,6 +407,46 @@ func ProvideBackupService( return svc } +// ProvideOpsService constructs OpsService and wires the SettingService-backed quota +// auto-pause cache sink. Mirrors the SetCleanupReloader pattern: OpsService doesn't +// hold a *SettingService reference, but wire injects a tiny callback so writes to +// ops_advanced_settings immediately propagate into the scheduler hot-path cache. +func ProvideOpsService( + opsRepo OpsRepository, + settingRepo SettingRepository, + cfg *config.Config, + accountRepo AccountRepository, + userRepo UserRepository, + concurrencyService *ConcurrencyService, + gatewayService *GatewayService, + openAIGatewayService *OpenAIGatewayService, + geminiCompatService *GeminiMessagesCompatService, + antigravityGatewayService *AntigravityGatewayService, + systemLogSink *OpsSystemLogSink, + settingService *SettingService, +) *OpsService { + svc := NewOpsService( + opsRepo, + settingRepo, + cfg, + accountRepo, + userRepo, + concurrencyService, + gatewayService, + openAIGatewayService, + geminiCompatService, + antigravityGatewayService, + systemLogSink, + ) + if settingService != nil { + svc.SetOpenAIQuotaAutoPauseSettingsSink(settingService.SetOpenAIQuotaAutoPauseSettings) + // Optional warm-up so the first scheduled request after process start observes + // a populated cache rather than zero defaults. Best-effort, sync-bounded. + settingService.WarmOpenAIQuotaAutoPauseSettings(context.Background()) + } + return svc +} + // ProvideSettingService wires SettingService with group reader and proxy repo. func ProvideSettingService(settingRepo SettingRepository, groupRepo GroupRepository, proxyRepo ProxyRepository, cfg *config.Config) *SettingService { svc := NewSettingService(settingRepo, cfg) @@ -461,7 +512,7 @@ var ProviderSet = wire.NewSet( NewOpenAIGatewayService, wire.Bind(new(AccountRuntimeBlocker), new(*OpenAIGatewayService)), NewOAuthService, - NewOpenAIOAuthService, + ProvideOpenAIOAuthService, NewGeminiOAuthService, NewGeminiQuotaService, NewCompositeTokenCacheInvalidator, @@ -481,7 +532,7 @@ var ProviderSet = wire.NewSet( NewDataManagementService, ProvideBackupService, ProvideOpsSystemLogSink, - NewOpsService, + ProvideOpsService, ProvideOpsMetricsCollector, ProvideOpsAggregationService, ProvideOpsAlertEvaluatorService, @@ -531,8 +582,16 @@ var ProviderSet = wire.NewSet( ProvideChannelMonitorService, ProvideChannelMonitorRunner, NewChannelMonitorRequestTemplateService, + ProvideUserPlatformQuotaUsageFlusher, ) +// ProvideUserPlatformQuotaUsageFlusher 创建并启动 UserPlatformQuotaUsageFlusher。 +func ProvideUserPlatformQuotaUsageFlusher(cfg *config.Config, cache BillingCache, quotaRepo UserPlatformQuotaRepository, tw *TimingWheelService) *UserPlatformQuotaUsageFlusher { + svc := NewUserPlatformQuotaUsageFlusher(cfg, cache, quotaRepo, tw) + svc.Start() + return svc +} + // ProvidePaymentConfigService wraps NewPaymentConfigService to accept the named // payment.EncryptionKey type instead of raw []byte, avoiding Wire ambiguity. func ProvidePaymentConfigService(entClient *dbent.Client, settingRepo SettingRepository, key payment.EncryptionKey) *PaymentConfigService { diff --git a/backend/resources/model-pricing/model_prices_and_context_window.json b/backend/resources/model-pricing/model_prices_and_context_window.json index 3cae8c8b58dc..e88ed2da2215 100644 --- a/backend/resources/model-pricing/model_prices_and_context_window.json +++ b/backend/resources/model-pricing/model_prices_and_context_window.json @@ -1,135 +1,4 @@ { - "claude-3-5-haiku-20241022": { - "cache_creation_input_token_cost": 1e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 8e-08, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 8e-07, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 4e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 264 - }, - "claude-3-5-haiku-latest": { - "cache_creation_input_token_cost": 1.25e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 1e-07, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 1e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 5e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 264 - }, - "claude-3-5-sonnet-20240620": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 - }, - "claude-3-5-sonnet-20241022": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 - }, - "claude-3-5-sonnet-latest": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 - }, "claude-3-7-sonnet-20250219": { "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -159,34 +28,6 @@ "supports_web_search": true, "tool_use_system_prompt_tokens": 159 }, - "claude-3-7-sonnet-latest": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 - }, "claude-3-haiku-20240307": { "cache_creation_input_token_cost": 3e-07, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -226,28 +67,9 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 395 }, - "claude-3-opus-latest": { - "cache_creation_input_token_cost": 1.875e-05, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 1.5e-06, - "deprecation_date": "2025-03-01", - "input_cost_per_token": 1.5e-05, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 7.5e-05, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 395 - }, "claude-4-opus-20250514": { "cache_creation_input_token_cost": 1.875e-05, + "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "anthropic", @@ -274,6 +96,7 @@ }, "claude-4-sonnet-20250514": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, @@ -447,6 +270,7 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -474,6 +298,7 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -485,18 +310,14 @@ "claude-opus-4-6": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, - "cache_creation_input_token_cost_above_200k_tokens": 1.25e-05, "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 5e-06, - "input_cost_per_token_above_200k_tokens": 1e-05, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "output_cost_per_token_above_200k_tokens": 3.75e-05, "provider_specific_entry": { "fast": 6.0, "us": 1.1 @@ -506,9 +327,13 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -520,18 +345,14 @@ "claude-opus-4-6-20260205": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, - "cache_creation_input_token_cost_above_200k_tokens": 1.25e-05, "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 5e-06, - "input_cost_per_token_above_200k_tokens": 1e-05, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "output_cost_per_token_above_200k_tokens": 3.75e-05, "provider_specific_entry": { "fast": 6.0, "us": 1.1 @@ -541,9 +362,13 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -555,18 +380,14 @@ "claude-opus-4-6-thinking": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, - "cache_creation_input_token_cost_above_200k_tokens": 1.25e-05, "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 5e-06, - "input_cost_per_token_above_200k_tokens": 1e-05, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "output_cost_per_token_above_200k_tokens": 3.75e-05, "provider_specific_entry": { "fast": 6.0, "us": 1.1 @@ -576,15 +397,127 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346 + }, + "claude-opus-4-7": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "provider_specific_entry": { + "fast": 6.0, + "us": 1.1 + }, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "tool_use_system_prompt_tokens": 346 + }, + "claude-opus-4-7-20260416": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "provider_specific_entry": { + "fast": 6.0, + "us": 1.1 + }, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "tool_use_system_prompt_tokens": 346 + }, + "claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "provider_specific_entry": { + "fast": 6.0, + "us": 1.1 + }, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346 }, "claude-sonnet-4-20250514": { @@ -621,6 +554,7 @@ }, "claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, @@ -651,6 +585,7 @@ }, "claude-sonnet-4-5-20250929": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, @@ -682,6 +617,7 @@ }, "claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, @@ -707,26 +643,27 @@ }, "claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, - "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token": 3e-06, - "input_cost_per_token_above_200k_tokens": 6e-06, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_200k_tokens": 2.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": true, + "supports_output_config": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -735,6 +672,54 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346 }, + "codex-auto-review": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_priority": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_priority": 6e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "deepseek-chat": { "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.8e-07, @@ -792,493 +777,250 @@ "supports_reasoning": true, "supports_tool_choice": true }, - "gemini-1.0-pro": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-001": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-002": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-vision": { - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.0-pro-vision-001": { - "deprecation_date": "2025-04-09", - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.0-ultra": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "As of Jun, 2024. There is no available doc on vertex ai pricing gemini-1.0-ultra-001. Using gemini-1.0-pro pricing. Got max_tokens info here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-ultra-001": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "As of Jun, 2024. There is no available doc on vertex ai pricing gemini-1.0-ultra-001. Using gemini-1.0-pro pricing. Got max_tokens info here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.5-flash": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, + "gemini-2.0-flash": { + "cache_read_input_token_cost": 2.5e-08, + "deprecation_date": "2026-06-01", + "input_cost_per_audio_token": 7e-07, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, + "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_pdf_size_mb": 30, "max_tokens": 8192, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", + "output_cost_per_token": 4e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/pricing#2_0flash", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true }, - "gemini-1.5-flash-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, + "gemini-2.0-flash-001": { + "cache_read_input_token_cost": 3.75e-08, + "deprecation_date": "2026-06-01", + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, + "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_pdf_size_mb": 30, "max_tokens": 8192, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", + "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "gemini-2.0-flash-exp-image-generation": { + "input_cost_per_token": 0.0, + "litellm_provider": "gemini", + "max_images_per_prompt": 3000, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.039, + "output_cost_per_token": 0.0, + "source": "https://ai.google.dev/pricing", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], "supports_vision": true }, - "gemini-1.5-flash-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, + "gemini-2.0-flash-lite": { + "cache_read_input_token_cost": 1.875e-08, + "deprecation_date": "2026-06-01", + "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_pdf_size_mb": 50, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-1.5-flash", + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, - "gemini-1.5-flash-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 4.688e-09, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, + "gemini-2.0-flash-lite-001": { + "cache_read_input_token_cost": 1.875e-08, + "deprecation_date": "2026-06-01", + "input_cost_per_audio_token": 7.5e-08, + "input_cost_per_token": 7.5e-08, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, + "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_pdf_size_mb": 50, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 1.875e-08, - "output_cost_per_character_above_128k_tokens": 3.75e-08, - "output_cost_per_token": 4.6875e-09, - "output_cost_per_token_above_128k_tokens": 9.375e-09, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", + "output_cost_per_token": 3e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, - "gemini-1.5-flash-preview-0514": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, + "gemini-2.5-computer-use-preview-10-2025": { + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", - "output_cost_per_character": 1.875e-08, - "output_cost_per_character_above_128k_tokens": 3.75e-08, - "output_cost_per_token": 4.6875e-09, - "output_cost_per_token_above_128k_tokens": 9.375e-09, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-1.5-pro", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/computer-use", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true }, - "gemini-1.5-pro-preview-0215": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gemini-1.5-pro-preview-0409": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "gemini-1.5-pro-preview-0514": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gemini-2.0-flash": { - "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2026-03-31", - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, + "gemini-2.5-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 8192, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/pricing#2_0flash", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1286,39 +1028,48 @@ "video" ], "supported_output_modalities": [ - "text", - "image" + "text" ], - "supports_audio_input": true, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-001": { - "cache_read_input_token_cost": 3.75e-08, - "deprecation_date": "2026-03-31", + "gemini-2.5-flash-image": { + "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 1.5e-07, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, + "max_input_tokens": 32768, + "max_output_tokens": 32768, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 32768, "max_video_length": 1, "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "mode": "image_generation", + "output_cost_per_image": 0.039, + "output_cost_per_image_token": 3e-05, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "rpm": 100000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1329,44 +1080,48 @@ "text", "image" ], - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": false, + "tpm": 8000000 }, - "gemini-2.0-flash-exp": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 1.5e-07, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, + "gemini-2.5-flash-lite": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 8192, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 6e-07, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "output_cost_per_reasoning_token": 4e-07, + "output_cost_per_token": 4e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1374,57 +1129,51 @@ "video" ], "supported_output_modalities": [ - "text", - "image" + "text" ], - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-exp-image-generation": { - "input_cost_per_token": 0.0, - "litellm_provider": "gemini", - "max_images_per_prompt": 3000, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_token": 0.0, - "source": "https://ai.google.dev/pricing", - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_vision": true - }, - "gemini-2.0-flash-lite": { - "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", - "input_cost_per_audio_token": 7.5e-08, - "input_cost_per_token": 7.5e-08, + "gemini-2.5-flash-lite-preview-06-17": { + "cache_read_input_token_cost": 2.5e-08, + "deprecation_date": "2025-11-18", + "input_cost_per_audio_token": 5e-07, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 50, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 3e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "output_cost_per_reasoning_token": 4e-07, + "output_cost_per_token": 4e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1434,34 +1183,48 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-lite-001": { - "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", - "input_cost_per_audio_token": 7.5e-08, - "input_cost_per_token": 7.5e-08, + "gemini-2.5-flash-lite-preview-09-2025": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 50, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 3e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ + "output_cost_per_reasoning_token": 4e-07, + "output_cost_per_token": 4e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://developers.googleblog.com/en/continuing-to-bring-you-our-latest-models-with-an-improved-gemini-2-5-flash-and-flash-lite-release/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ "text", "image", "audio", @@ -1470,181 +1233,119 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-live-preview-04-09": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 3e-06, - "input_cost_per_image": 3e-06, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 3e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, + "gemini-2.5-flash-native-audio-latest": { + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, + "max_output_tokens": 8192, + "max_tokens": 8192, "mode": "chat", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 2e-06, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/gemini#gemini-2-0-flash-live-preview-04-09", + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/pricing", "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" + "/v1/realtime" ], "supported_modalities": [ "text", - "image", - "audio", - "video" + "audio" ], "supported_output_modalities": [ "text", "audio" ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 + "supports_audio_input": true, + "supports_audio_output": true }, - "gemini-2.0-flash-preview-image-generation": { - "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2025-11-14", - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, + "gemini-2.5-flash-native-audio-preview-09-2025": { + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_pdf_size_mb": 30, "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/pricing#2_0flash", + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], "supported_modalities": [ "text", - "image", - "audio", - "video" + "audio" ], "supported_output_modalities": [ "text", - "image" + "audio" ], "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true + "supports_audio_output": true }, - "gemini-2.0-flash-thinking-exp": { - "cache_read_input_token_cost": 0.0, - "deprecation_date": "2025-12-02", - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, + "gemini-2.5-flash-native-audio-preview-12-2025": { + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_pdf_size_mb": 30, "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], "supported_modalities": [ "text", - "image", - "audio", - "video" + "audio" ], "supported_output_modalities": [ "text", - "image" + "audio" ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true + "supports_audio_input": true, + "supports_audio_output": true }, - "gemini-2.0-flash-thinking-exp-01-21": { - "cache_read_input_token_cost": 0.0, - "deprecation_date": "2025-12-02", - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, + "gemini-2.5-flash-preview-09-2025": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65536, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 65536, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://developers.googleblog.com/en/continuing-to-bring-you-our-latest-models-with-an-improved-gemini-2-5-flash-and-flash-lite-release/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], "supported_modalities": [ "text", "image", @@ -1652,37 +1353,55 @@ "video" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_audio_output": false, - "supports_function_calling": false, + "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": false, + "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-pro-exp-02-05": { - "cache_read_input_token_cost": 3.125e-07, + "gemini-2.5-flash-preview-tts": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", + "mode": "audio_speech", + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/pricing", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "gemini-2.5-pro": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 2097152, - "max_output_tokens": 8192, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -1699,45 +1418,65 @@ ], "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true }, - "gemini-2.5-computer-use-preview-10-2025": { + "gemini-2.5-pro-preview-tts": { + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/computer-use", + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", "supported_modalities": [ - "text", - "image" + "text" ], "supported_output_modalities": [ - "text" + "audio" ], - "supports_computer_use": true, + "supports_audio_output": false, "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, - "gemini-2.5-flash": { - "cache_read_input_token_cost": 3e-08, + "gemini-3-flash": { + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, @@ -1749,9 +1488,15 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_reasoning_token": 3e-06, + "output_cost_per_token": 3e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/pricing/gemini-3", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1768,38 +1513,47 @@ ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-image": { - "cache_read_input_token_cost": 3e-08, + "gemini-3-flash-preview": { + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, - "max_input_tokens": 32768, - "max_output_tokens": 32768, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, "max_pdf_size_mb": 30, - "max_tokens": 32768, + "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, - "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_image_token": 3e-05, - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image", + "mode": "chat", + "output_cost_per_reasoning_token": 3e-06, + "output_cost_per_token": 3e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/pricing/gemini-3", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1812,45 +1566,43 @@ "video" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": false, - "tpm": 8000000 + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-image-preview": { - "cache_read_input_token_cost": 7.5e-08, - "deprecation_date": "2026-01-15", - "input_cost_per_audio_token": 1e-06, - "input_cost_per_image_token": 3e-07, - "input_cost_per_token": 3e-07, + "gemini-3-pro-image-preview": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_image_token": 3e-05, - "output_cost_per_reasoning_token": 3e-05, - "output_cost_per_token": 3e-05, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1858,31 +1610,33 @@ ], "supported_modalities": [ "text", - "image", - "audio", - "video" + "image" ], "supported_output_modalities": [ "text", "image" ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, + "supports_function_calling": false, "supports_prompt_caching": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "tpm": 8000000 + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, + "gemini-3-pro-preview": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "deprecation_date": "2026-03-26", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, @@ -1894,9 +1648,17 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1911,38 +1673,130 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": false, + "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-lite-preview-06-17": { + "gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-3.1-flash-image-preview": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-3.1-flash-lite": { "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2025-11-18", + "cache_read_input_token_cost_batches": 1.25e-08, + "cache_read_input_token_cost_flex": 1.25e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_audio_token": 5e-07, - "input_cost_per_token": 1e-07, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token_flex": 1.25e-07, + "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_priority": 2.7e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -1957,37 +1811,50 @@ "supported_output_modalities": [ "text" ], + "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, + "gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, + "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "source": "https://developers.googleblog.com/en/continuing-to-bring-you-our-latest-models-with-an-improved-gemini-2-5-flash-and-flash-lite-release/", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/gemini-api/docs/models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2002,109 +1869,92 @@ "supported_output_modalities": [ "text" ], + "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-native-audio-latest": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "gemini-3.1-flash-live-preview": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], "supported_modalities": [ "text", - "audio" + "image", + "audio", + "video" ], "supported_output_modalities": [ "text", "audio" ], "supports_audio_input": true, - "supports_audio_output": true + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true }, - "gemini-2.5-flash-native-audio-preview-09-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", - "supported_endpoints": [ - "/v1/realtime" - ], - "supported_modalities": [ - "text", - "audio" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true - }, - "gemini-2.5-flash-native-audio-preview-12-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", - "supported_endpoints": [ - "/v1/realtime" - ], - "supported_modalities": [ - "text", - "audio" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true - }, - "gemini-2.5-flash-preview-04-17": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 1.5e-07, + "gemini-3.1-pro-high": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 3.5e-06, - "output_cost_per_token": 6e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_image": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2119,37 +1969,56 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": false, + "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-preview-05-20": { - "cache_read_input_token_cost": 7.5e-08, - "deprecation_date": "2025-11-18", - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "gemini-3.1-pro-low": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_image": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2164,37 +2033,56 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": false, + "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-flash-preview-09-2025": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "gemini-3.1-pro-preview": { + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://developers.googleblog.com/en/continuing-to-bring-you-our-latest-models-with-an-improved-gemini-2-5-flash-and-flash-lite-release/", + "output_cost_per_image": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2209,52 +2097,54 @@ "supported_output_modalities": [ "text" ], - "supports_audio_output": false, + "supports_audio_input": true, "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, + "supports_video_input": true, "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-flash-preview-tts": { - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "mode": "audio_speech", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", - "supported_endpoints": [ - "/v1/audio/speech" - ] + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-pro": { + "gemini-3.1-pro-preview-customtools": { "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "output_cost_per_image": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -2267,21 +2157,26 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-pro-exp-03-25": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_priority": 2.7e-07, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, @@ -2293,12 +2188,19 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_token_priority": 1.62e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/pricing/gemini-3", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -2310,25 +2212,67 @@ "text" ], "supports_audio_input": true, + "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "web_search_billing_unit": "per_query" }, - "gemini-2.5-pro-preview-03-25": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "deprecation_date": "2025-12-02", - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", + "gemini-embedding-001": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "vertex_ai-embedding-models", + "max_input_tokens": 2048, + "max_tokens": 2048, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" + }, + "gemini-embedding-2": { + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.00079, + "litellm_provider": "vertex_ai-embedding-models", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_multimodal": true, + "uses_embed_content": true + }, + "gemini-embedding-2-preview": { + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.00079, + "litellm_provider": "vertex_ai-embedding-models", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "uses_embed_content": true + }, + "gemini-exp-1206": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, @@ -2339,10 +2283,16 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "rpm": 100000, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", "/v1/batch" @@ -2365,17 +2315,29 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "tpm": 8000000 }, - "gemini-2.5-pro-preview-05-06": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "deprecation_date": "2025-12-02", - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "gemini-flash-experimental": { + "input_cost_per_character": 0, + "input_cost_per_token": 0, "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "uses_embed_content": true + }, + "gemini-flash-latest": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, @@ -2386,8 +2348,14 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "rpm": 100000, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", "supported_endpoints": [ "/v1/chat/completions", @@ -2403,9 +2371,6 @@ "supported_output_modalities": [ "text" ], - "supported_regions": [ - "global" - ], "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -2415,16 +2380,16 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "tpm": 8000000 }, - "gemini-2.5-pro-preview-06-05": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", + "gemini-flash-lite-latest": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, + "input_cost_per_token": 1e-07, + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, @@ -2435,9 +2400,15 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "output_cost_per_reasoning_token": 4e-07, + "output_cost_per_token": 4e-07, + "rpm": 15, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-lite", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -2461,52 +2432,15 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-pro-preview-tts": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "audio" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "tpm": 250000 }, - "gemini-3-flash-preview": { - "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_priority": 9e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_audio_token_priority": 1.8e-06, - "input_cost_per_token": 5e-07, - "input_cost_per_token_priority": 9e-07, + "gemini-live-2.5-flash-preview-native-audio-09-2025": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, @@ -2517,15 +2451,17 @@ "max_tokens": 65535, "max_video_length": 1, "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 3e-06, - "output_cost_per_token": 3e-06, - "output_cost_per_token_priority": 5.4e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/vertex_ai/live" ], "supported_modalities": [ "text", @@ -2534,69 +2470,28 @@ "video" ], "supported_output_modalities": [ - "text" + "text", + "audio" ], - "supports_audio_output": false, + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, - "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, "supports_web_search": true }, - "gemini-3-pro-image-preview": { - "input_cost_per_image": 0.0011, - "input_cost_per_token": 2e-06, - "input_cost_per_token_batches": 1e-06, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 65536, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "image_generation", - "output_cost_per_image": 0.134, - "output_cost_per_image_token": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_function_calling": false, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-3-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_priority": 3.6e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "input_cost_per_token_batches": 1e-06, - "input_cost_per_token_priority": 3.6e-06, - "litellm_provider": "vertex_ai-language-models", + "gemini-pro-latest": { + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, @@ -2607,16 +2502,18 @@ "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "output_cost_per_token_batches": 6e-06, - "output_cost_per_token_priority": 2.16e-05, + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "rpm": 2000, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035 + }, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/v1/completions" ], "supported_modalities": [ "text", @@ -2629,576 +2526,43 @@ ], "supports_audio_input": true, "supports_function_calling": true, - "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "tpm": 800000 }, - "gemini-3.1-pro-high": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_priority": 3.6e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "input_cost_per_token_batches": 1e-06, - "input_cost_per_token_priority": 3.6e-06, + "gemini-robotics-er-1.5-preview": { + "cache_read_input_token_cost": 0, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, + "max_output_tokens": 65535, + "max_tokens": 65535, "mode": "chat", - "output_cost_per_image": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "output_cost_per_token_batches": 6e-06, - "output_cost_per_token_priority": 2.16e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-robotics-er-1-5-preview", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions", - "/v1/batch" + "/v1/completions" ], "supported_modalities": [ "text", "image", - "audio", - "video" + "video", + "audio" ], "supported_output_modalities": [ "text" ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-3.1-pro-low": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_priority": 3.6e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "input_cost_per_token_batches": 1e-06, - "input_cost_per_token_priority": 3.6e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_image": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "output_cost_per_token_batches": 6e-06, - "output_cost_per_token_priority": 2.16e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_priority": 3.6e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "input_cost_per_token_batches": 1e-06, - "input_cost_per_token_priority": 3.6e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_image": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "output_cost_per_token_batches": 6e-06, - "output_cost_per_token_priority": 2.16e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-3.1-pro-preview-customtools": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_read_input_token_cost": 2e-07, - "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "input_cost_per_token": 2e-06, - "input_cost_per_token_above_200k_tokens": 4e-06, - "input_cost_per_token_batches": 1e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_image": 0.00012, - "output_cost_per_token": 1.2e-05, - "output_cost_per_token_above_200k_tokens": 1.8e-05, - "output_cost_per_token_batches": 6e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-embedding-001": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 2048, - "max_tokens": 2048, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 3072, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" - }, - "gemini-exp-1206": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 8000000 - }, - "gemini-flash-experimental": { - "input_cost_per_character": 0, - "input_cost_per_token": 0, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_token": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/gemini-experimental", - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-flash-latest": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 8000000 - }, - "gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-lite", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini-live-2.5-flash-preview-native-audio-09-2025": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 2e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-pro": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-pro-experimental": { - "input_cost_per_character": 0, - "input_cost_per_token": 0, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_token": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/gemini-experimental", - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 800000 - }, - "gemini-pro-vision": { - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-robotics-er-1.5-preview": { - "cache_read_input_token_cost": 0, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-robotics-er-1-5-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "video", - "audio" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, + "supports_audio_output": false, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": false, @@ -3236,31 +2600,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-3.5-turbo-0301": { - "input_cost_per_token": 1.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-3.5-turbo-0613": { - "input_cost_per_token": 1.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-3.5-turbo-1106": { "deprecation_date": "2026-09-28", "input_cost_per_token": 1e-06, @@ -3288,18 +2627,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-3.5-turbo-16k-0613": { - "input_cost_per_token": 3e-06, - "litellm_provider": "openai", - "max_input_tokens": 16385, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 4e-06, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-3.5-turbo-instruct": { "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", @@ -3347,6 +2674,7 @@ "supports_tool_choice": true }, "gpt-4-0314": { + "deprecation_date": "2026-03-26", "input_cost_per_token": 3e-05, "litellm_provider": "openai", "max_input_tokens": 8192, @@ -3354,7 +2682,6 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-05, - "supports_prompt_caching": true, "supports_system_messages": true, "supports_tool_choice": true }, @@ -3387,57 +2714,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4-1106-vision-preview": { - "deprecation_date": "2024-12-06", - "input_cost_per_token": 1e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 3e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4-32k": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-4-32k-0314": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-4-32k-0613": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4-turbo": { "input_cost_per_token": 1e-05, "litellm_provider": "openai", @@ -3454,23 +2730,7 @@ "supports_tool_choice": true, "supports_vision": true }, - "gpt-4-turbo-2024-04-09": { - "input_cost_per_token": 1e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 3e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4-turbo-preview": { + "gpt-4-turbo-2024-04-09": { "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -3483,10 +2743,10 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, - "gpt-4-vision-preview": { - "deprecation_date": "2024-12-06", + "gpt-4-turbo-preview": { "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -3494,11 +2754,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 3e-05, + "supports_function_calling": true, + "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true + "supports_tool_choice": true }, "gpt-4.1": { "cache_read_input_token_cost": 5e-07, @@ -3535,7 +2796,8 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-2025-04-14": { "cache_read_input_token_cost": 5e-07, @@ -3569,7 +2831,8 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-mini": { "cache_read_input_token_cost": 1e-07, @@ -3606,7 +2869,8 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-mini-2025-04-14": { "cache_read_input_token_cost": 1e-07, @@ -3640,7 +2904,8 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-nano": { "cache_read_input_token_cost": 2.5e-08, @@ -3713,47 +2978,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "gpt-4.5-preview": { - "cache_read_input_token_cost": 3.75e-05, - "input_cost_per_token": 7.5e-05, - "input_cost_per_token_batches": 3.75e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 0.00015, - "output_cost_per_token_batches": 7.5e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4.5-preview-2025-02-27": { - "cache_read_input_token_cost": 3.75e-05, - "deprecation_date": "2025-07-14", - "input_cost_per_token": 7.5e-05, - "input_cost_per_token_batches": 3.75e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 0.00015, - "output_cost_per_token_batches": 7.5e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, "gpt-4o": { "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_priority": 2.125e-06, @@ -3857,23 +3081,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4o-audio-preview-2024-10-01": { - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 2.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 1e-05, - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4o-audio-preview-2024-12-17": { "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, @@ -4056,17 +3263,207 @@ "supports_vision": true, "supports_web_search": true }, - "gpt-4o-mini-search-preview-2025-03-11": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_token": 1.5e-07, - "input_cost_per_token_batches": 7.5e-08, + "gpt-4o-mini-search-preview-2025-03-11": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "gpt-4o-mini-transcribe": { + "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 16000, + "max_output_tokens": 2000, + "mode": "audio_transcription", + "output_cost_per_token": 5e-06, + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "gpt-4o-mini-transcribe-2025-03-20": { + "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 16000, + "max_output_tokens": 2000, + "mode": "audio_transcription", + "output_cost_per_token": 5e-06, + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "gpt-4o-mini-transcribe-2025-12-15": { + "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 16000, + "max_output_tokens": 2000, + "mode": "audio_transcription", + "output_cost_per_token": 5e-06, + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "gpt-4o-mini-tts": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "mode": "audio_speech", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_second": 0.00025, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "audio" + ] + }, + "gpt-4o-mini-tts-2025-03-20": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "mode": "audio_speech", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_second": 0.00025, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "audio" + ] + }, + "gpt-4o-mini-tts-2025-12-15": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "mode": "audio_speech", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_second": 0.00025, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "audio" + ] + }, + "gpt-4o-realtime-preview": { + "cache_read_input_token_cost": 2.5e-06, + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 2e-05, + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "gpt-4o-realtime-preview-2024-12-17": { + "cache_read_input_token_cost": 2.5e-06, + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 2e-05, + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "gpt-4o-realtime-preview-2025-06-03": { + "cache_read_input_token_cost": 2.5e-06, + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 2e-05, + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "gpt-4o-search-preview": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.05, + "search_context_size_low": 0.03, + "search_context_size_medium": 0.035 + }, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "gpt-4o-search-preview-2025-03-11": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 6e-07, - "output_cost_per_token_batches": 3e-07, + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -4076,186 +3473,469 @@ "supports_tool_choice": true, "supports_vision": true }, - "gpt-4o-mini-transcribe": { - "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 1.25e-06, + "gpt-4o-transcribe": { + "input_cost_per_audio_token": 2.5e-06, + "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "mode": "audio_transcription", - "output_cost_per_token": 5e-06, + "output_cost_per_token": 1e-05, "supported_endpoints": [ "/v1/audio/transcriptions" ] }, - "gpt-4o-mini-transcribe-2025-03-20": { - "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 1.25e-06, + "gpt-4o-transcribe-diarize": { + "input_cost_per_audio_token": 2.5e-06, + "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 16000, "max_output_tokens": 2000, "mode": "audio_transcription", - "output_cost_per_token": 5e-06, + "output_cost_per_token": 1e-05, "supported_endpoints": [ "/v1/audio/transcriptions" ] }, - "gpt-4o-mini-transcribe-2025-12-15": { - "input_cost_per_audio_token": 3e-06, + "gpt-5": { + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_flex": 6.25e-07, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", - "max_input_tokens": 16000, - "max_output_tokens": 2000, - "mode": "audio_transcription", - "output_cost_per_token": 5e-06, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_flex": 5e-06, + "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ - "/v1/audio/transcriptions" - ] + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-mini-tts": { - "input_cost_per_token": 2.5e-06, + "gpt-5-2025-08-07": { + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_priority": 2.5e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_flex": 6.25e-07, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", - "mode": "audio_speech", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_second": 0.00025, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_flex": 5e-06, + "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ - "/v1/audio/speech" + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" ], "supported_modalities": [ "text", - "audio" + "image" ], "supported_output_modalities": [ - "audio" - ] + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-mini-tts-2025-03-20": { - "input_cost_per_token": 2.5e-06, + "gpt-5-chat": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", - "mode": "audio_speech", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_second": 0.00025, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": false, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_xhigh_reasoning_effort": false + }, + "gpt-5-chat-latest": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": false, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_xhigh_reasoning_effort": false + }, + "gpt-5-codex": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", "output_cost_per_token": 1e-05, "supported_endpoints": [ - "/v1/audio/speech" + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false + }, + "gpt-5-mini": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_flex": 1.25e-08, + "cache_read_input_token_cost_priority": 4.5e-08, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_flex": 1.25e-07, + "input_cost_per_token_priority": 4.5e-07, + "litellm_provider": "openai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "output_cost_per_token_flex": 1e-06, + "output_cost_per_token_priority": 3.6e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false + }, + "gpt-5-mini-2025-08-07": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_flex": 1.25e-08, + "cache_read_input_token_cost_priority": 4.5e-08, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_flex": 1.25e-07, + "input_cost_per_token_priority": 4.5e-07, + "litellm_provider": "openai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "output_cost_per_token_flex": 1e-06, + "output_cost_per_token_priority": 3.6e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" ], "supported_modalities": [ "text", - "audio" + "image" ], "supported_output_modalities": [ - "audio" - ] + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_service_tier": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-mini-tts-2025-12-15": { - "input_cost_per_token": 2.5e-06, + "gpt-5-nano": { + "cache_read_input_token_cost": 5e-09, + "cache_read_input_token_cost_flex": 2.5e-09, + "input_cost_per_token": 5e-08, + "input_cost_per_token_flex": 2.5e-08, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", - "mode": "audio_speech", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_second": 0.00025, - "output_cost_per_token": 1e-05, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-07, + "output_cost_per_token_flex": 2e-07, "supported_endpoints": [ - "/v1/audio/speech" + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" ], "supported_modalities": [ "text", - "audio" + "image" ], "supported_output_modalities": [ - "audio" - ] - }, - "gpt-4o-realtime-preview": { - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 5e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, + "text" + ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-realtime-preview-2024-10-01": { - "cache_creation_input_audio_token_cost": 2e-05, - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 0.0001, - "input_cost_per_token": 5e-06, + "gpt-5-nano-2025-08-07": { + "cache_read_input_token_cost": 5e-09, + "cache_read_input_token_cost_flex": 2.5e-09, + "input_cost_per_token": 5e-08, + "input_cost_per_token_flex": 2.5e-08, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", - "output_cost_per_audio_token": 0.0002, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, + "output_cost_per_token": 4e-07, + "output_cost_per_token_flex": 2e-07, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 5e-06, + "gpt-5-pro": { + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, + "max_output_tokens": 272000, + "max_tokens": 272000, + "mode": "responses", + "output_cost_per_token": 0.00012, + "output_cost_per_token_batches": 6e-05, + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": false, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-realtime-preview-2025-06-03": { - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 5e-06, + "gpt-5-pro-2025-10-06": { + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, + "max_output_tokens": 272000, + "max_tokens": 272000, + "mode": "responses", + "output_cost_per_token": 0.00012, + "output_cost_per_token_batches": 6e-05, + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": false, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-search-preview": { - "cache_read_input_token_cost": 1.25e-06, - "input_cost_per_token": 2.5e-06, - "input_cost_per_token_batches": 1.25e-06, + "gpt-5-search-api": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, - "output_cost_per_token_batches": 5e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 - }, "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4263,58 +3943,34 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-4o-search-preview-2025-03-11": { - "cache_read_input_token_cost": 1.25e-06, - "input_cost_per_token": 2.5e-06, - "input_cost_per_token_batches": 1.25e-06, + "gpt-5-search-api-2025-10-14": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, - "output_cost_per_token_batches": 5e-06, "supports_function_calling": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4o-transcribe": { - "input_cost_per_audio_token": 6e-06, - "input_cost_per_token": 2.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 16000, - "max_output_tokens": 2000, - "mode": "audio_transcription", - "output_cost_per_token": 1e-05, - "supported_endpoints": [ - "/v1/audio/transcriptions" - ] - }, - "gpt-4o-transcribe-diarize": { - "input_cost_per_audio_token": 6e-06, - "input_cost_per_token": 2.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 16000, - "max_output_tokens": 2000, - "mode": "audio_transcription", - "output_cost_per_token": 1e-05, - "supported_endpoints": [ - "/v1/audio/transcriptions" - ] + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5": { + "gpt-5.1": { "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, - "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -4322,11 +3978,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, - "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4334,10 +3988,13 @@ "image" ], "supported_output_modalities": [ - "text" + "text", + "image" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4346,14 +4003,14 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-2025-08-07": { + "gpt-5.1-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, - "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -4361,11 +4018,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, - "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4373,10 +4028,13 @@ "image" ], "supported_output_modalities": [ - "text" + "text", + "image" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4385,20 +4043,24 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-chat": { + "gpt-5.1-chat-latest": { "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4406,42 +4068,13 @@ "image" ], "supported_output_modalities": [ - "text" - ], - "supports_function_calling": false, - "supports_native_streaming": true, - "supports_parallel_function_calling": false, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": false, - "supports_vision": true - }, - "gpt-5-chat-latest": { - "cache_read_input_token_cost": 1.25e-07, - "input_cost_per_token": 1.25e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 1e-05, - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/batch", - "/v1/responses" - ], - "supported_modalities": [ "text", "image" ], - "supported_output_modalities": [ - "text" - ], "supports_function_calling": false, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4449,17 +4082,22 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": false, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-codex": { + "gpt-5.1-codex": { "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 2e-05, "supported_endpoints": [ "/v1/responses" ], @@ -4471,7 +4109,9 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4479,26 +4119,20 @@ "supports_response_schema": true, "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-mini": { - "cache_read_input_token_cost": 2.5e-08, - "cache_read_input_token_cost_flex": 1.25e-08, - "cache_read_input_token_cost_priority": 4.5e-08, - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_flex": 1.25e-07, - "input_cost_per_token_priority": 4.5e-07, + "gpt-5.1-codex-max": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 2e-06, - "output_cost_per_token_flex": 1e-06, - "output_cost_per_token_priority": 3.6e-06, + "mode": "responses", + "output_cost_per_token": 1e-05, "supported_endpoints": [ - "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4509,35 +4143,33 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, + "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5-mini-2025-08-07": { + "gpt-5.1-codex-mini": { "cache_read_input_token_cost": 2.5e-08, - "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_token": 2.5e-07, - "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", + "mode": "responses", "output_cost_per_token": 2e-06, - "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, "supported_endpoints": [ - "/v1/chat/completions", - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4548,30 +4180,32 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, - "supports_system_messages": true, + "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-nano": { - "cache_read_input_token_cost": 5e-09, - "cache_read_input_token_cost_flex": 2.5e-09, - "input_cost_per_token": 5e-08, - "input_cost_per_token_flex": 2.5e-08, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.2": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 4e-07, - "output_cost_per_token_flex": 2e-07, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -4582,31 +4216,37 @@ "image" ], "supported_output_modalities": [ - "text" + "text", + "image" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5-nano-2025-08-07": { - "cache_read_input_token_cost": 5e-09, - "cache_read_input_token_cost_flex": 2.5e-09, - "input_cost_per_token": 5e-08, - "input_cost_per_token_flex": 2.5e-08, + "gpt-5.2-2025-12-11": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 4e-07, - "output_cost_per_token_flex": 2e-07, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -4617,31 +4257,39 @@ "image" ], "supported_output_modalities": [ - "text" + "text", + "image" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5-pro": { - "input_cost_per_token": 1.5e-05, - "input_cost_per_token_batches": 7.5e-06, + "gpt-5.2-chat-latest": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 272000, - "max_tokens": 272000, - "mode": "responses", - "output_cost_per_token": 0.00012, - "output_cost_per_token_batches": 6e-05, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ - "/v1/batch", + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -4652,7 +4300,9 @@ "text" ], "supports_function_calling": true, - "supports_native_streaming": false, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4661,20 +4311,22 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5-pro-2025-10-06": { - "input_cost_per_token": 1.5e-05, - "input_cost_per_token_batches": 7.5e-06, + "gpt-5.2-codex": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 272000, - "max_tokens": 272000, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00012, - "output_cost_per_token_batches": 6e-05, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ - "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4685,69 +4337,30 @@ "text" ], "supports_function_calling": true, - "supports_native_streaming": false, + "supports_minimal_reasoning_effort": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gpt-5-search-api": { - "cache_read_input_token_cost": 1.25e-07, - "input_cost_per_token": 1.25e-06, - "litellm_provider": "openai", - "max_input_tokens": 272000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gpt-5-search-api-2025-10-14": { - "cache_read_input_token_cost": 1.25e-07, - "input_cost_per_token": 1.25e-06, - "litellm_provider": "openai", - "max_input_tokens": 272000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, + "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.1": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_priority": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.2-pro": { + "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_priority": 2e-05, + "mode": "responses", + "output_cost_per_token": 0.000168, "supported_endpoints": [ - "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4755,35 +4368,33 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.1-2025-11-13": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_priority": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.2-pro-2025-12-11": { + "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_priority": 2e-05, + "mode": "responses", + "output_cost_per_token": 0.000168, "supported_endpoints": [ - "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4791,33 +4402,35 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.1-chat-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_priority": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.3-chat-latest": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_priority": 2e-05, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -4827,32 +4440,35 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], - "supports_function_calling": false, + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, - "supports_parallel_function_calling": false, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": false, - "supports_vision": true + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5.1-codex": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_priority": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_priority": 2.5e-06, + "gpt-5.3-codex": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1e-05, - "output_cost_per_token_priority": 2e-05, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/responses" ], @@ -4864,7 +4480,9 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4872,17 +4490,22 @@ "supports_response_schema": true, "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5.1-codex-max": { - "cache_read_input_token_cost": 1.25e-07, - "input_cost_per_token": 1.25e-06, + "gpt-5.3-codex-spark": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1e-05, + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, "supported_endpoints": [ "/v1/responses" ], @@ -4894,7 +4517,9 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -4902,21 +4527,33 @@ "supports_response_schema": true, "supports_system_messages": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": false }, - "gpt-5.1-codex-mini": { - "cache_read_input_token_cost": 2.5e-08, - "cache_read_input_token_cost_priority": 4.5e-08, - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_priority": 4.5e-07, + "gpt-5.4": { + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost_flex": 1.3e-07, + "cache_read_input_token_cost_priority": 5e-07, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 5e-06, + "input_cost_per_token_batches": 1.25e-06, + "input_cost_per_token_flex": 1.25e-06, + "input_cost_per_token_priority": 5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "responses", - "output_cost_per_token": 2e-06, - "output_cost_per_token_priority": 3.6e-06, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, + "output_cost_per_token_batches": 7.5e-06, + "output_cost_per_token_flex": 7.5e-06, + "output_cost_per_token_priority": 3e-05, "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -4927,28 +4564,40 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_system_messages": false, + "supports_service_tier": true, + "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.4-2026-03-05": { + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost_flex": 1.3e-07, + "cache_read_input_token_cost_priority": 5e-07, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 5e-06, + "input_cost_per_token_batches": 1.25e-06, + "input_cost_per_token_flex": 1.25e-06, + "input_cost_per_token_priority": 5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, + "output_cost_per_token_batches": 7.5e-06, + "output_cost_per_token_flex": 7.5e-06, + "output_cost_per_token_priority": 3e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -4959,8 +4608,7 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, "supports_native_streaming": true, @@ -4974,18 +4622,24 @@ "supports_tool_choice": true, "supports_vision": true }, - "gpt-5.2-2025-12-11": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.4-mini": { + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_batches": 3.75e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.5e-07, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "output_cost_per_token": 4.5e-06, + "output_cost_per_token_batches": 2.25e-06, + "output_cost_per_token_flex": 2.25e-06, + "output_cost_per_token_priority": 9e-06, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -4996,11 +4650,12 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -5009,22 +4664,31 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2-chat-latest": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.4-mini-2026-03-17": { + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_batches": 3.75e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.5e-07, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "output_cost_per_token": 4.5e-06, + "output_cost_per_token_batches": 2.25e-06, + "output_cost_per_token_flex": 2.25e-06, + "output_cost_per_token_priority": 9e-06, "supported_endpoints": [ "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -5035,29 +4699,39 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2-codex": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_batches": 1e-08, + "cache_read_input_token_cost_flex": 1e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_token_flex": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "responses", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "output_cost_per_token_batches": 6.25e-07, + "output_cost_per_token_flex": 6.25e-07, "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -5068,25 +4742,38 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_system_messages": false, + "supports_service_tier": true, + "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2-pro": { - "input_cost_per_token": 2.1e-05, + "gpt-5.4-nano-2026-03-17": { + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_batches": 1e-08, + "cache_read_input_token_cost_flex": 1e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_token_flex": 1e-07, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "responses", - "output_cost_per_token": 0.000168, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "output_cost_per_token_batches": 6.25e-07, + "output_cost_per_token_flex": 6.25e-07, "supported_endpoints": [ + "/v1/chat/completions", "/v1/batch", "/v1/responses" ], @@ -5098,28 +4785,40 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.2-pro-2025-12-11": { - "input_cost_per_token": 2.1e-05, + "gpt-5.4-pro": { + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_batches": 1.5e-05, + "input_cost_per_token_flex": 1.5e-05, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.000168, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_batches": 9e-05, + "output_cost_per_token_flex": 9e-05, "supported_endpoints": [ - "/v1/batch", - "/v1/responses" + "/v1/responses", + "/v1/batch" ], "supported_modalities": [ "text", @@ -5129,61 +4828,87 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.4": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 2.5e-06, + "gpt-5.4-pro-2026-03-05": { + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_batches": 1.5e-05, + "input_cost_per_token_flex": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1.5e-05, + "mode": "responses", + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_batches": 9e-05, + "output_cost_per_token_flex": 9e-05, "supported_endpoints": [ - "/v1/chat/completions", - "/v1/responses" + "/v1/responses", + "/v1/batch" ], "supported_modalities": [ "text", "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "codex-auto-review": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 2.5e-06, + "gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_priority": 1e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_priority": 6e-05, "supported_endpoints": [ "/v1/chat/completions", + "/v1/batch", "/v1/responses" ], "supported_modalities": [ @@ -5191,11 +4916,12 @@ "image" ], "supported_output_modalities": [ - "text", - "image" + "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -5204,17 +4930,30 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.4-mini": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_token": 7.5e-07, + "gpt-5.5-2026-04-23": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_priority": 1e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 4.5e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_priority": 6e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -5228,7 +4967,9 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -5237,21 +4978,29 @@ "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.4-nano": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 2e-07, + "gpt-5.5-pro": { + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_batches": 1.5e-05, + "input_cost_per_token_flex": 1.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 1.25e-06, + "mode": "responses", + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_batches": 9e-05, + "output_cost_per_token_flex": 9e-05, "supported_endpoints": [ - "/v1/chat/completions", - "/v1/batch", - "/v1/responses" + "/v1/responses", + "/v1/batch" ], "supported_modalities": [ "text", @@ -5261,30 +5010,41 @@ "text" ], "supports_function_calling": true, + "supports_low_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, + "supports_service_tier": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, - "gpt-5.3-codex": { - "cache_read_input_token_cost": 1.75e-07, - "cache_read_input_token_cost_priority": 3.5e-07, - "input_cost_per_token": 1.75e-06, - "input_cost_per_token_priority": 3.5e-06, + "gpt-5.5-pro-2026-04-23": { + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_batches": 1.5e-05, + "input_cost_per_token_flex": 1.5e-05, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1.4e-05, - "output_cost_per_token_priority": 2.8e-05, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_batches": 9e-05, + "output_cost_per_token_flex": 9e-05, "supported_endpoints": [ - "/v1/responses" + "/v1/responses", + "/v1/batch" ], "supported_modalities": [ "text", @@ -5294,15 +5054,21 @@ "text" ], "supports_function_calling": true, + "supports_low_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": false, + "supports_response_schema": false, + "supports_service_tier": true, + "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true }, "gpt-audio": { "input_cost_per_audio_token": 3.2e-05, @@ -5340,6 +5106,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "gpt-audio-1.5": { + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "gpt-audio-2025-08-28": { "input_cost_per_audio_token": 3.2e-05, "input_cost_per_token": 2.5e-06, @@ -5540,6 +5339,38 @@ "supports_pdf_input": true, "supports_vision": true }, + "gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, + "gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, "gpt-realtime": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -5572,6 +5403,70 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "gpt-realtime-1.5": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "openai", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "gpt-realtime-2": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "openai", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -5720,62 +5615,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "o1-mini": { - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 1.1e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 4.4e-06, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_vision": true - }, - "o1-mini-2024-09-12": { - "cache_read_input_token_cost": 1.5e-06, - "deprecation_date": "2025-10-27", - "input_cost_per_token": 3e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, - "o1-preview": { - "cache_read_input_token_cost": 7.5e-06, - "input_cost_per_token": 1.5e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, - "o1-preview-2024-09-12": { - "cache_read_input_token_cost": 7.5e-06, - "input_cost_per_token": 1.5e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, "o1-pro": { "input_cost_per_token": 0.00015, "input_cost_per_token_batches": 7.5e-05, @@ -5876,7 +5715,8 @@ "supports_response_schema": true, "supports_service_tier": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-2025-04-16": { "cache_read_input_token_cost": 5e-07, @@ -5908,7 +5748,8 @@ "supports_response_schema": true, "supports_service_tier": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, @@ -5941,7 +5782,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-deep-research-2025-06-26": { "cache_read_input_token_cost": 2.5e-06, @@ -5974,7 +5816,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-mini": { "cache_read_input_token_cost": 5.5e-07, @@ -6038,7 +5881,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-pro-2025-06-10": { "input_cost_per_token": 2e-05, @@ -6068,7 +5912,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini": { "cache_read_input_token_cost": 2.75e-07, @@ -6093,7 +5938,8 @@ "supports_response_schema": true, "supports_service_tier": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, @@ -6112,7 +5958,8 @@ "supports_response_schema": true, "supports_service_tier": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-deep-research": { "cache_read_input_token_cost": 5e-07, @@ -6145,7 +5992,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-deep-research-2025-06-26": { "cache_read_input_token_cost": 5e-07, @@ -6178,6 +6026,7 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true } } diff --git a/frontend/src/api/admin/accounts.ts b/frontend/src/api/admin/accounts.ts index 6fd23c475d4e..bb75e30251d3 100644 --- a/frontend/src/api/admin/accounts.ts +++ b/frontend/src/api/admin/accounts.ts @@ -487,6 +487,23 @@ export async function syncUpstreamModels(id: number): Promise { + const { data } = await apiClient.post('/admin/accounts/models/sync-upstream-preview', params) + return data +} + export interface CRSPreviewAccount { crs_account_id: string kind: string @@ -703,6 +720,7 @@ export const accountsAPI = { setSchedulable, getAvailableModels, syncUpstreamModels, + syncUpstreamModelsPreview, generateAuthUrl, exchangeCode, refreshOpenAIToken, diff --git a/frontend/src/api/admin/ops.ts b/frontend/src/api/admin/ops.ts index 6923566867a9..847fc8c92bc7 100644 --- a/frontend/src/api/admin/ops.ts +++ b/frontend/src/api/admin/ops.ts @@ -778,9 +778,15 @@ export interface OpsAlertRuntimeSettings { thresholds: OpsMetricThresholds // 指标阈值配置 } +export interface OpsOpenAIAccountQuotaAutoPauseSettings { + default_threshold_5h: number // 0~1，0 表示不启用全局默认 5h 阈值 + default_threshold_7d: number // 0~1，0 表示不启用全局默认 7d 阈值 +} + export interface OpsAdvancedSettings { data_retention: OpsDataRetentionSettings aggregation: OpsAggregationSettings + openai_account_quota_auto_pause: OpsOpenAIAccountQuotaAutoPauseSettings ignore_count_tokens_errors: boolean ignore_context_canceled: boolean ignore_no_available_accounts: boolean diff --git a/frontend/src/api/admin/usage.ts b/frontend/src/api/admin/usage.ts index 7ad00742fbde..d933ac6351f0 100644 --- a/frontend/src/api/admin/usage.ts +++ b/frontend/src/api/admin/usage.ts @@ -27,6 +27,7 @@ export interface AdminUsageStatsResponse { export interface SimpleUser { id: number email: string + deleted: boolean } export interface SimpleApiKey { @@ -120,6 +121,7 @@ export async function getStats(params: { start_date?: string end_date?: string timezone?: string + nocache?: number }): Promise { const { data } = await apiClient.get('/admin/usage/stats', { params diff --git a/frontend/src/api/admin/users.ts b/frontend/src/api/admin/users.ts index bfe5e3baec38..61879d2f55eb 100644 --- a/frontend/src/api/admin/users.ts +++ b/frontend/src/api/admin/users.ts @@ -100,10 +100,12 @@ export async function list( /** * Get user by ID * @param id - User ID + * @param includeDeleted - Whether to include soft-deleted users * @returns User details */ -export async function getById(id: number): Promise { - const { data } = await apiClient.get(`/admin/users/${id}`) +export async function getById(id: number, includeDeleted = false): Promise { + const url = includeDeleted ? `/admin/users/${id}?include_deleted=true` : `/admin/users/${id}` + const { data } = await apiClient.get(url) return data } @@ -115,8 +117,11 @@ export async function getById(id: number): Promise { export async function create(userData: { email: string password: string + username?: string + notes?: string balance?: number concurrency?: number + rpm_limit?: number allowed_groups?: number[] | null }): Promise { const { data } = await apiClient.post('/admin/users', userData) diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index a8abfd9f4191..d1a7729a58d3 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -1124,7 +1124,7 @@
- +

{{ t('admin.accounts.selectedModels', { count: allowedModels.length }) }} {{ @@ -1562,7 +1562,7 @@

- +

{{ t('admin.accounts.selectedModels', { count: allowedModels.length }) }} {{ t('admin.accounts.supportsAllModels') }} @@ -1813,7 +1813,7 @@

- +

{{ t('admin.accounts.selectedModels', { count: allowedModels.length }) }} {{ @@ -3361,6 +3361,17 @@ const accountCategory = ref<'oauth-based' | 'apikey' | 'bedrock' | 'service_acco const addMethod = ref('oauth') // For oauth-based: 'oauth' or 'setup-token' const apiKeyBaseUrl = ref('https://api.anthropic.com') const apiKeyValue = ref('') + +const syncPreviewCredentials = computed(() => { + if (!apiKeyValue.value) return undefined + return { + platform: form.platform, + type: form.type, + base_url: apiKeyBaseUrl.value || undefined, + api_key: apiKeyValue.value + } +}) + const editQuotaLimit = ref(null) const editQuotaDailyLimit = ref(null) const editQuotaWeeklyLimit = ref(null) diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 515ba58f7c7e..8dc85d0eb841 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -1787,6 +1787,84 @@

+
+
+
+ + +
+

{{ t('admin.accounts.autoPauseDisabledHint') }}

+
+
+ + +

{{ t('admin.accounts.autoPauseThresholdHint') }}

+
+
+
+ + +
+

{{ t('admin.accounts.autoPauseDisabledHint') }}

+
+
+ + +

{{ t('admin.accounts.autoPauseThresholdHint') }}

+
+
+
([]) const customErrorCodeInput = ref(null) const interceptWarmupRequests = ref(false) const autoPauseOnExpired = ref(false) +const autoPause5hThreshold = ref(null) +const autoPause7dThreshold = ref(null) +const autoPause5hDisabled = ref(false) +const autoPause7dDisabled = ref(false) const mixedScheduling = ref(false) // For antigravity accounts: enable mixed scheduling const allowOverages = ref(false) // For antigravity accounts: enable AI Credits overages const antigravityModelRestrictionMode = ref<'whitelist' | 'mapping'>('whitelist') @@ -2862,9 +2944,13 @@ const syncFormFromAccount = (newAccount: Account | null) => { // Load mixed scheduling setting (only for antigravity accounts) mixedScheduling.value = false allowOverages.value = false - const extra = newAccount.extra as Record | undefined - mixedScheduling.value = extra?.mixed_scheduling === true - allowOverages.value = extra?.allow_overages === true + const extra = newAccount.extra as Record | undefined + mixedScheduling.value = extra?.mixed_scheduling === true + allowOverages.value = extra?.allow_overages === true + autoPause5hThreshold.value = typeof extra?.auto_pause_5h_threshold === 'number' ? extra.auto_pause_5h_threshold * 100 : null + autoPause7dThreshold.value = typeof extra?.auto_pause_7d_threshold === 'number' ? extra.auto_pause_7d_threshold * 100 : null + autoPause5hDisabled.value = extra?.auto_pause_5h_disabled === true + autoPause7dDisabled.value = extra?.auto_pause_7d_disabled === true // Load OpenAI passthrough toggle (OpenAI OAuth/API Key) openaiPassthroughEnabled.value = false @@ -3987,9 +4073,9 @@ const handleSubmit = async () => { } // For OpenAI OAuth/API Key accounts, handle passthrough mode in extra - if (props.account.platform === 'openai' && (props.account.type === 'oauth' || props.account.type === 'apikey')) { - const currentExtra = (props.account.extra as Record) || {} - const newExtra: Record = { ...currentExtra } + if (props.account.platform === 'openai' && (props.account.type === 'oauth' || props.account.type === 'apikey')) { + const currentExtra = (props.account.extra as Record) || {} + const newExtra: Record = { ...currentExtra } const hadCodexCLIOnlyEnabled = currentExtra.codex_cli_only === true if (props.account.type === 'oauth') { newExtra.openai_oauth_responses_websockets_v2_mode = openaiOAuthResponsesWebSocketV2Mode.value @@ -4011,15 +4097,35 @@ const handleSubmit = async () => { } else { newExtra.openai_compact_mode = openAICompactMode.value } - if (props.account.type === 'apikey') { + if (props.account.type === 'apikey') { if (!openAITextGenerationCapabilityEnabled.value || openAIResponsesMode.value === 'auto') { delete newExtra.openai_responses_mode } else { newExtra.openai_responses_mode = openAIResponsesMode.value } - } - - delete newExtra.codex_image_generation_bridge_enabled + } + if (autoPause5hThreshold.value != null && autoPause5hThreshold.value > 0) { + newExtra.auto_pause_5h_threshold = autoPause5hThreshold.value / 100 + } else { + delete newExtra.auto_pause_5h_threshold + } + if (autoPause7dThreshold.value != null && autoPause7dThreshold.value > 0) { + newExtra.auto_pause_7d_threshold = autoPause7dThreshold.value / 100 + } else { + delete newExtra.auto_pause_7d_threshold + } + if (autoPause5hDisabled.value) { + newExtra.auto_pause_5h_disabled = true + } else { + delete newExtra.auto_pause_5h_disabled + } + if (autoPause7dDisabled.value) { + newExtra.auto_pause_7d_disabled = true + } else { + delete newExtra.auto_pause_7d_disabled + } + + delete newExtra.codex_image_generation_bridge_enabled if (codexImageGenerationBridgeMode.value === 'inherit') { delete newExtra.codex_image_generation_bridge } else { diff --git a/frontend/src/components/account/ModelWhitelistSelector.vue b/frontend/src/components/account/ModelWhitelistSelector.vue index 9a0d6af80f39..d4d726d06ab8 100644 --- a/frontend/src/components/account/ModelWhitelistSelector.vue +++ b/frontend/src/components/account/ModelWhitelistSelector.vue @@ -133,6 +133,7 @@ import { ref, computed } from 'vue' import { useI18n } from 'vue-i18n' import { useAppStore } from '@/stores/app' import { accountsAPI } from '@/api/admin/accounts' +import type { SyncUpstreamPreviewParams } from '@/api/admin/accounts' import ModelIcon from '@/components/common/ModelIcon.vue' import Icon from '@/components/icons/Icon.vue' import { allModels, getModelsByPlatform } from '@/composables/useModelWhitelist' @@ -144,6 +145,12 @@ const props = defineProps<{ platform?: string platforms?: string[] accountId?: number + syncCredentials?: { + platform: string + type: string + base_url?: string + api_key: string + } }>() const emit = defineEmits<{ @@ -176,9 +183,14 @@ const normalizedPlatforms = computed(() => { const upstreamSyncPlatforms = new Set(['anthropic', 'openai', 'gemini', 'antigravity']) const canSyncUpstream = computed(() => { - if (!props.accountId) return false - if (normalizedPlatforms.value.length === 0) return true - return normalizedPlatforms.value.some(platform => upstreamSyncPlatforms.has(platform.toLowerCase())) + if (props.accountId) { + if (normalizedPlatforms.value.length === 0) return true + return normalizedPlatforms.value.some(platform => upstreamSyncPlatforms.has(platform.toLowerCase())) + } + if (props.syncCredentials) { + return upstreamSyncPlatforms.has(props.syncCredentials.platform.toLowerCase()) + } + return false }) const availableOptions = computed(() => { @@ -249,11 +261,20 @@ const fillRelated = () => { } const syncUpstreamModels = async () => { - if (!props.accountId || isSyncingUpstream.value) return + if (isSyncingUpstream.value) return + if (!props.accountId && !props.syncCredentials) return isSyncingUpstream.value = true try { - const result = await accountsAPI.syncUpstreamModels(props.accountId) + let result + if (props.accountId) { + result = await accountsAPI.syncUpstreamModels(props.accountId) + } else if (props.syncCredentials) { + result = await accountsAPI.syncUpstreamModelsPreview(props.syncCredentials as SyncUpstreamPreviewParams) + } else { + return + } + const upstreamModels = result.models.map(model => model.trim()).filter(Boolean) if (upstreamModels.length === 0) { appStore.showInfo(t('admin.accounts.syncUpstreamModelsEmpty')) diff --git a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts index 4561924f675d..f4865de9e6c5 100644 --- a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts @@ -330,6 +330,49 @@ describe('EditAccountModal', () => { ]) }) + it('submits OpenAI quota auto-pause thresholds in extra', async () => { + const account = buildAccount() + account.extra = { + auto_pause_5h_threshold: 0.9, + auto_pause_7d_threshold: 0.8 + } + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + + await wrapper.get('[data-testid="auto-pause-5h-threshold"]').setValue('95') + await wrapper.get('[data-testid="auto-pause-7d-threshold"]').setValue('96') + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.auto_pause_5h_threshold).toBe(0.95) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.auto_pause_7d_threshold).toBe(0.96) + }) + + it('submits OpenAI quota auto-pause disable flag in extra', async () => { + // Toggling the per-account disable flag must persist as auto_pause_5h_disabled + // so an admin can exempt one account from auto-pause even when a global default + // threshold is configured (otherwise leaving the threshold blank would silently + // fall back to the global default). + const account = buildAccount() + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + + await wrapper.get('[data-testid="auto-pause-5h-disabled"]').trigger('click') + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.auto_pause_5h_disabled).toBe(true) + expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.auto_pause_7d_disabled).toBeUndefined() + }) + it('keeps at least one OpenAI APIKey endpoint capability selected', async () => { const account = buildAccount() updateAccountMock.mockReset() diff --git a/frontend/src/components/admin/usage/UsageFilters.vue b/frontend/src/components/admin/usage/UsageFilters.vue index 66c2b4fa6504..a800f1903d1b 100644 --- a/frontend/src/components/admin/usage/UsageFilters.vue +++ b/frontend/src/components/admin/usage/UsageFilters.vue @@ -35,7 +35,7 @@ @click="selectUser(u)" class="w-full px-4 py-2 text-left hover:bg-gray-100 dark:hover:bg-gray-700" > - {{ u.email }} + {{ u.email }}（{{ t('admin.usage.userDeletedBadge') }}） #{{ u.id }}
@@ -168,7 +168,7 @@ diff --git a/frontend/src/views/admin/__tests__/AccountsView.usageWindowsHint.spec.ts b/frontend/src/views/admin/__tests__/AccountsView.usageWindowsHint.spec.ts new file mode 100644 index 000000000000..81e7d87e0eb7 --- /dev/null +++ b/frontend/src/views/admin/__tests__/AccountsView.usageWindowsHint.spec.ts @@ -0,0 +1,164 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { flushPromises, mount } from '@vue/test-utils' + +import AccountsView from '../AccountsView.vue' + +const { + listAccounts, + listWithEtag, + getBatchTodayStats, + getAllProxies, + getAllGroups +} = vi.hoisted(() => ({ + listAccounts: vi.fn(), + listWithEtag: vi.fn(), + getBatchTodayStats: vi.fn(), + getAllProxies: vi.fn(), + getAllGroups: vi.fn() +})) + +vi.mock('@/api/admin', () => ({ + adminAPI: { + accounts: { + list: listAccounts, + listWithEtag, + getBatchTodayStats, + delete: vi.fn(), + batchClearError: vi.fn(), + batchRefresh: vi.fn(), + toggleSchedulable: vi.fn() + }, + proxies: { + getAll: getAllProxies + }, + groups: { + getAll: getAllGroups + } + } +})) + +vi.mock('@/stores/app', () => ({ + useAppStore: () => ({ + showError: vi.fn(), + showSuccess: vi.fn(), + showInfo: vi.fn() + }) +})) + +vi.mock('@/stores/auth', () => ({ + useAuthStore: () => ({ + token: 'test-token' + }) +})) + +vi.mock('vue-i18n', async () => { + const actual = await vi.importActual('vue-i18n') + return { + ...actual, + useI18n: () => ({ + t: (key: string) => key + }) + } +}) + +// Render the per-column header slots so we can assert the usage-window header hint. +const DataTableStub = { + props: ['columns', 'data'], + template: ` +
+ +
+ ` +} + +// Expose the content passed to HelpTooltip without dealing with its . +const HelpTooltipStub = { + props: ['content', 'widthClass'], + template: '{{ content }}' +} + +function mountView() { + return mount(AccountsView, { + global: { + stubs: { + AppLayout: { template: '
' }, + TablePageLayout: { + template: '
' + }, + DataTable: DataTableStub, + HelpTooltip: HelpTooltipStub, + Pagination: true, + ConfirmDialog: true, + AccountTableActions: { template: '
' }, + AccountTableFilters: { template: '
' }, + AccountBulkActionsBar: true, + AccountActionMenu: true, + ImportDataModal: true, + ReAuthAccountModal: true, + AccountTestModal: true, + AccountStatsModal: true, + ScheduledTestsPanel: true, + SyncFromCrsModal: true, + TempUnschedStatusModal: true, + ErrorPassthroughRulesModal: true, + TLSFingerprintProfilesModal: true, + CreateAccountModal: true, + EditAccountModal: true, + BulkEditAccountModal: true, + PlatformTypeBadge: true, + AccountCapacityCell: true, + AccountStatusIndicator: true, + AccountTodayStatsCell: true, + AccountGroupsCell: true, + AccountUsageCell: true, + Icon: true + } + } + }) +} + +describe('admin AccountsView usage windows hint', () => { + beforeEach(() => { + localStorage.clear() + + listAccounts.mockReset() + listWithEtag.mockReset() + getBatchTodayStats.mockReset() + getAllProxies.mockReset() + getAllGroups.mockReset() + + listAccounts.mockResolvedValue({ + items: [], + total: 0, + page: 1, + page_size: 20, + pages: 0 + }) + listWithEtag.mockResolvedValue({ + notModified: true, + etag: null, + data: null + }) + getBatchTodayStats.mockResolvedValue({ stats: {} }) + getAllProxies.mockResolvedValue([]) + getAllGroups.mockResolvedValue([]) + }) + + it('renders an explanatory tooltip next to the usage windows column header', async () => { + const wrapper = mountView() + await flushPromises() + + const header = wrapper.find('[data-test="usage-header"]') + expect(header.exists()).toBe(true) + // Column label is still shown alongside the help icon. + expect(header.text()).toContain('admin.accounts.columns.usageWindows') + + const hint = wrapper.find('[data-test="usage-windows-hint"]') + expect(hint.exists()).toBe(true) + expect(hint.text()).toBe('admin.accounts.usageWindowsHint') + }) +}) diff --git a/frontend/src/views/admin/__tests__/UsageView.spec.ts b/frontend/src/views/admin/__tests__/UsageView.spec.ts index 1a5d285a8603..8c644a7536d2 100644 --- a/frontend/src/views/admin/__tests__/UsageView.spec.ts +++ b/frontend/src/views/admin/__tests__/UsageView.spec.ts @@ -3,7 +3,7 @@ import { flushPromises, mount } from '@vue/test-utils' import UsageView from '../UsageView.vue' -const { list, getStats, getSnapshotV2, getById } = vi.hoisted(() => { +const { list, getStats, getSnapshotV2, getById, getModelStats } = vi.hoisted(() => { vi.stubGlobal('localStorage', { getItem: vi.fn(() => null), setItem: vi.fn(), @@ -15,6 +15,7 @@ const { list, getStats, getSnapshotV2, getById } = vi.hoisted(() => { getStats: vi.fn(), getSnapshotV2: vi.fn(), getById: vi.fn(), + getModelStats: vi.fn(), } }) @@ -40,6 +41,7 @@ vi.mock('@/api/admin', () => ({ }, dashboard: { getSnapshotV2, + getModelStats, }, users: { getById, @@ -84,6 +86,10 @@ vi.mock('vue-router', () => ({ const AppLayoutStub = { template: '
' } const UsageFiltersStub = { template: '
' } +const UsageTableStub = { + emits: ['userClick'], + template: '
', +} const ModelDistributionChartStub = { props: ['metric'], emits: ['update:metric'], @@ -112,6 +118,7 @@ describe('admin UsageView distribution metric toggles', () => { getStats.mockReset() getSnapshotV2.mockReset() getById.mockReset() + getModelStats.mockReset() list.mockResolvedValue({ items: [], @@ -133,12 +140,44 @@ describe('admin UsageView distribution metric toggles', () => { models: [], groups: [], }) + getModelStats.mockResolvedValue({ models: [] }) }) afterEach(() => { vi.useRealTimers() }) + it('keeps previous model stats visible during refresh until new data arrives', async () => { + // 首次加载返回 A + getModelStats.mockResolvedValueOnce({ models: [{ model: 'A', total_tokens: 10 }] }) + + const wrapper = mount(UsageView, { + global: { stubs: { + AppLayout: AppLayoutStub, UsageStatsCards: true, UsageFilters: UsageFiltersStub, + UsageTable: true, UsageExportProgress: true, UsageCleanupDialog: true, + UserBalanceHistoryModal: true, AuditLogModal: true, Pagination: true, Select: true, + DateRangePicker: true, Icon: true, TokenUsageTrend: true, + ModelDistributionChart: ModelDistributionChartStub, GroupDistributionChart: GroupDistributionChartStub, + EndpointDistributionChart: true, + } }, + }) + vi.advanceTimersByTime(120) + await flushPromises() + expect((wrapper.vm as any).requestedModelStats).toEqual([{ model: 'A', total_tokens: 10 }]) + + // 刷新:让第二次 getModelStats 处于 pending,断言旧数据 A 仍在(不被清空成 []) + let resolveSecond: (v: any) => void = () => {} + getModelStats.mockReturnValueOnce(new Promise((res) => { resolveSecond = res })) + ;(wrapper.vm as any).refreshData() + await flushPromises() + expect((wrapper.vm as any).requestedModelStats).toEqual([{ model: 'A', total_tokens: 10 }]) + + // 新数据到达后替换为 B + resolveSecond({ models: [{ model: 'B', total_tokens: 20 }] }) + await flushPromises() + expect((wrapper.vm as any).requestedModelStats).toEqual([{ model: 'B', total_tokens: 20 }]) + }) + it('keeps model and group metric toggles independent without refetching chart data', async () => { const wrapper = mount(UsageView, { global: { @@ -194,3 +233,59 @@ describe('admin UsageView distribution metric toggles', () => { expect(getSnapshotV2).toHaveBeenCalledTimes(1) }) }) + +describe('admin UsageView handleUserClick', () => { + beforeEach(() => { + vi.useFakeTimers() + list.mockReset() + getStats.mockReset() + getSnapshotV2.mockReset() + getById.mockReset() + + list.mockResolvedValue({ items: [], total: 0, pages: 0 }) + getStats.mockResolvedValue({ + total_requests: 0, total_input_tokens: 0, total_output_tokens: 0, + total_cache_tokens: 0, total_tokens: 0, total_cost: 0, total_actual_cost: 0, average_duration_ms: 0, + }) + getSnapshotV2.mockResolvedValue({ trend: [], models: [], groups: [] }) + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it('opens user via include_deleted when clicking a usage row user', async () => { + getById.mockResolvedValue({ id: 2, email: 'd@test.com', deleted_at: '2026-05-28T00:00:00Z' }) + + const wrapper = mount(UsageView, { + global: { + stubs: { + AppLayout: AppLayoutStub, + UsageStatsCards: true, + UsageFilters: UsageFiltersStub, + UsageTable: UsageTableStub, + UsageExportProgress: true, + UsageCleanupDialog: true, + UserBalanceHistoryModal: true, + AuditLogModal: true, + Pagination: true, + Select: true, + DateRangePicker: true, + Icon: true, + TokenUsageTrend: true, + ModelDistributionChart: true, + GroupDistributionChart: true, + EndpointDistributionChart: true, + }, + }, + }) + + vi.advanceTimersByTime(120) + await flushPromises() + + await wrapper.find('[data-test="usage-table"] .user-click').trigger('click') + await flushPromises() + + expect(getById).toHaveBeenCalledWith(2, true) + }) +}) diff --git a/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue b/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue index 5dba5b1dc4c3..bfb7a65f85cf 100644 --- a/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue +++ b/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue @@ -50,6 +50,10 @@ async function loadAllSettings() { runtimeSettings.value = runtime emailConfig.value = email advancedSettings.value = advanced + // 兼容旧 payload：后端未返回该字段时补默认值，保证表单可绑定 + if (advancedSettings.value && !advancedSettings.value.openai_account_quota_auto_pause) { + advancedSettings.value.openai_account_quota_auto_pause = { default_threshold_5h: 0, default_threshold_7d: 0 } + } // 如果后端返回了阈值，使用后端的值；否则保持默认值 if (thresholds && Object.keys(thresholds).length > 0) { metricThresholds.value = { @@ -119,6 +123,28 @@ function removeRecipient(target: 'alert' | 'report', email: string) { if (idx >= 0) list.splice(idx, 1) } +// OpenAI 账号配额自动暂停：后端按 0~1 分数存储，UI 按百分比(0~100)展示 +const quotaAutoPause5hPercent = computed({ + get() { + const v = advancedSettings.value?.openai_account_quota_auto_pause?.default_threshold_5h + return v && v > 0 ? Math.round(v * 1000) / 10 : null + }, + set(val) { + if (!advancedSettings.value?.openai_account_quota_auto_pause) return + advancedSettings.value.openai_account_quota_auto_pause.default_threshold_5h = val != null && val > 0 ? val / 100 : 0 + } +}) +const quotaAutoPause7dPercent = computed({ + get() { + const v = advancedSettings.value?.openai_account_quota_auto_pause?.default_threshold_7d + return v && v > 0 ? Math.round(v * 1000) / 10 : null + }, + set(val) { + if (!advancedSettings.value?.openai_account_quota_auto_pause) return + advancedSettings.value.openai_account_quota_auto_pause.default_threshold_7d = val != null && val > 0 ? val / 100 : 0 + } +}) + // 验证 const validation = computed(() => { const errors: string[] = [] @@ -145,6 +171,11 @@ const validation = computed(() => { if (hourly_metrics_retention_days < 0 || hourly_metrics_retention_days > 365) { errors.push(t('admin.ops.settings.validation.retentionDaysRange')) } + + const { default_threshold_5h, default_threshold_7d } = advancedSettings.value.openai_account_quota_auto_pause + if (default_threshold_5h < 0 || default_threshold_5h > 1 || default_threshold_7d < 0 || default_threshold_7d > 1) { + errors.push(t('admin.ops.settings.validation.openaiQuotaAutoPauseRange')) + } } // 验证指标阈值 @@ -473,6 +504,40 @@ async function saveAllSettings() {
+ +
+
{{ t('admin.ops.settings.openaiQuotaAutoPause') }}
+

{{ t('admin.ops.settings.openaiQuotaAutoPauseHint') }}

+ +
+
+ + +
+
+ + +
+
+

{{ t('admin.ops.settings.openaiQuotaAutoPauseThresholdHint') }}

+
+
{{ t('admin.ops.settings.errorFiltering') }}