diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index c7ab9eeb88..a58ac974b1 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -255,7 +255,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig) userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient) userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig) - gatewayHandler := handler.NewGatewayHandler(gatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService) + gatewayHandler := handler.NewGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService) openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, configConfig) handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService) totpHandler := handler.NewTotpHandler(totpService) diff --git a/backend/internal/handler/admin/channel_handler.go b/backend/internal/handler/admin/channel_handler.go index 30d208c0d2..f5442448c8 100644 --- a/backend/internal/handler/admin/channel_handler.go +++ b/backend/internal/handler/admin/channel_handler.go @@ -64,6 +64,7 @@ type channelModelPricingRequest struct { OutputPrice *float64 `json:"output_price" binding:"omitempty,min=0"` CacheWritePrice *float64 `json:"cache_write_price" binding:"omitempty,min=0"` CacheReadPrice *float64 `json:"cache_read_price" binding:"omitempty,min=0"` + ImageInputPrice *float64 `json:"image_input_price" binding:"omitempty,min=0"` ImageOutputPrice *float64 `json:"image_output_price" binding:"omitempty,min=0"` PerRequestPrice *float64 `json:"per_request_price" binding:"omitempty,min=0"` Intervals []pricingIntervalRequest `json:"intervals"` @@ -115,6 +116,7 @@ type channelModelPricingResponse struct { OutputPrice *float64 `json:"output_price"` CacheWritePrice *float64 `json:"cache_write_price"` CacheReadPrice *float64 `json:"cache_read_price"` + ImageInputPrice *float64 `json:"image_input_price"` ImageOutputPrice *float64 `json:"image_output_price"` PerRequestPrice *float64 `json:"per_request_price"` Intervals []pricingIntervalResponse `json:"intervals"` @@ -222,6 +224,7 @@ func pricingToResponse(p *service.ChannelModelPricing) channelModelPricingRespon OutputPrice: p.OutputPrice, CacheWritePrice: p.CacheWritePrice, CacheReadPrice: p.CacheReadPrice, + ImageInputPrice: p.ImageInputPrice, ImageOutputPrice: p.ImageOutputPrice, PerRequestPrice: p.PerRequestPrice, Intervals: intervals, @@ -273,6 +276,7 @@ func pricingRequestToService(reqs []channelModelPricingRequest) []service.Channe OutputPrice: r.OutputPrice, CacheWritePrice: r.CacheWritePrice, CacheReadPrice: r.CacheReadPrice, + ImageInputPrice: r.ImageInputPrice, ImageOutputPrice: r.ImageOutputPrice, PerRequestPrice: r.PerRequestPrice, Intervals: intervals, @@ -498,6 +502,7 @@ func (h *ChannelHandler) GetModelDefaultPricing(c *gin.Context) { "output_price": pricing.OutputPricePerToken, "cache_write_price": pricing.CacheCreationPricePerToken, "cache_read_price": pricing.CacheReadPricePerToken, + "image_input_price": pricing.ImageInputPricePerToken, "image_output_price": pricing.ImageOutputPricePerToken, }) } diff --git a/backend/internal/handler/available_channel_handler.go b/backend/internal/handler/available_channel_handler.go index 8da4c07973..05d64e2002 100644 --- a/backend/internal/handler/available_channel_handler.go +++ b/backend/internal/handler/available_channel_handler.go @@ -72,6 +72,7 @@ type userSupportedModelPricing struct { OutputPrice *float64 `json:"output_price"` CacheWritePrice *float64 `json:"cache_write_price"` CacheReadPrice *float64 `json:"cache_read_price"` + ImageInputPrice *float64 `json:"image_input_price"` ImageOutputPrice *float64 `json:"image_output_price"` PerRequestPrice *float64 `json:"per_request_price"` Intervals []userPricingIntervalDTO `json:"intervals"` @@ -284,6 +285,7 @@ func toUserPricing(p *service.ChannelModelPricing) *userSupportedModelPricing { OutputPrice: p.OutputPrice, CacheWritePrice: p.CacheWritePrice, CacheReadPrice: p.CacheReadPrice, + ImageInputPrice: p.ImageInputPrice, ImageOutputPrice: p.ImageOutputPrice, PerRequestPrice: p.PerRequestPrice, Intervals: intervals, diff --git a/backend/internal/handler/chat_completions_image_model_test.go b/backend/internal/handler/chat_completions_image_model_test.go new file mode 100644 index 0000000000..06c0b45488 --- /dev/null +++ b/backend/internal/handler/chat_completions_image_model_test.go @@ -0,0 +1,96 @@ +package handler + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestChatCompletionsRejectsGPTImageModelsBeforeScheduling(t *testing.T) { + gin.SetMode(gin.TestMode) + + for _, model := range []string{"gpt-image-1", "gpt-image-1.5", "gpt-image-2"} { + for _, tc := range []struct { + name string + call func(*gin.Context) + }{ + { + name: "gateway", + call: (&GatewayHandler{}).ChatCompletions, + }, + { + name: "openai_gateway", + call: newOpenAIImageChatRejectionHandler(t).ChatCompletions, + }, + } { + t.Run(tc.name+"/"+model, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"` + model + `","messages":[{"role":"user","content":"draw"}]}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + setImageChatTestAuth(c) + + tc.call(c) + + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Equal(t, "invalid_request_error", gjson.Get(recorder.Body.String(), "error.type").String()) + require.Contains(t, gjson.Get(recorder.Body.String(), "error.message").String(), "Chat Completions") + _, selected := c.Get(opsAccountIDKey) + require.False(t, selected, "rejection must happen before account selection") + }) + } + } +} + +func TestOpenAIChatCompletionsImageModelRejectionDoesNotAcquireConcurrency(t *testing.T) { + var acquireCalls atomic.Int64 + cache := &concurrencyCacheMock{ + acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) { + acquireCalls.Add(1) + return true, nil + }, + } + h := newOpenAIImageChatRejectionHandlerWithCache(t, cache) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString( + `{"model":"gpt-image-2","messages":[{"role":"user","content":"draw"}]}`, + )) + setImageChatTestAuth(c) + + h.ChatCompletions(c) + + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Zero(t, acquireCalls.Load(), "rejection must happen before user/account concurrency and scheduling") +} + +func newOpenAIImageChatRejectionHandler(t *testing.T) *OpenAIGatewayHandler { + t.Helper() + return newOpenAIImageChatRejectionHandlerWithCache(t, &concurrencyCacheMock{}) +} + +func newOpenAIImageChatRejectionHandlerWithCache(t *testing.T, cache *concurrencyCacheMock) *OpenAIGatewayHandler { + t.Helper() + return &OpenAIGatewayHandler{ + gatewayService: &service.OpenAIGatewayService{}, + billingCacheService: &service.BillingCacheService{}, + apiKeyService: &service.APIKeyService{}, + concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(cache), SSEPingFormatNone, time.Second), + } +} + +func setImageChatTestAuth(c *gin.Context) { + apiKey := &service.APIKey{ID: 4348, UserID: 4348, User: &service.User{ID: 4348}} + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 1}) +} diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index 3c45c3b95e..7aa12d7a40 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -634,6 +634,8 @@ func usageLogFromServiceUser(l *service.UsageLog) UsageLog { ImageSize: l.ImageSize, ImageInputSize: l.ImageInputSize, ImageOutputSize: l.ImageOutputSize, + ImageInputTokens: l.ImageInputTokens, + ImageInputCost: l.ImageInputCost, ImageOutputTokens: l.ImageOutputTokens, ImageOutputCost: l.ImageOutputCost, ImageSizeSource: l.ImageSizeSource, diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index 619926c1e4..89d3bc5634 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -503,6 +503,8 @@ type UsageLog struct { ImageSize *string `json:"image_size"` ImageInputSize *string `json:"image_input_size"` ImageOutputSize *string `json:"image_output_size"` + ImageInputTokens int `json:"image_input_tokens"` + ImageInputCost float64 `json:"image_input_cost"` ImageOutputTokens int `json:"image_output_tokens"` ImageOutputCost float64 `json:"image_output_cost"` ImageSizeSource *string `json:"image_size_source"` diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 1c51713656..7a23984d90 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -39,6 +39,7 @@ var gatewayCompatibilityMetricsLogCounter atomic.Uint64 // GatewayHandler handles API gateway requests type GatewayHandler struct { gatewayService *service.GatewayService + openAIGatewayService *service.OpenAIGatewayService geminiCompatService *service.GeminiMessagesCompatService antigravityGatewayService *service.AntigravityGatewayService userService *service.UserService @@ -59,6 +60,7 @@ type GatewayHandler struct { // NewGatewayHandler creates a new GatewayHandler func NewGatewayHandler( gatewayService *service.GatewayService, + openAIGatewayService *service.OpenAIGatewayService, geminiCompatService *service.GeminiMessagesCompatService, antigravityGatewayService *service.AntigravityGatewayService, userService *service.UserService, @@ -94,6 +96,7 @@ func NewGatewayHandler( return &GatewayHandler{ gatewayService: gatewayService, + openAIGatewayService: openAIGatewayService, geminiCompatService: geminiCompatService, antigravityGatewayService: antigravityGatewayService, userService: userService, diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index cc9ff14282..84456a9049 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -80,6 +80,10 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", invalidStreamFieldTypeMessage) return } + if service.IsGPTImageGenerationModel(reqModel) { + h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "This model is not supported on the Chat Completions endpoint") + return + } reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream)) setOpsRequestContext(c, reqModel, reqStream) diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index a5060b7263..ce1ae40811 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -85,7 +85,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { setOpsRequestContext(c, reqModel, reqStream) setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(reqStream, false))) requestCtx := c.Request.Context() - if service.IsImageGenerationIntent("/v1/responses", reqModel, body) { + if service.IsImageGenerationIntentForPlatform("/v1/responses", reqModel, body, openAICompatibleRequestPlatform(apiKey)) { requestCtx = service.WithOpenAIImageGenerationIntent(requestCtx) } diff --git a/backend/internal/handler/gateway_key_billing.go b/backend/internal/handler/gateway_key_billing.go new file mode 100644 index 0000000000..2c075d779d --- /dev/null +++ b/backend/internal/handler/gateway_key_billing.go @@ -0,0 +1,108 @@ +package handler + +import ( + "net/http" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" +) + +const keyBillingInfoSchemaVersion = 1 + +type keyBillingInfoResponse struct { + Object string `json:"object"` + SchemaVersion int `json:"schema_version"` + BillingScope string `json:"billing_scope"` + GroupRateMultiplier float64 `json:"group_rate_multiplier"` + UserRateMultiplier *float64 `json:"user_rate_multiplier,omitempty"` + ResolvedRateMultiplier float64 `json:"resolved_rate_multiplier"` + PeakRateEnabled bool `json:"peak_rate_enabled"` + PeakStart *string `json:"peak_start,omitempty"` + PeakEnd *string `json:"peak_end,omitempty"` + PeakRateMultiplier *float64 `json:"peak_rate_multiplier,omitempty"` + AppliedPeakMultiplier *float64 `json:"applied_peak_multiplier,omitempty"` + EffectiveRateMultiplier float64 `json:"effective_rate_multiplier"` + Timezone *string `json:"timezone,omitempty"` + ObservedAt time.Time `json:"observed_at"` +} + +// KeyBillingInfo returns the token billing multiplier effective for the authenticated API key. +// GET /v1/sub2api/billing +func (h *GatewayHandler) KeyBillingInfo(c *gin.Context) { + apiKey, ok := middleware2.GetAPIKeyFromContext(c) + if !ok { + h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key") + return + } + if h.cfg != nil && h.cfg.RunMode == config.RunModeSimple { + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Billing information is not supported in simple mode") + return + } + if apiKey.GroupID == nil { + h.errorResponse(c, http.StatusForbidden, "permission_error", "API key is not assigned to a group") + return + } + if apiKey.Group == nil { + h.errorResponse(c, http.StatusInternalServerError, "api_error", "Billing information is unavailable") + return + } + + resolvedRate, ok := h.resolveKeyBillingRate(c, apiKey) + if !ok { + h.errorResponse(c, http.StatusInternalServerError, "api_error", "Billing information is unavailable") + return + } + + c.Header("Cache-Control", "no-store") + c.JSON(http.StatusOK, buildKeyBillingInfo(apiKey, resolvedRate, timezone.Now())) +} + +func (h *GatewayHandler) resolveKeyBillingRate(c *gin.Context, apiKey *service.APIKey) (float64, bool) { + groupRate := apiKey.Group.RateMultiplier + switch apiKey.Group.Platform { + case service.PlatformOpenAI, service.PlatformGrok: + if h.openAIGatewayService == nil { + return 0, false + } + return h.openAIGatewayService.ResolveUserGroupRateMultiplier(c.Request.Context(), apiKey.UserID, *apiKey.GroupID, groupRate), true + default: + if h.gatewayService == nil { + return 0, false + } + return h.gatewayService.ResolveUserGroupRateMultiplier(c.Request.Context(), apiKey.UserID, *apiKey.GroupID, groupRate), true + } +} + +func buildKeyBillingInfo(apiKey *service.APIKey, resolvedRate float64, now time.Time) keyBillingInfoResponse { + groupRate := apiKey.Group.RateMultiplier + var userRate *float64 + if resolvedRate != groupRate { + userRate = &resolvedRate + } + appliedPeak := apiKey.Group.PeakMultiplierAt(now) + + response := keyBillingInfoResponse{ + Object: "sub2api.key_billing", + SchemaVersion: keyBillingInfoSchemaVersion, + BillingScope: "token", + GroupRateMultiplier: groupRate, + UserRateMultiplier: userRate, + ResolvedRateMultiplier: resolvedRate, + PeakRateEnabled: apiKey.Group.PeakRateEnabled, + EffectiveRateMultiplier: resolvedRate * appliedPeak, + ObservedAt: now.UTC(), + } + if apiKey.Group.PeakRateEnabled { + response.PeakStart = &apiKey.Group.PeakStart + response.PeakEnd = &apiKey.Group.PeakEnd + response.PeakRateMultiplier = &apiKey.Group.PeakRateMultiplier + response.AppliedPeakMultiplier = &appliedPeak + tz := timezone.Location().String() + response.Timezone = &tz + } + return response +} diff --git a/backend/internal/handler/gateway_key_billing_test.go b/backend/internal/handler/gateway_key_billing_test.go new file mode 100644 index 0000000000..ecd85bd56f --- /dev/null +++ b/backend/internal/handler/gateway_key_billing_test.go @@ -0,0 +1,311 @@ +package handler + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type keyBillingUserGroupRateRepo struct { + service.UserGroupRateRepository + rate *float64 + err error + gotUserID int64 + gotGroupID int64 + lookupCalls int +} + +func (r *keyBillingUserGroupRateRepo) GetByUserAndGroup(_ context.Context, userID, groupID int64) (*float64, error) { + r.gotUserID = userID + r.gotGroupID = groupID + r.lookupCalls++ + return r.rate, r.err +} + +func newKeyBillingHandler(repo service.UserGroupRateRepository) *GatewayHandler { + return &GatewayHandler{ + gatewayService: newKeyBillingGatewayService(repo), + openAIGatewayService: newKeyBillingOpenAIGatewayService(repo), + } +} + +func newKeyBillingGatewayService(repo service.UserGroupRateRepository) *service.GatewayService { + return service.NewGatewayService( + nil, nil, nil, nil, nil, nil, repo, nil, nil, nil, nil, nil, nil, nil, + nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + ) +} + +func newKeyBillingOpenAIGatewayService(repo service.UserGroupRateRepository) *service.OpenAIGatewayService { + return service.NewOpenAIGatewayService( + nil, nil, nil, nil, nil, repo, nil, nil, nil, nil, nil, + nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + ) +} + +func newKeyBillingContext(apiKey *service.APIKey) (*gin.Context, *httptest.ResponseRecorder) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil) + if apiKey != nil { + c.Set(string(middleware2.ContextKeyAPIKey), apiKey) + } + return c, w +} + +func TestGatewayHandlerKeyBillingInfoUsesGroupRate(t *testing.T) { + groupID := int64(7) + apiKey := &service.APIKey{ + UserID: 11, + GroupID: &groupID, + Key: "sk-sensitive-value", + Group: &service.Group{ + ID: groupID, + Name: "private-group-name", + RateMultiplier: 0.75, + }, + } + c, w := newKeyBillingContext(apiKey) + + newKeyBillingHandler(nil).KeyBillingInfo(c) + + require.Equal(t, http.StatusOK, w.Code) + require.Equal(t, "no-store", w.Header().Get("Cache-Control")) + var got keyBillingInfoResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got)) + require.Equal(t, "sub2api.key_billing", got.Object) + require.Equal(t, 1, got.SchemaVersion) + require.Equal(t, "token", got.BillingScope) + require.Equal(t, 0.75, got.GroupRateMultiplier) + require.Nil(t, got.UserRateMultiplier) + require.Equal(t, 0.75, got.ResolvedRateMultiplier) + require.False(t, got.PeakRateEnabled) + require.Nil(t, got.PeakStart) + require.Nil(t, got.PeakEnd) + require.Nil(t, got.PeakRateMultiplier) + require.Nil(t, got.AppliedPeakMultiplier) + require.Equal(t, 0.75, got.EffectiveRateMultiplier) + require.Nil(t, got.Timezone) + require.False(t, got.ObservedAt.IsZero()) + var fields map[string]json.RawMessage + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &fields)) + require.NotContains(t, fields, "user_rate_multiplier") + require.NotContains(t, fields, "peak_start") + require.NotContains(t, fields, "peak_end") + require.NotContains(t, fields, "peak_rate_multiplier") + require.NotContains(t, fields, "applied_peak_multiplier") + require.NotContains(t, fields, "timezone") + require.NotContains(t, w.Body.String(), apiKey.Key) + require.NotContains(t, w.Body.String(), apiKey.Group.Name) +} + +func TestGatewayHandlerKeyBillingInfoUsesUserOverride(t *testing.T) { + groupID := int64(7) + userRate := 0.5 + apiKey := &service.APIKey{ + UserID: 11, + GroupID: &groupID, + Group: &service.Group{ID: groupID, RateMultiplier: 0.75}, + } + c, w := newKeyBillingContext(apiKey) + repo := &keyBillingUserGroupRateRepo{rate: &userRate} + + newKeyBillingHandler(repo).KeyBillingInfo(c) + + require.Equal(t, http.StatusOK, w.Code) + require.Equal(t, 1, repo.lookupCalls) + require.Equal(t, apiKey.UserID, repo.gotUserID) + require.Equal(t, groupID, repo.gotGroupID) + var got keyBillingInfoResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got)) + require.NotNil(t, got.UserRateMultiplier) + require.Equal(t, 0.5, *got.UserRateMultiplier) + require.Equal(t, 0.5, got.ResolvedRateMultiplier) + require.Equal(t, 0.5, got.EffectiveRateMultiplier) +} + +func TestBuildKeyBillingInfoAppliesPeakMultiplier(t *testing.T) { + groupID := int64(7) + apiKey := &service.APIKey{ + GroupID: &groupID, + Group: &service.Group{ + ID: groupID, + RateMultiplier: 1.2, + SubscriptionType: service.SubscriptionTypeSubscription, + PeakRateEnabled: true, + PeakStart: "09:00", + PeakEnd: "18:00", + PeakRateMultiplier: 1.5, + }, + } + now := time.Date(2026, time.July, 12, 10, 0, 0, 0, timezone.Location()) + userRate := 0.8 + + got := buildKeyBillingInfo(apiKey, userRate, now) + + require.Equal(t, 1.2, got.GroupRateMultiplier) + require.NotNil(t, got.UserRateMultiplier) + require.Equal(t, 0.8, *got.UserRateMultiplier) + require.Equal(t, 0.8, got.ResolvedRateMultiplier) + require.True(t, got.PeakRateEnabled) + require.NotNil(t, got.PeakStart) + require.Equal(t, "09:00", *got.PeakStart) + require.NotNil(t, got.PeakEnd) + require.Equal(t, "18:00", *got.PeakEnd) + require.NotNil(t, got.PeakRateMultiplier) + require.Equal(t, 1.5, *got.PeakRateMultiplier) + require.NotNil(t, got.AppliedPeakMultiplier) + require.Equal(t, 1.5, *got.AppliedPeakMultiplier) + require.InDelta(t, 1.2, got.EffectiveRateMultiplier, 1e-12) + require.NotNil(t, got.Timezone) + require.Equal(t, timezone.Location().String(), *got.Timezone) + require.Equal(t, now.UTC(), got.ObservedAt) + + encoded, err := json.Marshal(got) + require.NoError(t, err) + var fields map[string]json.RawMessage + require.NoError(t, json.Unmarshal(encoded, &fields)) + for _, field := range []string{ + "user_rate_multiplier", + "peak_start", + "peak_end", + "peak_rate_multiplier", + "applied_peak_multiplier", + "timezone", + } { + require.Contains(t, fields, field) + } +} + +func TestKeyBillingInfoJSONKeepsZeroPeakMultiplierWhenEnabled(t *testing.T) { + groupID := int64(7) + apiKey := &service.APIKey{ + GroupID: &groupID, + Group: &service.Group{ + ID: groupID, + SubscriptionType: service.SubscriptionTypeSubscription, + PeakRateEnabled: true, + PeakStart: "00:00", + PeakEnd: "23:59", + PeakRateMultiplier: 0, + }, + } + now := time.Date(2026, time.July, 12, 12, 0, 0, 0, timezone.Location()) + encoded, err := json.Marshal(buildKeyBillingInfo(apiKey, apiKey.Group.RateMultiplier, now)) + require.NoError(t, err) + + var fields map[string]json.RawMessage + require.NoError(t, json.Unmarshal(encoded, &fields)) + require.JSONEq(t, "0", string(fields["peak_rate_multiplier"])) + require.JSONEq(t, "0", string(fields["applied_peak_multiplier"])) +} + +func TestGatewayHandlerKeyBillingInfoErrorsAreSafe(t *testing.T) { + t.Run("missing API key", func(t *testing.T) { + c, w := newKeyBillingContext(nil) + newKeyBillingHandler(nil).KeyBillingInfo(c) + require.Equal(t, http.StatusUnauthorized, w.Code) + }) + + t.Run("ungrouped API key", func(t *testing.T) { + c, w := newKeyBillingContext(&service.APIKey{}) + newKeyBillingHandler(nil).KeyBillingInfo(c) + require.Equal(t, http.StatusForbidden, w.Code) + }) + + t.Run("missing billing service", func(t *testing.T) { + groupID := int64(7) + c, w := newKeyBillingContext(&service.APIKey{ + UserID: 11, + GroupID: &groupID, + Group: &service.Group{ID: groupID, RateMultiplier: 1}, + }) + (&GatewayHandler{}).KeyBillingInfo(c) + require.Equal(t, http.StatusInternalServerError, w.Code) + }) + + t.Run("rate lookup failure matches billing fallback", func(t *testing.T) { + groupID := int64(7) + c, w := newKeyBillingContext(&service.APIKey{ + UserID: 11, + GroupID: &groupID, + Group: &service.Group{ID: groupID, RateMultiplier: 1}, + }) + newKeyBillingHandler(&keyBillingUserGroupRateRepo{err: errors.New("database password leaked")}).KeyBillingInfo(c) + require.Equal(t, http.StatusOK, w.Code) + var got keyBillingInfoResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got)) + require.Equal(t, 1.0, got.ResolvedRateMultiplier) + require.NotContains(t, w.Body.String(), "database password leaked") + }) +} + +func TestGatewayHandlerKeyBillingInfoSharesBillingResolverCacheByPlatform(t *testing.T) { + for _, tc := range []struct { + name string + platform string + openAI bool + }{ + {name: "anthropic", platform: service.PlatformAnthropic}, + {name: "openai", platform: service.PlatformOpenAI, openAI: true}, + {name: "grok", platform: service.PlatformGrok, openAI: true}, + } { + t.Run(tc.name, func(t *testing.T) { + groupID := int64(7) + oldRate, newRate := 0.5, 1.8 + repo := &keyBillingUserGroupRateRepo{rate: &oldRate} + gatewayService := newKeyBillingGatewayService(repo) + openAIGatewayService := newKeyBillingOpenAIGatewayService(repo) + h := &GatewayHandler{ + gatewayService: gatewayService, + openAIGatewayService: openAIGatewayService, + } + apiKey := &service.APIKey{ + UserID: 11, + GroupID: &groupID, + Group: &service.Group{ + ID: groupID, + Platform: tc.platform, + RateMultiplier: 0.75, + }, + } + + if tc.openAI { + require.Equal(t, oldRate, openAIGatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier)) + } else { + require.Equal(t, oldRate, gatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier)) + } + repo.rate = &newRate + + for range 2 { + c, w := newKeyBillingContext(apiKey) + h.KeyBillingInfo(c) + require.Equal(t, http.StatusOK, w.Code) + var got keyBillingInfoResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got)) + require.Equal(t, oldRate, got.ResolvedRateMultiplier) + require.Equal(t, oldRate, got.EffectiveRateMultiplier) + } + + var billedRate float64 + if tc.openAI { + billedRate = openAIGatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier) + } else { + billedRate = gatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier) + } + require.Equal(t, oldRate, billedRate) + require.Equal(t, 1, repo.lookupCalls) + }) + } +} diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 2b293fbbbe..eea6149e8c 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -274,7 +274,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } if failoverErr.ShouldReportAccountScheduleFailure() { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil) } if c.Writer.Size() != writerSizeBeforeForward { h.handleFailoverExhausted(c, failoverErr, true) @@ -322,7 +322,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. ) continue } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil) if c.Writer.Size() == writerSizeBeforeForward { h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed") } @@ -333,7 +333,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, nil) if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" { if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil { reqLog.Warn("grok_media.bind_video_request_account_failed", diff --git a/backend/internal/handler/openai_alpha_search.go b/backend/internal/handler/openai_alpha_search.go index 2daf44571d..1457c4de8e 100644 --- a/backend/internal/handler/openai_alpha_search.go +++ b/backend/internal/handler/openai_alpha_search.go @@ -118,7 +118,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { requestedModel, failedAccountIDs, service.OpenAIUpstreamTransportHTTPSSE, - service.OpenAIEndpointCapabilityChatCompletions, + service.OpenAIEndpointCapabilityAlphaSearch, false, false, false, @@ -164,7 +164,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { service.SetOpsLatencyMs(c, service.OpsResponseLatencyMsKey, time.Since(forwardStart).Milliseconds()) if err == nil { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestedModel), true, nil) if result != nil { h.recordAlphaSearchUsage(c, apiKey, account, subscription, channelMapping, requestedModel, body, result, subject.UserID) } @@ -173,7 +173,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { var failoverErr *service.UpstreamFailoverError if !errors.As(err, &failoverErr) { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestedModel), false, nil) if c.Writer.Size() == writerSizeBeforeForward { h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed") } @@ -181,7 +181,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestedModel), false, nil) if c.Writer.Size() != writerSizeBeforeForward { h.handleFailoverExhausted(c, failoverErr, true) return diff --git a/backend/internal/handler/openai_body_limit_failover_test.go b/backend/internal/handler/openai_body_limit_failover_test.go new file mode 100644 index 0000000000..edcb5a946b --- /dev/null +++ b/backend/internal/handler/openai_body_limit_failover_test.go @@ -0,0 +1,59 @@ +package handler + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestOpenAIBodyLimitFailoverExhausted_ReturnsRedactedJSON413(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + + (&OpenAIGatewayHandler{}).handleFailoverExhausted(c, bodyLimitFailoverTestError(), false) + + require.Equal(t, http.StatusRequestEntityTooLarge, rec.Code) + var envelope map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &envelope)) + errBody, ok := envelope["error"].(map[string]any) + require.True(t, ok) + require.Equal(t, "invalid_request_error", errBody["type"]) + require.Equal(t, "Request payload is too large", errBody["message"]) + require.NotContains(t, rec.Body.String(), "must-not-leak") +} + +func TestOpenAIBodyLimitFailoverExhausted_ReturnsRedactedResponsesSSE(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + + (&OpenAIGatewayHandler{}).handleFailoverExhausted(c, bodyLimitFailoverTestError(), true) + + body := rec.Body.String() + require.True(t, strings.HasPrefix(body, "event: response.failed\n")) + require.Contains(t, body, `"code":"invalid_request"`) + require.Contains(t, body, `"message":"Request payload is too large"`) + require.NotContains(t, body, "must-not-leak") +} + +func bodyLimitFailoverTestError() *service.UpstreamFailoverError { + return &service.UpstreamFailoverError{ + StatusCode: http.StatusRequestEntityTooLarge, + ResponseBody: []byte(`{"error":{"message":"proxy limit secret=must-not-leak"}}`), + Scope: service.GatewayFailureScopeAccount, + Reason: service.GatewayFailureReason("openai_request_body_too_large"), + NextAccountAction: service.NextAccountRetry, + ClientStatusCode: http.StatusRequestEntityTooLarge, + ClientMessage: "Request payload is too large", + } +} diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index 4e60b1a66a..84db86ce85 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -80,6 +80,10 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", invalidStreamFieldTypeMessage) return } + if service.IsGPTImageGenerationModel(reqModel) { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "This model is not supported on the Chat Completions endpoint") + return + } reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream)) @@ -252,7 +256,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { return } if failoverErr.ShouldReportAccountScheduleFailure() { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil) } if !failoverErr.ShouldRetryNextAccount() { h.handleFailoverExhausted(c, failoverErr, streamStarted) @@ -297,7 +301,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { ) continue } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil) upstreamErrorAlreadyCommunicated := openAIForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err) wroteFallback := false if !upstreamErrorAlreadyCommunicated { @@ -313,9 +317,9 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { } } if result != nil { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, result.FirstTokenMs) } else { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, nil) } userAgent := c.GetHeader("User-Agent") diff --git a/backend/internal/handler/openai_embeddings.go b/backend/internal/handler/openai_embeddings.go index 1d67565f11..473bfea5ef 100644 --- a/backend/internal/handler/openai_embeddings.go +++ b/backend/internal/handler/openai_embeddings.go @@ -193,7 +193,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { h.handleFailoverExhausted(c, failoverErr, true) return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil) if failoverClientGone(c) { reqLog.Info("openai_embeddings.failover_aborted_client_disconnected", zap.Int64("account_id", account.ID), @@ -217,7 +217,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { ) continue } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil) if c.Writer.Size() == writerSizeBeforeForward { h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed") } @@ -228,7 +228,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, nil) userAgent := c.GetHeader("User-Agent") clientIP := ip.GetClientIP(c) inboundEndpoint := GetInboundEndpoint(c) diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 90976a1fe2..b13fb01c55 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -42,6 +42,10 @@ type OpenAIGatewayHandler struct { const maxOpenAIFirstOutputTimeoutSwitches = 1 +func openAIForwardSucceededForScheduling(result *service.OpenAIForwardResult) bool { + return result.SucceededForScheduling() +} + func resolveOpenAIMessagesDispatchMappedModel(apiKey *service.APIKey, requestedModel string) string { if apiKey == nil || apiKey.Group == nil { return "" @@ -272,7 +276,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { return } - imageIntent := service.IsImageGenerationIntent("/v1/responses", reqModel, body) + imageIntent := service.IsImageGenerationIntentForPlatform("/v1/responses", reqModel, body, openAICompatibleRequestPlatform(apiKey)) if imageIntent && !service.GroupAllowsImageGeneration(apiKey.Group) { h.errorResponse(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage()) return @@ -482,7 +486,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { streamStarted = true } if failoverErr.ShouldReportAccountScheduleFailure() { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil) } if !failoverErr.ShouldRetryNextAccount() { h.handleFailoverExhausted(c, failoverErr, streamStarted) @@ -542,7 +546,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { reqLog.Warn("openai.upstream_failover_switching", failoverSwitchFields...) continue } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil) upstreamErrorAlreadyCommunicated := openAIForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err) wroteFallback := false if !upstreamErrorAlreadyCommunicated { @@ -567,9 +571,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { if account.Type == service.AccountTypeOAuth && !account.IsShadow() { h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(c.Request.Context(), account.ID, result.ResponseHeaders) } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), result.FirstTokenMs) } else { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), nil) } // 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context） @@ -1011,7 +1015,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { return } if failoverErr.ShouldReportAccountScheduleFailure() { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), false, nil) } if !failoverErr.ShouldRetryNextAccount() { h.handleAnthropicFailoverExhausted(c, failoverErr, streamStarted) @@ -1063,7 +1067,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { ) return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), false, nil) wroteFallback := h.ensureAnthropicErrorResponse(c, streamStarted) reqLog.Warn("openai_messages.forward_failed", zap.Int64("account_id", account.ID), @@ -1074,9 +1078,9 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { } } if result != nil { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, result.FirstTokenMs) } else { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, nil) } userAgent := c.GetHeader("User-Agent") @@ -1362,6 +1366,28 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { reqLog.Info("openai.websocket_ingress_started") clientIP := ip.GetClientIP(c) userAgent := strings.TrimSpace(c.GetHeader("User-Agent")) + ctx := c.Request.Context() + maxIngressConnections := 0 + if h.cfg != nil { + maxIngressConnections = h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey + } + ingressLease, ingressLeaseAcquired, ingressLeaseErr := h.concurrencyHelper.AcquireOpenAIWSIngressLease(ctx, apiKey.ID, maxIngressConnections) + if ingressLeaseErr != nil { + reqLog.Error("openai.websocket_ingress_lease_acquire_failed", zap.Error(ingressLeaseErr)) + h.errorResponse(c, http.StatusServiceUnavailable, "service_unavailable", "WebSocket ingress capacity is temporarily unavailable") + return + } + if !ingressLeaseAcquired { + reqLog.Info("openai.websocket_ingress_capacity_rejected", zap.Int("max_ingress_connections_per_api_key", maxIngressConnections)) + c.Header("Retry-After", "5") + h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Too many open WebSocket connections, please retry later") + return + } + if ingressLease != nil { + defer ingressLease.Release() + ctx = ingressLease.Context() + c.Request = c.Request.WithContext(ctx) + } wsConn, err := coderws.Accept(c.Writer, c.Request, &coderws.AcceptOptions{ CompressionMode: coderws.CompressionContextTakeover, @@ -1383,32 +1409,14 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { }() wsConn.SetReadLimit(service.ResolveOpenAIWSClientReadLimitBytes(h.cfg)) - ctx := c.Request.Context() - maxIngressConnections := 0 - if h.cfg != nil { - maxIngressConnections = h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey - } - ingressLease, ingressLeaseAcquired, ingressLeaseErr := h.concurrencyHelper.AcquireOpenAIWSIngressLease(ctx, apiKey.ID, maxIngressConnections) - if ingressLeaseErr != nil { - reqLog.Error("openai.websocket_ingress_lease_acquire_failed", zap.Error(ingressLeaseErr)) - closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to reserve websocket ingress capacity") - return - } - if !ingressLeaseAcquired { - reqLog.Info("openai.websocket_ingress_capacity_rejected", zap.Int("max_ingress_connections_per_api_key", maxIngressConnections)) - closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "too many open websocket connections, please retry later") - return - } - if ingressLease != nil { - defer ingressLease.Release() - ctx = ingressLease.Context() - c.Request = c.Request.WithContext(ctx) - } - firstMessageTimeout := service.ResolveOpenAIWSClientFirstMessageTimeout(h.cfg) - readCtx, cancel := context.WithTimeout(ctx, firstMessageTimeout) - msgType, firstMessage, err := wsConn.Read(readCtx) - cancel() + msgType, firstMessage, err := service.ReadOpenAIWSClientMessage( + ctx, + wsConn, + firstMessageTimeout, + coderws.StatusPolicyViolation, + "missing first response.create message", + ) if err != nil { if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) { reqLog.Warn("openai.websocket_ingress_lease_lost_before_first_message", zap.Error(err)) @@ -1463,7 +1471,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return } - imageIntent := service.IsImageGenerationIntent("/v1/responses", reqModel, firstMessage) + imageIntent := service.IsImageGenerationIntentForPlatform("/v1/responses", reqModel, firstMessage, openAICompatibleRequestPlatform(apiKey)) if imageIntent && !service.GroupAllowsImageGeneration(apiKey.Group) { closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, service.ImageGenerationPermissionMessage()) return @@ -1557,7 +1565,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return false } if failoverErr.ShouldReportAccountScheduleFailure() { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil) } releaseAccountSlot() if !failoverErr.ShouldRetryNextAccount() { @@ -1777,7 +1785,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { if account.Type == service.AccountTypeOAuth && !account.IsShadow() { h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(ctx, account.ID, result.ResponseHeaders) } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), result.FirstTokenMs) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) @@ -1857,7 +1865,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil) closeStatus, closeReason := summarizeWSCloseErrorForLog(err) proxyFailedFields := []zap.Field{ zap.Int64("account_id", account.ID), @@ -2081,6 +2089,17 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE h.handleFailoverExhaustedSimple(c, http.StatusBadGateway, streamStarted) return } + if failoverErr.IsOpenAIRequestBodyTooLarge() { + service.SetOpsUpstreamError(c, http.StatusRequestEntityTooLarge, service.OpenAIRequestBodyTooLargeClientMessage, "") + h.handleStreamingAwareError( + c, + http.StatusRequestEntityTooLarge, + "invalid_request_error", + service.OpenAIRequestBodyTooLargeClientMessage, + streamStarted, + ) + return + } copyFailoverRetryAfter(c, failoverErr.ResponseHeaders) if failoverErr.IsCredentialFailure() { status, message := credentialFailoverClientResponse(failoverErr) diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index d855f364a7..931f3b3f65 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -96,6 +96,19 @@ func TestOpenAIHandleStreamingAwareError_JSONEscaping(t *testing.T) { } } +func TestOpenAIForwardSucceededForScheduling(t *testing.T) { + require.True(t, openAIForwardSucceededForScheduling(nil)) + require.True(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{})) + require.True(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{ + OpenAIWSMode: true, + UpstreamTerminalEvent: "response.completed", + })) + require.False(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{ + OpenAIWSMode: true, + UpstreamTerminalEvent: "response.failed", + })) +} + func TestResolveOpenAIMessagesMetadataSession_DoesNotDerivePromptCacheKey(t *testing.T) { body := []byte(`{"model":"claude-sonnet-4-5","metadata":{"user_id":"claude-code-session"},"messages":[{"role":"user","content":"hello"}]}`) @@ -221,6 +234,43 @@ func TestOpenAIEnsureForwardErrorResponse_ResponsesRouteAfterWrittenEmitsRespons assert.Contains(t, body, "Upstream request failed") } +func TestOpenAIEnsureForwardErrorResponse_AfterDeltaAppendsSingleValidResponseFailed(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil) + + delta := `{"type":"response.output_text.delta","delta":"ok","sequence_number":1}` + _, err := c.Writer.WriteString("event: response.output_text.delta\ndata: " + delta + "\n\n") + require.NoError(t, err) + + h := &OpenAIGatewayHandler{} + require.True(t, h.ensureForwardErrorResponse(c, true)) + + frames := strings.Split(strings.TrimSuffix(w.Body.String(), "\n\n"), "\n\n") + require.Len(t, frames, 2) + errorEvents := 0 + for _, frame := range frames { + lines := strings.Split(frame, "\n") + require.Len(t, lines, 2) + require.True(t, strings.HasPrefix(lines[0], "event: ")) + require.True(t, strings.HasPrefix(lines[1], "data: ")) + + eventType := strings.TrimPrefix(lines[0], "event: ") + data := strings.TrimPrefix(lines[1], "data: ") + require.True(t, json.Valid([]byte(data)), "each downstream SSE frame must contain valid JSON") + var event struct { + Type string `json:"type"` + } + require.NoError(t, json.Unmarshal([]byte(data), &event)) + require.Equal(t, eventType, event.Type) + if eventType == "response.failed" { + errorEvents++ + } + } + require.Equal(t, 1, errorEvents) +} + func TestOpenAIEnsureForwardErrorResponse_ImageJSONKeepaliveWritesSingleJSONFallback(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() @@ -820,17 +870,36 @@ func TestOpenAIResponsesWebSocket_IngressCapacityRejected(t *testing.T) { defer wsServer.Close() dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) - clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil) + clientConn, response, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil) cancelDial() - require.NoError(t, err) - defer func() { _ = clientConn.CloseNow() }() + require.Error(t, err) + require.Nil(t, clientConn) + require.NotNil(t, response) + require.Equal(t, http.StatusTooManyRequests, response.StatusCode) + _ = response.Body.Close() +} - readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) - _, _, err = clientConn.Read(readCtx) - cancelRead() - var closeErr coderws.CloseError - require.ErrorAs(t, err, &closeErr) - require.Equal(t, coderws.StatusTryAgainLater, closeErr.Code) +func TestOpenAIResponsesWebSocket_IngressLeaseBackendUnavailableBeforeUpgrade(t *testing.T) { + gin.SetMode(gin.TestMode) + cache := &concurrencyCacheMock{ + acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) { + return false, errors.New("redis unavailable") + }, + } + h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache) + h.cfg = &config.Config{} + h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1 + wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1}) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, response, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil) + cancelDial() + require.Error(t, err) + require.Nil(t, clientConn) + require.NotNil(t, response) + require.Equal(t, http.StatusServiceUnavailable, response.StatusCode) + _ = response.Body.Close() } func TestOpenAIResponsesWebSocket_FirstMessageTimeoutUsesConfig(t *testing.T) { @@ -902,6 +971,33 @@ func TestOpenAIResponsesWebSocket_IngressLeaseReleasedOnEarlyReturn(t *testing.T }, time.Second, 10*time.Millisecond) } +func TestOpenAIResponsesWebSocket_IngressLeaseReleasedWhenUpgradeFails(t *testing.T) { + gin.SetMode(gin.TestMode) + cache := &concurrencyCacheMock{ + acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) { + return true, nil + }, + } + h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache) + h.cfg = &config.Config{} + h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1 + wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1}) + defer wsServer.Close() + + req, err := http.NewRequest(http.MethodGet, wsServer.URL+"/openai/v1/responses", nil) + require.NoError(t, err) + req.Header.Set("Upgrade", "websocket") + req.Header.Set("Connection", "Upgrade") + req.Header.Set("Sec-WebSocket-Version", "13") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + _ = resp.Body.Close() + require.NotEqual(t, http.StatusSwitchingProtocols, resp.StatusCode) + require.Eventually(t, func() bool { + return atomic.LoadInt32(&cache.releaseIngressCalled) == 1 + }, time.Second, 10*time.Millisecond) +} + func TestOpenAIResponsesWebSocket_RejectsMessageIDAsPreviousResponseID(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/handler/openai_grok_image_intent_gate_test.go b/backend/internal/handler/openai_grok_image_intent_gate_test.go new file mode 100644 index 0000000000..3c20160c05 --- /dev/null +++ b/backend/internal/handler/openai_grok_image_intent_gate_test.go @@ -0,0 +1,101 @@ +package handler + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestOpenAIGatewayHandlerResponses_GrokPassiveImageToolDeclarationBypassesPermissionGate(t *testing.T) { + body := `{"model":"grok-4.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}],"tool_choice":"auto","input":"write code"}` + rec := runOpenAIResponsesImagePermissionGateTest(t, service.PlatformGrok, body) + + require.NotEqual(t, http.StatusForbidden, rec.Code) + require.NotContains(t, rec.Body.String(), service.ImageGenerationPermissionMessage()) +} + +func TestOpenAIGatewayHandlerResponses_GrokResponsesLiteImageToolDeclarationBypassesPermissionGate(t *testing.T) { + body := `{"model":"grok-4.5","tool_choice":"auto","input":[{"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]},{"type":"message","role":"user","content":"write code"}]}` + rec := runOpenAIResponsesImagePermissionGateTest(t, service.PlatformGrok, body) + + require.NotEqual(t, http.StatusForbidden, rec.Code) + require.NotContains(t, rec.Body.String(), service.ImageGenerationPermissionMessage()) +} + +func TestOpenAIGatewayHandlerResponses_ImagePermissionHardSignalsStillRejected(t *testing.T) { + passiveNamespace := `{"model":"gpt-5.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}],"tool_choice":"auto","input":"write code"}` + tests := []struct { + name string + platform string + body string + }{ + { + name: "OpenAI keeps declaration semantics", + platform: service.PlatformOpenAI, + body: passiveNamespace, + }, + { + name: "Grok native image_generation declaration", + platform: service.PlatformGrok, + body: `{"model":"grok-4.5","tools":[{"type":"image_generation"}],"input":"draw"}`, + }, + { + name: "Grok explicit image_gen tool choice", + platform: service.PlatformGrok, + body: `{"model":"grok-4.5","tools":[{"type":"namespace","name":"image_gen"}],"tool_choice":{"type":"namespace","name":"image_gen"},"input":"draw"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rec := runOpenAIResponsesImagePermissionGateTest(t, tt.platform, tt.body) + + require.Equal(t, http.StatusForbidden, rec.Code) + require.Contains(t, rec.Body.String(), service.ImageGenerationPermissionMessage()) + }) + } +} + +func runOpenAIResponsesImagePermissionGateTest(t *testing.T, platform string, body string) *httptest.ResponseRecorder { + t.Helper() + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + groupID := int64(6301) + userID := int64(6302) + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ + ID: 6303, + GroupID: &groupID, + Group: &service.Group{ + ID: groupID, + Platform: platform, + AllowImageGeneration: false, + }, + User: &service.User{ID: userID, Status: service.StatusActive}, + }) + c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: userID, Concurrency: 1}) + + h := &OpenAIGatewayHandler{ + gatewayService: &service.OpenAIGatewayService{}, + billingCacheService: service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, &config.Config{RunMode: config.RunModeSimple}, nil), + apiKeyService: &service.APIKeyService{}, + concurrencyHelper: &ConcurrencyHelper{concurrencyService: service.NewConcurrencyService( + &helperConcurrencyCacheStub{userSeq: []bool{true}}, + )}, + cfg: &config.Config{}, + imageLimiter: &imageConcurrencyLimiter{}, + } + + h.Responses(c) + return rec +} diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index f5709225d1..d585c159a8 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -253,7 +253,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { var imageUpstreamErr *service.OpenAIImagesUpstreamError if errors.As(err, &imageUpstreamErr) { retryableServerError := service.IsOpenAIImagesRetryableUpstreamError(imageUpstreamErr) - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, !retryableServerError, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), !retryableServerError, nil) logEvent := "openai.images.upstream_user_error" if retryableServerError { logEvent = "openai.images.upstream_server_error_after_flush" @@ -269,7 +269,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { } var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil) if service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward { reqLog.Warn("openai.images.upstream_failover_skipped_after_flush", zap.Int64("account_id", account.ID), @@ -323,7 +323,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { ) continue } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil) upstreamErrorAlreadyCommunicated := openAIForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err) wroteFallback := false if !upstreamErrorAlreadyCommunicated { @@ -348,9 +348,9 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { if account.Type == service.AccountTypeOAuth && !account.IsShadow() { h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(c.Request.Context(), account.ID, result.ResponseHeaders) } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, result.FirstTokenMs) } else { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, nil) } userAgent := c.GetHeader("User-Agent") diff --git a/backend/internal/pkg/tlsfingerprint/dialer_integration_test.go b/backend/internal/pkg/tlsfingerprint/dialer_integration_test.go index 38cddd0d9f..9912ddb3bb 100644 --- a/backend/internal/pkg/tlsfingerprint/dialer_integration_test.go +++ b/backend/internal/pkg/tlsfingerprint/dialer_integration_test.go @@ -28,6 +28,9 @@ func skipIfExternalServiceUnavailable(t *testing.T, err error) { if strings.Contains(errStr, "certificate has expired") || strings.Contains(errStr, "certificate is not yet valid") || strings.Contains(errStr, "connection refused") || + strings.Contains(errStr, "connection reset") || + strings.Contains(errStr, "broken pipe") || + strings.Contains(errStr, "EOF") || strings.Contains(errStr, "no such host") || strings.Contains(errStr, "network is unreachable") || strings.Contains(errStr, "timeout") || @@ -74,14 +77,15 @@ func TestJA3Fingerprint(t *testing.T) { defer func() { _ = resp.Body.Close() }() body, err := io.ReadAll(resp.Body) - if err != nil { - t.Fatalf("failed to read response: %v", err) + skipIfExternalServiceUnavailable(t, err) + if resp.StatusCode != http.StatusOK { + t.Skipf("skipping test: external service unhealthy: status %d", resp.StatusCode) } var fpResp FingerprintResponse if err := json.Unmarshal(body, &fpResp); err != nil { t.Logf("Response body: %s", string(body)) - t.Fatalf("failed to parse fingerprint response: %v", err) + t.Skipf("skipping test: external service returned unparsable response: %v", err) } t.Logf("JA3: %s", fpResp.TLS.JA3) @@ -207,16 +211,15 @@ func fetchFingerprint(t *testing.T, profile *Profile) *TLSInfo { defer func() { _ = resp.Body.Close() }() body, err := io.ReadAll(resp.Body) - if err != nil { - t.Fatalf("failed to read response: %v", err) - return nil + skipIfExternalServiceUnavailable(t, err) + if resp.StatusCode != http.StatusOK { + t.Skipf("skipping test: external service unhealthy: status %d", resp.StatusCode) } var fpResp FingerprintResponse if err := json.Unmarshal(body, &fpResp); err != nil { t.Logf("Response body: %s", string(body)) - t.Fatalf("failed to parse fingerprint response: %v", err) - return nil + t.Skipf("skipping test: external service returned unparsable response: %v", err) } return &fpResp.TLS diff --git a/backend/internal/repository/channel_repo_pricing.go b/backend/internal/repository/channel_repo_pricing.go index 6dcf3c9175..995621139e 100644 --- a/backend/internal/repository/channel_repo_pricing.go +++ b/backend/internal/repository/channel_repo_pricing.go @@ -16,7 +16,7 @@ import ( func (r *channelRepository) ListModelPricing(ctx context.Context, channelID int64) ([]service.ChannelModelPricing, error) { rows, err := r.db.QueryContext(ctx, - `SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_output_price, per_request_price, created_at, updated_at + `SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price, created_at, updated_at FROM channel_model_pricing WHERE channel_id = $1 ORDER BY id`, channelID, ) if err != nil { @@ -57,10 +57,10 @@ func (r *channelRepository) UpdateModelPricing(ctx context.Context, pricing *ser } result, err := r.db.ExecContext(ctx, `UPDATE channel_model_pricing - SET models = $1, billing_mode = $2, input_price = $3, output_price = $4, cache_write_price = $5, cache_read_price = $6, image_output_price = $7, per_request_price = $8, platform = $9, updated_at = NOW() - WHERE id = $10`, + SET models = $1, billing_mode = $2, input_price = $3, output_price = $4, cache_write_price = $5, cache_read_price = $6, image_input_price = $7, image_output_price = $8, per_request_price = $9, platform = $10, updated_at = NOW() + WHERE id = $11`, modelsJSON, billingMode, pricing.InputPrice, pricing.OutputPrice, pricing.CacheWritePrice, pricing.CacheReadPrice, - pricing.ImageOutputPrice, pricing.PerRequestPrice, pricing.Platform, pricing.ID, + pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice, pricing.Platform, pricing.ID, ) if err != nil { return fmt.Errorf("update model pricing: %w", err) @@ -91,7 +91,7 @@ func (r *channelRepository) ReplaceModelPricing(ctx context.Context, channelID i // batchLoadModelPricing 批量加载多个渠道的模型定价（含区间） func (r *channelRepository) batchLoadModelPricing(ctx context.Context, channelIDs []int64) (map[int64][]service.ChannelModelPricing, error) { rows, err := r.db.QueryContext(ctx, - `SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_output_price, per_request_price, created_at, updated_at + `SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price, created_at, updated_at FROM channel_model_pricing WHERE channel_id = ANY($1) ORDER BY channel_id, id`, pq.Array(channelIDs), ) @@ -172,7 +172,7 @@ func scanModelPricingRows(rows *sql.Rows) ([]service.ChannelModelPricing, []int6 if err := rows.Scan( &p.ID, &p.ChannelID, &p.Platform, &modelsJSON, &p.BillingMode, &p.InputPrice, &p.OutputPrice, &p.CacheWritePrice, &p.CacheReadPrice, - &p.ImageOutputPrice, &p.PerRequestPrice, &p.CreatedAt, &p.UpdatedAt, + &p.ImageInputPrice, &p.ImageOutputPrice, &p.PerRequestPrice, &p.CreatedAt, &p.UpdatedAt, ); err != nil { return nil, nil, fmt.Errorf("scan model pricing: %w", err) } @@ -229,11 +229,11 @@ func createModelPricingExec(ctx context.Context, exec dbExec, pricing *service.C platform = "anthropic" } err = exec.QueryRowContext(ctx, - `INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_output_price, per_request_price) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING id, created_at, updated_at`, + `INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) RETURNING id, created_at, updated_at`, pricing.ChannelID, platform, modelsJSON, billingMode, pricing.InputPrice, pricing.OutputPrice, pricing.CacheWritePrice, pricing.CacheReadPrice, - pricing.ImageOutputPrice, pricing.PerRequestPrice, + pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice, ).Scan(&pricing.ID, &pricing.CreatedAt, &pricing.UpdatedAt) if err != nil { return fmt.Errorf("insert model pricing: %w", err) diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index 5037f30a8f..90a7711e4f 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -395,6 +395,43 @@ func (c *schedulerCache) SetSnapshot(ctx context.Context, bucket service.Schedul return c.activateSnapshotVersion(ctx, bucket, token, version) } +// SetSnapshotAndReturnAccountIDs 完整发布快照，并返回 writeAccounts 实际接受的有序账号 ID。 +// 该可选能力只供同一重建批次复用，返回前仍会完成版本激活与 fencing 校验。 +func (c *schedulerCache) SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accounts []service.Account) ([]int64, error) { + if !token.ValidFor(bucket) { + return nil, fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String()) + } + // 分配版本与激活指针是两个 fencing 边界；中间写入的数据只有通过第二次校验才能发布。 + version, err := c.allocateSnapshotVersion(ctx, bucket, token) + if err != nil { + return nil, err + } + accountIDs, err := c.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, accounts) + if err != nil { + return nil, err + } + if err := c.activateSnapshotVersion(ctx, bucket, token, version); err != nil { + return nil, err + } + return accountIDs, nil +} + +// SetSnapshotByAccountIDs 复用同批次首次完整写入后得到的账号成员。 +// 每个桶仍独立分配版本、写入有序集合并执行激活 fencing，只省略重复的账号 JSON 与全局键写入。 +func (c *schedulerCache) SetSnapshotByAccountIDs(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accountIDs []int64) error { + if !token.ValidFor(bucket) { + return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String()) + } + version, err := c.allocateSnapshotVersion(ctx, bucket, token) + if err != nil { + return err + } + if err := c.writeSnapshotAccountIDs(ctx, bucket, version, accountIDs); err != nil { + return err + } + return c.activateSnapshotVersion(ctx, bucket, token, version) +} + func (c *schedulerCache) allocateSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken) (string, error) { result, err := allocateSnapshotVersionScript.Run(ctx, c.rdb, []string{ schedulerBucketKey(schedulerEpochPrefix, bucket), @@ -411,35 +448,74 @@ func (c *schedulerCache) allocateSnapshotVersion(ctx context.Context, bucket ser } func (c *schedulerCache) writeSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) error { - snapshotKey := schedulerSnapshotKey(bucket, version) cacheableAccounts, err := c.writeAccounts(ctx, accounts) if err != nil { return err } + return c.writeSnapshotAccounts(ctx, bucket, version, cacheableAccounts) +} - if len(cacheableAccounts) > 0 { - // 使用序号作为 score，保持数据库返回的排序语义。 - members := make([]redis.Z, 0, len(cacheableAccounts)) - for idx, account := range cacheableAccounts { - members = append(members, redis.Z{ - Score: float64(idx), - Member: strconv.FormatInt(account.ID, 10), - }) - } - pipe := c.rdb.Pipeline() - for start := 0; start < len(members); start += c.writeChunkSize { - end := start + c.writeChunkSize - if end > len(members) { - end = len(members) - } - pipe.ZAdd(ctx, snapshotKey, members[start:end]...) - } - if _, err := pipe.Exec(ctx); err != nil { - return err - } +func (c *schedulerCache) writeSnapshotVersionAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) ([]int64, error) { + accountIDs, err := c.writeAccountIDs(ctx, accounts) + if err != nil { + return nil, err } + if err := c.writeSnapshotAccountIDs(ctx, bucket, version, accountIDs); err != nil { + return nil, err + } + return accountIDs, nil +} - return nil +func (c *schedulerCache) writeSnapshotAccounts(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) error { + if len(accounts) == 0 { + return nil + } + members := make([]redis.Z, 0, len(accounts)) + for idx, account := range accounts { + members = append(members, redis.Z{ + Score: float64(idx), + Member: strconv.FormatInt(account.ID, 10), + }) + } + return c.writeSnapshotMembers(ctx, bucket, version, members) +} + +func (c *schedulerCache) writeSnapshotAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accountIDs []int64) error { + members := schedulerSnapshotMembers(accountIDs) + return c.writeSnapshotMembers(ctx, bucket, version, members) +} + +func schedulerSnapshotMembers(accountIDs []int64) []redis.Z { + if len(accountIDs) == 0 { + return nil + } + // 使用序号作为 score，保持数据库返回的排序语义；重复 ID 继续交由 Redis ZADD + // 按最后一个 score 覆盖，与直接从账号切片构造成员时的行为一致。 + members := make([]redis.Z, 0, len(accountIDs)) + for idx, accountID := range accountIDs { + members = append(members, redis.Z{ + Score: float64(idx), + Member: strconv.FormatInt(accountID, 10), + }) + } + return members +} + +func (c *schedulerCache) writeSnapshotMembers(ctx context.Context, bucket service.SchedulerBucket, version string, members []redis.Z) error { + if len(members) == 0 { + return nil + } + snapshotKey := schedulerSnapshotKey(bucket, version) + pipe := c.rdb.Pipeline() + for start := 0; start < len(members); start += c.writeChunkSize { + end := start + c.writeChunkSize + if end > len(members) { + end = len(members) + } + pipe.ZAdd(ctx, snapshotKey, members[start:end]...) + } + _, err := pipe.Exec(ctx) + return err } func (c *schedulerCache) activateSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, version string) error { @@ -644,12 +720,28 @@ func decodeCachedAccount(val any) (*service.Account, error) { } func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.Account) ([]service.Account, error) { + cacheableAccounts, _, err := c.writeAccountPayloads(ctx, accounts, false) + return cacheableAccounts, err +} + +func (c *schedulerCache) writeAccountIDs(ctx context.Context, accounts []service.Account) ([]int64, error) { + _, accountIDs, err := c.writeAccountPayloads(ctx, accounts, true) + return accountIDs, err +} + +func (c *schedulerCache) writeAccountPayloads(ctx context.Context, accounts []service.Account, collectIDs bool) ([]service.Account, []int64, error) { if len(accounts) == 0 { - return nil, nil + return nil, nil, nil } pipe := c.rdb.Pipeline() - cacheableAccounts := make([]service.Account, 0, len(accounts)) + var cacheableAccounts []service.Account + var accountIDs []int64 + if collectIDs { + accountIDs = make([]int64, 0, len(accounts)) + } else { + cacheableAccounts = make([]service.Account, 0, len(accounts)) + } pending := 0 flush := func() error { if pending == 0 { @@ -676,19 +768,24 @@ func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.A id := strconv.FormatInt(account.ID, 10) pipe.Set(ctx, schedulerAccountKey(id), fullPayload, 0) pipe.Set(ctx, schedulerAccountMetaKey(id), metaPayload, 0) - cacheableAccounts = append(cacheableAccounts, account) + // 复用路径只保留有序 ID，避免先物化完整账号切片再做第二次扫描。 + if collectIDs { + accountIDs = append(accountIDs, account.ID) + } else { + cacheableAccounts = append(cacheableAccounts, account) + } pending++ if pending >= c.writeChunkSize { if err := flush(); err != nil { - return nil, err + return nil, nil, err } } } if err := flush(); err != nil { - return nil, err + return nil, nil, err } - return cacheableAccounts, nil + return cacheableAccounts, accountIDs, nil } func marshalSchedulerCacheAccount(account service.Account) ([]byte, []byte, error) { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index f933603c47..230dc59903 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -5,6 +5,9 @@ package repository import ( "context" "encoding/hex" + "encoding/json" + "fmt" + "strconv" "strings" "testing" "time" @@ -82,6 +85,170 @@ func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) { require.Nil(t, cached) } +func TestSchedulerCacheSnapshotAccountIDReusePreservesPayloadAndMembers(t *testing.T) { + ctx := context.Background() + cache, _ := newSchedulerCacheUnitWithRedis(t) + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + validOne := service.Account{ + ID: 701, + Name: "first", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Credentials: map[string]any{"model_mapping": map[string]any{"z": "last", "a": "first"}}, + Extra: map[string]any{"mixed_scheduling": true}, + GroupIDs: []int64{17}, + } + validTwo := service.Account{ID: 702, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey} + invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime} + accounts := []service.Account{validOne, invalid, validTwo, validOne} + + single := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + singleToken, err := cache.CaptureBucketWriteToken(ctx, single) + require.NoError(t, err) + accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts) + require.NoError(t, err) + require.Equal(t, []int64{701, 702, 701}, accountIDs, "应保留可编码账号的原顺序和重复项") + + wantFull, err := json.Marshal(validOne) + require.NoError(t, err) + wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(validOne)) + require.NoError(t, err) + fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes() + require.NoError(t, err) + metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes() + require.NoError(t, err) + require.Equal(t, wantFull, fullBefore) + require.Equal(t, wantMeta, metaBefore) + + forced := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced} + forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs)) + + fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes() + require.NoError(t, err) + metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes() + require.NoError(t, err) + require.Equal(t, fullBefore, fullAfter, "ID-only 路径不得重写完整账号键") + require.Equal(t, metaBefore, metaAfter, "ID-only 路径不得重写调度元数据键") + + for _, bucket := range []service.SchedulerBucket{single, forced} { + version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result() + require.NoError(t, err) + members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result() + require.NoError(t, err) + require.Equal(t, []string{"702", "701"}, members, bucket.String()) + } + missing, err := cache.GetAccount(ctx, invalid.ID) + require.NoError(t, err) + require.Nil(t, missing) +} + +func TestSchedulerCacheSnapshotAccountIDReuseKeepsEmptySnapshotSemantics(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + accounts := []service.Account{{ID: 811, Platform: service.PlatformOpenAI, ExpiresAt: &invalidTime}} + + single := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + singleToken, err := cache.CaptureBucketWriteToken(ctx, single) + require.NoError(t, err) + accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts) + require.NoError(t, err) + require.Empty(t, accountIDs) + + forced := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced} + forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs)) + + for _, bucket := range []service.SchedulerBucket{single, forced} { + ready, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerReadyPrefix, bucket)).Result() + require.NoError(t, err) + require.Equal(t, "1", ready) + snapshot, hit, err := cache.GetSnapshot(ctx, bucket) + require.NoError(t, err) + require.False(t, hit, bucket.String()) + require.Nil(t, snapshot) + } +} + +func TestSchedulerCacheSetSnapshotByAccountIDsKeepsFencing(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + bucket := service.SchedulerBucket{GroupID: 19, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced} + + err := cache.SetSnapshotByAccountIDs(ctx, bucket, service.SchedulerBucketWriteToken{}, []int64{901}) + require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced) + _, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Result() + require.ErrorIs(t, err, redis.Nil) + + token, err := cache.CaptureBucketWriteToken(ctx, bucket) + require.NoError(t, err) + require.NoError(t, cache.RetireBucket(ctx, bucket)) + err = cache.SetSnapshotByAccountIDs(ctx, bucket, token, []int64{901}) + require.ErrorIs(t, err, service.ErrSchedulerBucketRetired) +} + +func TestSchedulerCacheSetSnapshotByAccountIDsDoesNotResurrectDeletedAccount(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + account := service.Account{ID: 902, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth} + single := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + singleToken, err := cache.CaptureBucketWriteToken(ctx, single) + require.NoError(t, err) + accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, []service.Account{account}) + require.NoError(t, err) + require.Equal(t, []int64{account.ID}, accountIDs) + require.NoError(t, cache.DeleteAccount(ctx, account.ID)) + + forced := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced} + forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs)) + + full, err := cache.GetAccount(ctx, account.ID) + require.NoError(t, err) + require.Nil(t, full, "ID-only 发布不得复活已删除的完整账号键") + snapshot, hit, err := cache.GetSnapshot(ctx, forced) + require.NoError(t, err) + require.False(t, hit, "元数据缺失时必须安全回源，而不是返回残缺快照") + require.Nil(t, snapshot) +} + +func TestMarshalSchedulerCacheAccountKeepsEncodingJSONWireFormat(t *testing.T) { + cases := []struct { + name string + account service.Account + }{ + {name: "nil collections", account: service.Account{ID: 801}}, + {name: "empty collections", account: service.Account{ + ID: 802, + Credentials: map[string]any{}, + Extra: map[string]any{}, + GroupIDs: []int64{}, + Groups: []*service.Group{}, + }}, + {name: "nested maps and escaping", account: service.Account{ + ID: 803, + Credentials: map[string]any{"model_mapping": map[string]any{"z": "", "a": "&first"}}, + Extra: map[string]any{"mixed_scheduling": true}, + }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + full, meta, err := marshalSchedulerCacheAccount(tc.account) + require.NoError(t, err) + wantFull, err := json.Marshal(tc.account) + require.NoError(t, err) + wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(tc.account)) + require.NoError(t, err) + require.Equal(t, wantFull, full) + require.Equal(t, wantMeta, meta) + }) + } +} + func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) { account := service.Account{ ID: 42, @@ -611,3 +778,126 @@ func TestSchedulerCacheGroupLifecycleLeaseRejectsInvalidInput(t *testing.T) { require.NoError(t, err) require.Zero(t, keys) } + +var schedulerCachePayloadBenchmarkSink int + +func BenchmarkSchedulerCacheAccountPayloadReuse(b *testing.B) { + for _, size := range []int{1, 100, 10_000} { + accounts := schedulerCacheBenchmarkAccounts(size) + b.Run(fmt.Sprintf("pair_baseline_%d_accounts", size), func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + first, err := benchmarkSchedulerLegacySnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + second, err := benchmarkSchedulerLegacySnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + schedulerCachePayloadBenchmarkSink = first + second + } + }) + b.Run(fmt.Sprintf("pair_reuse_%d_accounts", size), func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + // 第二个桶仍构造成员，只跳过账号 JSON 与全局账号键。 + total += len(schedulerSnapshotMembers(ids)) + schedulerCachePayloadBenchmarkSink = total + } + }) + b.Run(fmt.Sprintf("first_baseline_%d_accounts", size), func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + total, err := benchmarkSchedulerLegacySnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + schedulerCachePayloadBenchmarkSink = total + } + }) + b.Run(fmt.Sprintf("first_reuse_%d_accounts", size), func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + total += len(ids) + schedulerCachePayloadBenchmarkSink = total + } + }) + } +} + +func benchmarkSchedulerLegacySnapshotPayload(accounts []service.Account) (int, error) { + cacheable := make([]service.Account, 0, len(accounts)) + total := 0 + for _, account := range accounts { + full, meta, err := marshalSchedulerCacheAccount(account) + if err != nil { + continue + } + total += len(full) + len(meta) + cacheable = append(cacheable, account) + } + members := make([]redis.Z, 0, len(cacheable)) + for idx, account := range cacheable { + members = append(members, redis.Z{Score: float64(idx), Member: strconv.FormatInt(account.ID, 10)}) + } + return total + len(members), nil +} + +func benchmarkSchedulerReusableSnapshotPayload(accounts []service.Account) ([]int64, int, error) { + accountIDs := make([]int64, 0, len(accounts)) + total := 0 + for _, account := range accounts { + full, meta, err := marshalSchedulerCacheAccount(account) + if err != nil { + continue + } + total += len(full) + len(meta) + accountIDs = append(accountIDs, account.ID) + } + total += len(schedulerSnapshotMembers(accountIDs)) + return accountIDs, total, nil +} + +func schedulerCacheBenchmarkAccounts(size int) []service.Account { + largeValue := strings.Repeat("x", 4096) + credentials := map[string]any{ + "api_key": "benchmark-key", + "model_mapping": map[string]any{"z-model": "z-target", "a-model": "a-target"}, + "large_value": largeValue, + } + extra := map[string]any{ + "mixed_scheduling": true, + "model_rate_limits": map[string]any{ + "z-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"}, + "a-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"}, + }, + "large_value": largeValue, + } + accounts := make([]service.Account, size) + for i := range accounts { + id := int64(i + 1) + accounts[i] = service.Account{ + ID: id, + Name: "benchmark-account", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Credentials: credentials, + Extra: extra, + GroupIDs: []int64{7, 9}, + AccountGroups: []service.AccountGroup{ + {AccountID: id, GroupID: 7, Priority: 1}, + {AccountID: id, GroupID: 9, Priority: 2}, + }, + } + } + return accounts +} diff --git a/backend/internal/repository/usage_log_repo_insert.go b/backend/internal/repository/usage_log_repo_insert.go index ec09b308a0..54d24b5957 100644 --- a/backend/internal/repository/usage_log_repo_insert.go +++ b/backend/internal/repository/usage_log_repo_insert.go @@ -41,6 +41,8 @@ var usageLogInsertArgTypes = [...]string{ "integer", // cache_creation_1h_tokens "integer", // image_output_tokens "numeric", // image_output_cost + "integer", // image_input_tokens + "numeric", // image_input_cost "numeric", // input_cost "numeric", // output_cost "numeric", // cache_creation_cost @@ -234,6 +236,8 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -277,7 +281,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54 + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56 ) ON CONFLICT (request_id, api_key_id) DO NOTHING RETURNING id, created_at @@ -686,6 +690,8 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -725,7 +731,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage created_at ) AS (VALUES `) - args := make([]any, 0, len(keys)*54) + args := make([]any, 0, len(keys)*56) argPos := 1 for idx, key := range keys { if idx > 0 { @@ -771,6 +777,8 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -827,6 +835,8 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -923,6 +933,8 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -962,7 +974,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( created_at ) AS (VALUES `) - args := make([]any, 0, len(preparedList)*54) + args := make([]any, 0, len(preparedList)*56) argPos := 1 for idx, prepared := range preparedList { if idx > 0 { @@ -1005,6 +1017,8 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -1061,6 +1075,8 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -1125,6 +1141,8 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -1168,7 +1186,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54 + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56 ) ON CONFLICT (request_id, api_key_id) DO NOTHING `, prepared.args...) @@ -1243,6 +1261,8 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared { log.CacheCreation1hTokens, log.ImageOutputTokens, log.ImageOutputCost, + log.ImageInputTokens, + log.ImageInputCost, log.InputCost, log.OutputCost, log.CacheCreationCost, diff --git a/backend/internal/repository/usage_log_repo_query.go b/backend/internal/repository/usage_log_repo_query.go index 1fdedd8665..fb0e03ed60 100644 --- a/backend/internal/repository/usage_log_repo_query.go +++ b/backend/internal/repository/usage_log_repo_query.go @@ -19,7 +19,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/service" ) -const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" +const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, image_input_tokens, image_input_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" func (r *usageLogRepository) GetByID(ctx context.Context, id int64) (log *service.UsageLog, err error) { query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE id = $1" @@ -443,6 +443,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e cacheCreation1h int imageOutputTokens int imageOutputCost float64 + imageInputTokens int + imageInputCost float64 inputCost float64 outputCost float64 cacheCreationCost float64 @@ -501,6 +503,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e &cacheCreation1h, &imageOutputTokens, &imageOutputCost, + &imageInputTokens, + &imageInputCost, &inputCost, &outputCost, &cacheCreationCost, @@ -557,6 +561,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e CacheCreation1hTokens: cacheCreation1h, ImageOutputTokens: imageOutputTokens, ImageOutputCost: imageOutputCost, + ImageInputTokens: imageInputTokens, + ImageInputCost: imageInputCost, InputCost: inputCost, OutputCost: outputCost, CacheCreationCost: cacheCreationCost, diff --git a/backend/internal/repository/usage_log_repo_request_type_test.go b/backend/internal/repository/usage_log_repo_request_type_test.go index 052c319183..f039029608 100644 --- a/backend/internal/repository/usage_log_repo_request_type_test.go +++ b/backend/internal/repository/usage_log_repo_request_type_test.go @@ -58,6 +58,8 @@ func TestUsageLogRepositoryCreateSyncRequestTypeAndLegacyFields(t *testing.T) { log.CacheCreation1hTokens, log.ImageOutputTokens, log.ImageOutputCost, + log.ImageInputTokens, + log.ImageInputCost, log.InputCost, log.OutputCost, log.CacheCreationCost, @@ -145,6 +147,8 @@ func TestUsageLogRepositoryCreate_PersistsServiceTier(t *testing.T) { log.CacheCreation1hTokens, log.ImageOutputTokens, log.ImageOutputCost, + log.ImageInputTokens, + log.ImageInputCost, log.InputCost, log.OutputCost, log.CacheCreationCost, @@ -267,11 +271,11 @@ func TestPrepareUsageLogInsert_PersistsImageSizeMetadata(t *testing.T) { CreatedAt: time.Date(2025, 1, 6, 12, 0, 0, 0, time.UTC), }) - require.Equal(t, sql.NullString{String: imageSize, Valid: true}, prepared.args[34]) - require.Equal(t, sql.NullString{String: inputSize, Valid: true}, prepared.args[35]) - require.Equal(t, sql.NullString{String: outputSize, Valid: true}, prepared.args[36]) - require.Equal(t, sql.NullString{String: source, Valid: true}, prepared.args[37]) - breakdownJSON, ok := prepared.args[38].(string) + require.Equal(t, sql.NullString{String: imageSize, Valid: true}, prepared.args[36]) + require.Equal(t, sql.NullString{String: inputSize, Valid: true}, prepared.args[37]) + require.Equal(t, sql.NullString{String: outputSize, Valid: true}, prepared.args[38]) + require.Equal(t, sql.NullString{String: source, Valid: true}, prepared.args[39]) + breakdownJSON, ok := prepared.args[40].(string) require.True(t, ok) require.JSONEq(t, `{"1K":1,"4K":1}`, breakdownJSON) } @@ -790,6 +794,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullInt64{}, 0, 0, 0, 0, 0, 0, 0, 0.0, // image_output_tokens, image_output_cost + 0, 0.0, // image_input_tokens, image_input_cost 0.0, 0.0, 0.0, 0.0, 0.8, 0.8, 1.0, sql.NullFloat64{}, @@ -857,6 +862,8 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { 6, // cache_creation_1h_tokens 0, // image_output_tokens 0.0, // image_output_cost + 0, // image_input_tokens + 0.0, // image_input_cost 0.1, // input_cost 0.2, // output_cost 0.3, // cache_creation_cost @@ -918,6 +925,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullInt64{}, 1, 2, 3, 4, 5, 6, 0, 0.0, // image_output_tokens, image_output_cost + 0, 0.0, // image_input_tokens, image_input_cost 0.1, 0.2, 0.3, 0.4, 1.0, 0.9, 1.0, sql.NullFloat64{}, @@ -974,6 +982,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullInt64{}, 1, 2, 3, 4, 5, 6, 0, 0.0, // image_output_tokens, image_output_cost + 0, 0.0, // image_input_tokens, image_input_cost 0.1, 0.2, 0.3, 0.4, 1.0, 0.9, 1.0, sql.NullFloat64{}, diff --git a/backend/internal/server/api_contract_test.go b/backend/internal/server/api_contract_test.go index 81adaac4b5..a9c865a5ab 100644 --- a/backend/internal/server/api_contract_test.go +++ b/backend/internal/server/api_contract_test.go @@ -603,6 +603,8 @@ func TestAPIContracts(t *testing.T) { "image_size": null, "image_input_size": null, "image_output_size": null, + "image_input_tokens": 0, + "image_input_cost": 0, "image_output_tokens": 0, "image_output_cost": 0, "image_size_source": null, diff --git a/backend/internal/server/middleware/api_key_auth.go b/backend/internal/server/middleware/api_key_auth.go index 082766bfd5..33bf65ea46 100644 --- a/backend/internal/server/middleware/api_key_auth.go +++ b/backend/internal/server/middleware/api_key_auth.go @@ -25,7 +25,8 @@ func NewAPIKeyAuthMiddleware(apiKeyService *service.APIKeyService, subscriptionS // - 鉴权（Authentication）：验证 Key 有效性、用户状态、IP 限制 —— 始终执行 // - 计费执行（Billing Enforcement）：过期/配额/订阅/余额检查 —— skipBilling 时整块跳过 // -// /v1/usage 端点只需鉴权，不需要计费执行（允许过期/配额耗尽的 Key 查询自身用量）。 +// /v1/usage 和 /v1/sub2api/billing 端点只需鉴权，不需要计费执行。 +// 前者允许过期/配额耗尽的 Key 查询自身用量，后者用于读取当前 Key 的倍率配置。 func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc { return func(c *gin.Context) { // ── 1. 提取 API Key ────────────────────────────────────────── @@ -128,6 +129,8 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti } ctx := context.WithValue(c.Request.Context(), ctxkey.UserID, apiKey.User.ID) c.Request = c.Request.WithContext(ctx) + billingInfoRequest := c.Request.URL.Path == "/v1/sub2api/billing" + skipBilling := c.Request.URL.Path == "/v1/usage" || billingInfoRequest // ── 4. SimpleMode → early return ───────────────────────────── @@ -139,20 +142,20 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti }) c.Set(string(ContextKeyUserRole), apiKey.User.Role) setGroupContext(c, apiKey.Group) - _ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID) + if !billingInfoRequest { + _ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID) + } c.Next() return } - // ── 5. 加载订阅（订阅模式时始终加载） ─────────────────────── - - // skipBilling: /v1/usage 只需鉴权，跳过所有计费执行 - skipBilling := c.Request.URL.Path == "/v1/usage" + // ── 5. 按端点需要加载订阅 ─────────────────────────────────── var subscription *service.UserSubscription isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType() - if isSubscriptionType && subscriptionService != nil { + // 倍率自省不需要订阅数据；/v1/usage 仍保留原有订阅读取行为。 + if isSubscriptionType && subscriptionService != nil && !billingInfoRequest { sub, subErr := subscriptionService.GetActiveSubscription( c.Request.Context(), apiKey.User.ID, @@ -237,7 +240,9 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti }) c.Set(string(ContextKeyUserRole), apiKey.User.Role) setGroupContext(c, apiKey.Group) - _ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID) + if !billingInfoRequest { + _ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID) + } c.Next() } diff --git a/backend/internal/server/middleware/api_key_auth_test.go b/backend/internal/server/middleware/api_key_auth_test.go index d9b5c7efc1..2a2e2f5588 100644 --- a/backend/internal/server/middleware/api_key_auth_test.go +++ b/backend/internal/server/middleware/api_key_auth_test.go @@ -1066,6 +1066,129 @@ func TestAPIKeyAuthTouchesLastUsedInStandardMode(t *testing.T) { require.Equal(t, 1, touchCalls) } +func TestAPIKeyAuthBillingInfoSkipsBillingAndSideEffects(t *testing.T) { + gin.SetMode(gin.TestMode) + + group := &service.Group{ + ID: 42, + Name: "subscription", + Status: service.StatusActive, + Hydrated: true, + SubscriptionType: service.SubscriptionTypeSubscription, + } + user := &service.User{ + ID: 7, + Role: service.RoleUser, + Status: service.StatusActive, + Balance: 0, + Concurrency: 3, + } + expiredAt := time.Now().Add(-time.Hour) + apiKey := &service.APIKey{ + ID: 100, + UserID: user.ID, + Key: "billing-info-auth-only", + Status: service.StatusAPIKeyQuotaExhausted, + User: user, + GroupID: &group.ID, + Group: group, + Quota: 1, + QuotaUsed: 1, + ExpiresAt: &expiredAt, + } + + touchCalls := 0 + subscriptionCalls := 0 + apiKeyRepo := &stubApiKeyRepo{ + getByKey: func(context.Context, string) (*service.APIKey, error) { + clone := *apiKey + return &clone, nil + }, + updateLastUsed: func(context.Context, int64, time.Time) error { + touchCalls++ + return nil + }, + } + subscriptionRepo := &stubUserSubscriptionRepo{ + getActive: func(context.Context, int64, int64) (*service.UserSubscription, error) { + subscriptionCalls++ + return nil, service.ErrSubscriptionNotFound + }, + } + cfg := &config.Config{RunMode: config.RunModeStandard} + apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg) + subscriptionService := service.NewSubscriptionService(nil, subscriptionRepo, nil, nil, cfg) + t.Cleanup(subscriptionService.Stop) + router := newAuthTestRouter(apiKeyService, subscriptionService, cfg) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil) + req.Header.Set("x-api-key", apiKey.Key) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Zero(t, subscriptionCalls) + require.Zero(t, touchCalls) +} + +func TestAPIKeyAuthBillingInfoSkipsLastUsedInSimpleMode(t *testing.T) { + gin.SetMode(gin.TestMode) + + user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive} + apiKey := &service.APIKey{ID: 100, UserID: user.ID, Key: "billing-info-simple", Status: service.StatusActive, User: user} + touchCalls := 0 + apiKeyRepo := &stubApiKeyRepo{ + getByKey: func(context.Context, string) (*service.APIKey, error) { + clone := *apiKey + return &clone, nil + }, + updateLastUsed: func(context.Context, int64, time.Time) error { + touchCalls++ + return nil + }, + } + cfg := &config.Config{RunMode: config.RunModeSimple} + apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg) + router := newAuthTestRouter(apiKeyService, nil, cfg) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil) + req.Header.Set("x-api-key", apiKey.Key) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Zero(t, touchCalls) +} + +func TestAPIKeyAuthUsageStillTouchesLastUsed(t *testing.T) { + gin.SetMode(gin.TestMode) + + user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive, Balance: 10} + apiKey := &service.APIKey{ID: 100, UserID: user.ID, Key: "usage-touch", Status: service.StatusActive, User: user} + touchCalls := 0 + apiKeyRepo := &stubApiKeyRepo{ + getByKey: func(context.Context, string) (*service.APIKey, error) { + clone := *apiKey + return &clone, nil + }, + updateLastUsed: func(context.Context, int64, time.Time) error { + touchCalls++ + return nil + }, + } + cfg := &config.Config{RunMode: config.RunModeStandard} + apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg) + router := newAuthTestRouter(apiKeyService, nil, cfg) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v1/usage", nil) + req.Header.Set("x-api-key", apiKey.Key) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Equal(t, 1, touchCalls) +} + func TestAPIKeyAuthAllowsBalanceBelowMinimumReserve(t *testing.T) { gin.SetMode(gin.TestMode) @@ -1155,9 +1278,12 @@ func TestAPIKeyAuthRejectsExhaustedBalance(t *testing.T) { func newAuthTestRouter(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) *gin.Engine { router := gin.New() router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, cfg))) - router.GET("/t", func(c *gin.Context) { + ok := func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true}) - }) + } + router.GET("/t", ok) + router.GET("/v1/usage", ok) + router.GET("/v1/sub2api/billing", ok) return router } diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index f73379981e..61b4c54ced 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -114,6 +114,7 @@ func RegisterGatewayRoutes( gateway.Use(opsErrorLogger) gateway.Use(endpointNorm) gateway.Use(gin.HandlerFunc(apiKeyAuth)) + gateway.GET("/sub2api/billing", h.Gateway.KeyBillingInfo) gateway.Use(requireGroupAnthropic) { // /v1/messages: auto-route based on group platform diff --git a/backend/internal/server/routes/gateway_key_billing_test.go b/backend/internal/server/routes/gateway_key_billing_test.go new file mode 100644 index 0000000000..ef69a5eb44 --- /dev/null +++ b/backend/internal/server/routes/gateway_key_billing_test.go @@ -0,0 +1,172 @@ +package routes + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/handler" + servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/Wei-Shaw/sub2api/internal/web" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type keyBillingRouteAPIKeyRepo struct { + service.APIKeyRepository + apiKey *service.APIKey +} + +func (r *keyBillingRouteAPIKeyRepo) GetByKeyForAuth(_ context.Context, key string) (*service.APIKey, error) { + if r.apiKey == nil || key != r.apiKey.Key { + return nil, service.ErrAPIKeyNotFound + } + clone := *r.apiKey + return &clone, nil +} + +type keyBillingRouteRateRepo struct { + service.UserGroupRateRepository + lookupCalls int +} + +func (r *keyBillingRouteRateRepo) GetByUserAndGroup(context.Context, int64, int64) (*float64, error) { + r.lookupCalls++ + return nil, nil +} + +func (r *keyBillingRouteRateRepo) GetRPMOverrideByUserAndGroup(context.Context, int64, int64) (*int, error) { + return nil, nil +} + +func newKeyBillingRouteTestRouter(runMode string) (*gin.Engine, *keyBillingRouteRateRepo, string) { + gin.SetMode(gin.TestMode) + group := &service.Group{ + ID: 42, + Status: service.StatusActive, + Hydrated: true, + Platform: service.PlatformOpenAI, + SubscriptionType: service.SubscriptionTypeStandard, + RateMultiplier: 0.75, + } + user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive, Balance: 10} + var groupID *int64 + var apiKeyGroup *service.Group + if runMode != config.RunModeSimple { + groupID = &group.ID + apiKeyGroup = group + } + apiKey := &service.APIKey{ + ID: 100, + UserID: user.ID, + Key: "billing-route-test-key", + Status: service.StatusActive, + User: user, + GroupID: groupID, + Group: apiKeyGroup, + } + cfg := &config.Config{RunMode: runMode} + rateRepo := &keyBillingRouteRateRepo{} + apiKeyService := service.NewAPIKeyService( + &keyBillingRouteAPIKeyRepo{apiKey: apiKey}, nil, nil, nil, rateRepo, nil, cfg, + ) + gatewayService := service.NewGatewayService( + nil, nil, nil, nil, nil, nil, rateRepo, nil, cfg, nil, nil, nil, nil, nil, + nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + ) + openAIGatewayService := service.NewOpenAIGatewayService( + nil, nil, nil, nil, nil, rateRepo, nil, cfg, nil, nil, nil, + nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + ) + gatewayHandler := handler.NewGatewayHandler( + gatewayService, openAIGatewayService, nil, nil, nil, nil, nil, nil, + apiKeyService, nil, nil, nil, nil, cfg, nil, + ) + + router := gin.New() + if web.HasEmbeddedFrontend() { + router.Use(web.ServeEmbeddedFrontend()) + } + RegisterGatewayRoutes( + router, + &handler.Handlers{Gateway: gatewayHandler, OpenAIGateway: &handler.OpenAIGatewayHandler{}}, + servermiddleware.NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg), + apiKeyService, + nil, + nil, + nil, + cfg, + ) + return router, rateRepo, apiKey.Key +} + +func TestGatewayRoutesKeyBillingInfoPathIsRegistered(t *testing.T) { + router := newGatewayRoutesTestRouter() + + for _, route := range router.Routes() { + if route.Method == http.MethodGet && route.Path == "/v1/sub2api/billing" { + return + } + } + + t.Fatal("GET /v1/sub2api/billing should be registered") +} + +func TestGatewayRoutesKeyBillingInfoEndToEnd(t *testing.T) { + t.Run("missing credentials", func(t *testing.T) { + router, rateRepo, _ := newKeyBillingRouteTestRouter(config.RunModeStandard) + w := httptest.NewRecorder() + + router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)) + + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Contains(t, w.Header().Get("Content-Type"), "application/json") + require.NotContains(t, strings.ToLower(w.Body.String()), "") + require.Zero(t, rateRepo.lookupCalls) + }) + + t.Run("standard mode", func(t *testing.T) { + router, rateRepo, key := newKeyBillingRouteTestRouter(config.RunModeStandard) + req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil) + req.Header.Set("Authorization", "Bearer "+key) + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Contains(t, w.Header().Get("Content-Type"), "application/json") + require.Equal(t, "no-store", w.Header().Get("Cache-Control")) + require.NotContains(t, strings.ToLower(w.Body.String()), "") + var body map[string]any + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body)) + require.Equal(t, "sub2api.key_billing", body["object"]) + require.Equal(t, 0.75, body["effective_rate_multiplier"]) + require.Equal(t, 1, rateRepo.lookupCalls) + }) + + t.Run("simple mode", func(t *testing.T) { + router, rateRepo, key := newKeyBillingRouteTestRouter(config.RunModeSimple) + req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil) + req.Header.Set("x-api-key", key) + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusNotFound, w.Code) + require.Contains(t, w.Header().Get("Content-Type"), "application/json") + require.NotContains(t, strings.ToLower(w.Body.String()), "") + require.JSONEq(t, `{ + "type": "error", + "error": { + "type": "not_found_error", + "message": "Billing information is not supported in simple mode" + } + }`, w.Body.String()) + require.Zero(t, rateRepo.lookupCalls) + }) +} diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 8890254069..6aa2ae8a2f 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -87,6 +87,7 @@ const openAILongContextBillingEnabledKey = "openai_long_context_billing_enabled" const ( OpenAIEndpointCapabilityChatCompletions OpenAIEndpointCapability = "chat_completions" OpenAIEndpointCapabilityEmbeddings OpenAIEndpointCapability = "embeddings" + OpenAIEndpointCapabilityAlphaSearch OpenAIEndpointCapability = "alpha_search" ) const openAIEndpointCapabilitiesCredentialKey = "openai_capabilities" @@ -1397,6 +1398,13 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa } switch capability { case OpenAIEndpointCapabilityChatCompletions: + case OpenAIEndpointCapabilityAlphaSearch: + // Codex alpha/search 是 ChatGPT/Codex 后端工具端点，必须使用 + // OAuth/PAT/AgentIdentity 这类 ChatGPT 账号凭据；API key 被发往 + // chatgpt.com/backend-api/codex/alpha/search 会稳定 401。 + if a.Type != AccountTypeOAuth { + return false + } case OpenAIEndpointCapabilityEmbeddings: if a.Type != AccountTypeAPIKey { return false @@ -1409,6 +1417,9 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa if !found { return true } + if capability == OpenAIEndpointCapabilityAlphaSearch && configured[string(OpenAIEndpointCapabilityChatCompletions)] { + return true + } return configured[string(capability)] } diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 7fa69d41ea..8ebe754bdb 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -153,7 +153,8 @@ type UsageTokens struct { // CostBreakdown 费用明细 type CostBreakdown struct { - InputCost float64 + InputCost float64 // 文本输入费用（不含图片输入，图片输入单独记入 ImageInputCost） + ImageInputCost float64 // 图片输入 token 费用（如 gpt-image-2 图片编辑） OutputCost float64 ImageOutputCost float64 CacheCreationCost float64 @@ -802,6 +803,7 @@ func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) { LongContextInputThreshold: litellmPricing.LongContextInputTokenThreshold, LongContextInputMultiplier: litellmPricing.LongContextInputCostMultiplier, LongContextOutputMultiplier: litellmPricing.LongContextOutputCostMultiplier, + ImageInputPricePerToken: litellmPricing.InputCostPerImageToken, ImageOutputPricePerToken: litellmPricing.OutputCostPerImageToken, }), nil } @@ -859,6 +861,7 @@ func (s *BillingService) GetModelPricingWithChannel(model string, channelPricing pricing.ImageOutputPricePerToken = 0 } pricing.ImageOutputPriceExplicit = true + applyChannelImageInputPrice(channelPricing, pricing) return pricing, nil } @@ -1001,7 +1004,8 @@ func (s *BillingService) computeTokenBreakdown( } bd := &CostBreakdown{} - // 分离图片输入 token 与文本输入 token（多模态 embedding 等图文不同价场景）。 + // 分离图片输入 token 与文本输入 token（多模态 embedding、图片编辑等图文不同价场景）。 + // InputCost 仅计文本输入，图片输入费用单独记入 ImageInputCost，便于对账；总额不变。 // ImageInputTokens 为 0 时（绝大多数 chat/vision 流量）走原始单价路径，行为不变。 if tokens.ImageInputTokens > 0 { imageInputTokens := tokens.ImageInputTokens @@ -1015,7 +1019,8 @@ func (s *BillingService) computeTokenBreakdown( // 未配置图片输入档时回退到文本 input 价（已含 priority / 长上下文调整） imageInputPrice = inputPrice } - bd.InputCost = float64(textInputTokens)*inputPrice + float64(imageInputTokens)*imageInputPrice + bd.InputCost = float64(textInputTokens) * inputPrice + bd.ImageInputCost = float64(imageInputTokens) * imageInputPrice } else { bd.InputCost = float64(tokens.InputTokens) * inputPrice } @@ -1043,13 +1048,14 @@ func (s *BillingService) computeTokenBreakdown( if tierMultiplier != 1.0 { bd.InputCost *= tierMultiplier + bd.ImageInputCost *= tierMultiplier bd.OutputCost *= tierMultiplier bd.ImageOutputCost *= tierMultiplier bd.CacheCreationCost *= tierMultiplier bd.CacheReadCost *= tierMultiplier } - bd.TotalCost = bd.InputCost + bd.OutputCost + bd.ImageOutputCost + + bd.TotalCost = bd.InputCost + bd.ImageInputCost + bd.OutputCost + bd.ImageOutputCost + bd.CacheCreationCost + bd.CacheReadCost bd.ActualCost = bd.TotalCost * rateMultiplier bd.LongContextBillingApplied = baselineCost != nil && bd.ActualCost > baselineCost.ActualCost @@ -1277,6 +1283,7 @@ func (s *BillingService) CalculateCostWithLongContext(model string, tokens Usage // 合并成本 return &CostBreakdown{ InputCost: inRangeCost.InputCost + outRangeCost.InputCost, + ImageInputCost: inRangeCost.ImageInputCost + outRangeCost.ImageInputCost, OutputCost: inRangeCost.OutputCost, ImageOutputCost: inRangeCost.ImageOutputCost, CacheCreationCost: inRangeCost.CacheCreationCost, diff --git a/backend/internal/service/billing_service_test.go b/backend/internal/service/billing_service_test.go index 885da194e3..265fbf4fa4 100644 --- a/backend/internal/service/billing_service_test.go +++ b/backend/internal/service/billing_service_test.go @@ -755,7 +755,7 @@ func TestGetModelPricing_DoubaoEmbeddingVisionImageInputRate(t *testing.T) { } } -// 验证双档计费：InputCost = 文本token×文本价 + 图片token×图片价； +// 验证双档计费：InputCost = 文本token×文本价（不含图片），ImageInputCost = 图片token×图片价； // 且 ImageInputTokens=0 时走原单价路径，ImageInputTokens>InputTokens 时不负计文本。 func TestCalculateCost_DoubaoEmbeddingVisionDifferentialInput(t *testing.T) { svc := newTestBillingService() @@ -764,22 +764,60 @@ func TestCalculateCost_DoubaoEmbeddingVisionDifferentialInput(t *testing.T) { mixed := UsageTokens{InputTokens: 1340, ImageInputTokens: 28} cost, err := svc.CalculateCost("doubao-embedding-vision", mixed, 1.0) require.NoError(t, err) - wantMixed := float64(1312)*0.098e-6 + float64(28)*0.252e-6 - require.InDelta(t, wantMixed, cost.InputCost, 1e-15) - require.InDelta(t, wantMixed, cost.TotalCost, 1e-15) + wantText := float64(1312) * 0.098e-6 + wantImage := float64(28) * 0.252e-6 + require.InDelta(t, wantText, cost.InputCost, 1e-15, "InputCost 仅计文本输入") + require.InDelta(t, wantImage, cost.ImageInputCost, 1e-15, "ImageInputCost 单独计图片输入") + require.InDelta(t, wantText+wantImage, cost.TotalCost, 1e-15, "TotalCost 口径不变") require.Zero(t, cost.OutputCost) - // 纯文本：全部按文本档计费，与原单价路径一致。 + // 纯文本：全部按文本档计费，与原单价路径一致，无图片输入费用。 textOnly := UsageTokens{InputTokens: 1340} costText, err := svc.CalculateCost("doubao-embedding-vision", textOnly, 1.0) require.NoError(t, err) require.InDelta(t, float64(1340)*0.098e-6, costText.InputCost, 1e-15) + require.Zero(t, costText.ImageInputCost) // 健壮性：ImageInputTokens 超过 InputTokens 时，文本置 0、计费 token 不超过 InputTokens。 weird := UsageTokens{InputTokens: 10, ImageInputTokens: 50} costWeird, err := svc.CalculateCost("doubao-embedding-vision", weird, 1.0) require.NoError(t, err) - require.InDelta(t, float64(10)*0.252e-6, costWeird.InputCost, 1e-15) + require.Zero(t, costWeird.InputCost, "全为图片输入时文本费用为 0") + require.InDelta(t, float64(10)*0.252e-6, costWeird.ImageInputCost, 1e-15) + require.InDelta(t, float64(10)*0.252e-6, costWeird.TotalCost, 1e-15) +} + +// 复现 issue #4386：gpt-image-2 /v1/images/edits 带 1 张输入图。 +// 上游 usage：input_tokens=371（image_tokens=352 + text_tokens=19）， +// output_tokens=439（全部图片输出）。官方定价：文本输入 $5/1M、图片输入 $8/1M、 +// 文本输出 $10/1M、图片输出 $30/1M。修复前图片输入被并入文本价，单次偏低 ~6.6%。 +func TestComputeTokenBreakdown_GptImage2ImageEditIssue4386(t *testing.T) { + svc := newTestBillingService() + + pricing := &ModelPricing{ + InputPricePerToken: 5e-6, + ImageInputPricePerToken: 8e-6, + OutputPricePerToken: 10e-6, + ImageOutputPricePerToken: 30e-6, + ImageOutputPriceExplicit: true, + } + tokens := UsageTokens{ + InputTokens: 371, + ImageInputTokens: 352, + OutputTokens: 439, + ImageOutputTokens: 439, + } + + cost := svc.computeTokenBreakdown(pricing, tokens, 1.0, "", false) + + wantTextInput := float64(19) * 5e-6 // 0.000095 + wantImageInput := float64(352) * 8e-6 // 0.002816 + wantImageOutput := float64(439) * 30e-6 // 0.013170 + require.InDelta(t, wantTextInput, cost.InputCost, 1e-15, "InputCost 仅含文本输入") + require.InDelta(t, wantImageInput, cost.ImageInputCost, 1e-15, "图片输入按 $8/1M 独立计费") + require.Zero(t, cost.OutputCost, "输出全部为图片，文本输出费用为 0") + require.InDelta(t, wantImageOutput, cost.ImageOutputCost, 1e-15) + require.InDelta(t, 0.016081, cost.TotalCost, 1e-9, "总额应为 $0.016081（修复前为 $0.015025）") } func TestCalculateCostWithLongContext_BelowThreshold(t *testing.T) { svc := newTestBillingService() diff --git a/backend/internal/service/channel.go b/backend/internal/service/channel.go index 1fd9e57068..4a3bdf304d 100644 --- a/backend/internal/service/channel.go +++ b/backend/internal/service/channel.go @@ -92,6 +92,7 @@ type ChannelModelPricing struct { OutputPrice *float64 // 每 token 输出价格（USD） CacheWritePrice *float64 // 缓存写入价格 CacheReadPrice *float64 // 缓存读取价格 + ImageInputPrice *float64 // 图片输入 token 价格（如 gpt-image-2 图片编辑）；未配置时回退文本输入价 ImageOutputPrice *float64 // 图片输出价格（向后兼容） PerRequestPrice *float64 // 默认按次计费价格（USD） Intervals []PricingInterval // 区间定价列表 diff --git a/backend/internal/service/channel_service.go b/backend/internal/service/channel_service.go index 5272fe0e93..25c7b587e3 100644 --- a/backend/internal/service/channel_service.go +++ b/backend/internal/service/channel_service.go @@ -645,6 +645,7 @@ func checkPricesNotNegative(p ChannelModelPricing) error { {"output_price", p.OutputPrice}, {"cache_write_price", p.CacheWritePrice}, {"cache_read_price", p.CacheReadPrice}, + {"image_input_price", p.ImageInputPrice}, {"image_output_price", p.ImageOutputPrice}, {"per_request_price", p.PerRequestPrice}, } diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 61ab3abd2f..71236b69bd 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -29,6 +29,11 @@ func (s *GatewayService) getUserGroupRateMultiplier(ctx context.Context, userID, return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier) } +// ResolveUserGroupRateMultiplier resolves the same cached multiplier used by usage billing. +func (s *GatewayService) ResolveUserGroupRateMultiplier(ctx context.Context, userID, groupID int64, groupDefaultMultiplier float64) float64 { + return s.getUserGroupRateMultiplier(ctx, userID, groupID, groupDefaultMultiplier) +} + // RecordUsageInput 记录使用量的输入参数。 // 异步 worker 只接收计费所需快照，不能持有 ParsedRequest/RequestBodyRef 这类大请求体引用。 type RecordUsageInput struct { @@ -662,7 +667,7 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage } if apiKey.GroupID != nil && apiKey.Group != nil { groupDefault := apiKey.Group.RateMultiplier - multiplier = s.getUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault) + multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault) } // token 倍率叠加高峰因子（token 计费含图片 token，图片按次倍率不受影响）。高峰因子按请求时刻现算， // 不并入上面的 getUserGroupRateMultiplier，以免污染 user:group 倍率缓存。 diff --git a/backend/internal/service/image_generation_intent.go b/backend/internal/service/image_generation_intent.go index 72f3be7db6..9351fcd682 100644 --- a/backend/internal/service/image_generation_intent.go +++ b/backend/internal/service/image_generation_intent.go @@ -79,6 +79,56 @@ func IsImageGenerationIntent(endpoint string, requestedModel string, body []byte return imageIntent } +// IsImageGenerationIntentForPlatform applies platform-specific intent rules. +// +// Codex advertises the image_gen namespace on ordinary Responses requests so +// that it is available if the model needs it. Grok strips namespace and +// Responses Lite additional_tools declarations before forwarding, so those +// declarations alone must not turn every Codex request into an image request. +// Native image_generation tools, explicit image selection and image models +// remain image intent. Other platforms retain the original declaration rule. +func IsImageGenerationIntentForPlatform(endpoint string, requestedModel string, body []byte, platform string) bool { + if !strings.EqualFold(strings.TrimSpace(platform), PlatformGrok) { + return IsImageGenerationIntent(endpoint, requestedModel, body) + } + return isExplicitGrokImageGenerationIntent(endpoint, requestedModel, body) +} + +func isExplicitGrokImageGenerationIntent(endpoint string, requestedModel string, body []byte) bool { + if IsImageGenerationEndpoint(endpoint) || isOpenAIImageGenerationModel(requestedModel) { + return true + } + if len(body) == 0 || !gjson.ValidBytes(body) { + return false + } + + var modelSeen, toolsSeen, toolChoiceSeen bool + imageIntent := false + parseRawJSONView(body).ForEach(func(key, value gjson.Result) bool { + switch key.Str { + case "model": + if !modelSeen { + modelSeen = true + imageIntent = isOpenAIImageGenerationModel(strings.TrimSpace(value.String())) + } + case "tools": + if !toolsSeen { + toolsSeen = true + // Grok removes namespace catalogs before forwarding. Native + // image_generation remains an explicit capability request. + imageIntent = openAIJSONToolsContainNativeImageGeneration(value) + } + case "tool_choice": + if !toolChoiceSeen { + toolChoiceSeen = true + imageIntent = openAIJSONToolChoiceSelectsExplicitImageGeneration(value) + } + } + return !imageIntent && (!modelSeen || !toolsSeen || !toolChoiceSeen) + }) + return imageIntent +} + // IsImageGenerationIntentMap is the map-backed variant used after service-side request mutation. func IsImageGenerationIntentMap(endpoint string, requestedModel string, reqBody map[string]any) bool { if IsImageGenerationEndpoint(endpoint) { @@ -140,6 +190,18 @@ func openAIJSONToolsContainImageGeneration(tools gjson.Result) bool { return found } +func openAIJSONToolsContainNativeImageGeneration(tools gjson.Result) bool { + if !tools.IsArray() { + return false + } + found := false + tools.ForEach(func(_, item gjson.Result) bool { + found = isOpenAIImageGenerationType(openAIJSONString(item.Get("type"))) + return !found + }) + return found +} + func isOpenAIImageGenerationType(value string) bool { return strings.TrimSpace(value) == "image_generation" } @@ -148,9 +210,8 @@ func isOpenAIImageGenNamespaceName(value string) bool { return strings.TrimSpace(value) == "image_gen" } -// isImageGenNamespaceTool detects the Codex namespace-style image generation -// tool declaration: { "type": "namespace", "name": "image_gen", ... }. -// Codex /image uses this instead of the flat { "type": "image_generation" }. +// isImageGenNamespaceTool detects the namespace advertised by Codex's built-in +// image-generation extension instead of a hosted image_generation tool. func isImageGenNamespaceTool(tool gjson.Result) bool { return openAIJSONString(tool.Get("type")) == "namespace" && isOpenAIImageGenNamespaceName(openAIJSONString(tool.Get("name"))) @@ -235,6 +296,45 @@ func openAIJSONToolChoiceSelectsImageGeneration(choice gjson.Result) bool { return false } +func openAIJSONToolChoiceSelectsExplicitImageGeneration(choice gjson.Result) bool { + if openAIJSONToolChoiceSelectsImageGeneration(choice) { + return true + } + if !choice.IsObject() { + return false + } + if tool := choice.Get("tool"); tool.IsObject() && openAIJSONToolChoiceSelectsExplicitImageGeneration(tool) { + return true + } + if isOpenAIImageGenFunctionReference( + openAIJSONString(choice.Get("namespace")), + openAIJSONString(choice.Get("name")), + ) { + return true + } + if fn := choice.Get("function"); fn.IsObject() { + return isOpenAIImageGenFunctionReference( + openAIJSONString(fn.Get("namespace")), + openAIJSONString(fn.Get("name")), + ) + } + return false +} + +func isOpenAIImageGenFunctionReference(namespace string, name string) bool { + namespace = strings.TrimSpace(namespace) + name = strings.TrimSpace(name) + if namespace == "image_gen" && name == "imagegen" { + return true + } + switch name { + case "image_gen.imagegen", "image_gen__imagegen": + return true + default: + return false + } +} + func openAIAnyToolChoiceSelectsImageGeneration(choice any) bool { switch v := choice.(type) { case string: diff --git a/backend/internal/service/image_generation_intent_grok_test.go b/backend/internal/service/image_generation_intent_grok_test.go new file mode 100644 index 0000000000..cfef208d45 --- /dev/null +++ b/backend/internal/service/image_generation_intent_grok_test.go @@ -0,0 +1,108 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestIsImageGenerationIntentForPlatform_GrokCodexDeclarations(t *testing.T) { + tests := []struct { + name string + body string + want bool + }{ + { + name: "top-level image_gen namespace is passive", + body: `{"model":"grok-4.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}],"tool_choice":"auto","input":"write code"}`, + }, + { + name: "responses lite additional_tools image_gen is passive", + body: `{"model":"grok-4.5","tool_choice":"auto","input":[{"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]},{"type":"message","role":"user","content":"write code"}]}`, + }, + { + name: "legacy adapter flattened image_gen function declaration is passive", + body: `{"model":"grok-4.5","tools":[{"type":"function","name":"image_gen.imagegen"}],"input":"write code"}`, + }, + { + name: "native image generation declaration remains explicit", + body: `{"model":"grok-4.5","tools":[{"type":"image_generation"}],"input":"draw a cat"}`, + want: true, + }, + { + name: "native image generation tool choice is explicit", + body: `{"model":"grok-4.5","tools":[{"type":"image_generation"}],"tool_choice":{"type":"image_generation"},"input":"draw a cat"}`, + want: true, + }, + { + name: "image_gen namespace tool choice is explicit", + body: `{"model":"grok-4.5","tools":[{"type":"namespace","name":"image_gen"}],"tool_choice":{"type":"namespace","name":"image_gen"},"input":"draw a cat"}`, + want: true, + }, + { + name: "flattened image_gen function tool choice is explicit", + body: `{"model":"grok-4.5","tools":[{"type":"function","name":"image_gen.imagegen"}],"tool_choice":{"type":"function","name":"image_gen.imagegen"},"input":"draw a cat"}`, + want: true, + }, + { + name: "wrapped image_gen function tool choice is explicit", + body: `{"model":"grok-4.5","tools":[{"type":"function","name":"image_gen.imagegen"}],"tool_choice":{"tool":{"type":"function","name":"image_gen.imagegen"}},"input":"draw a cat"}`, + want: true, + }, + { + name: "vision input is not image generation", + body: `{"model":"grok-4.5","input":[{"type":"message","role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}]}`, + }, + { + name: "ordinary function declaration is not image generation", + body: `{"model":"grok-4.5","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto","input":"write code"}`, + }, + { + name: "image_gen function call history is not current intent", + body: `{"model":"grok-4.5","input":[{"type":"function_call","namespace":"image_gen","name":"imagegen","arguments":"{}"}]}`, + }, + { + name: "image generation call history is not current intent", + body: `{"model":"grok-4.5","input":[{"type":"image_generation_call","id":"ig_1"}]}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, IsImageGenerationIntentForPlatform( + openAIResponsesEndpoint, + "grok-4.5", + []byte(tt.body), + PlatformGrok, + )) + }) + } +} + +func TestIsImageGenerationIntentForPlatform_GrokPreservesHardSignals(t *testing.T) { + require.True(t, IsImageGenerationIntentForPlatform( + "/v1/images/generations", + "grok-4.5", + []byte(`{"input":"draw"}`), + PlatformGrok, + )) + require.True(t, IsImageGenerationIntentForPlatform( + openAIResponsesEndpoint, + "gpt-image-2", + []byte(`{"input":"draw"}`), + PlatformGrok, + )) + require.True(t, IsImageGenerationIntentForPlatform( + openAIResponsesEndpoint, + "grok-4.5", + []byte(`{"model":"gpt-image-2","input":"draw"}`), + PlatformGrok, + )) +} + +func TestIsImageGenerationIntentForPlatform_OtherPlatformsKeepDeclarationSemantics(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}],"input":"write code"}`) + + require.True(t, IsImageGenerationIntentForPlatform(openAIResponsesEndpoint, "gpt-5.5", body, PlatformOpenAI)) + require.True(t, IsImageGenerationIntentForPlatform(openAIResponsesEndpoint, "gpt-5.5", body, PlatformAnthropic)) +} diff --git a/backend/internal/service/model_pricing_resolver.go b/backend/internal/service/model_pricing_resolver.go index a9f603113b..ea5f4a29d1 100644 --- a/backend/internal/service/model_pricing_resolver.go +++ b/backend/internal/service/model_pricing_resolver.go @@ -161,6 +161,7 @@ func (r *ModelPricingResolver) applyTokenOverrides(chPricing *ChannelModelPricin resolved.BasePricing.ImageOutputPricePerToken = 0 } resolved.BasePricing.ImageOutputPriceExplicit = true + applyChannelImageInputPrice(chPricing, resolved.BasePricing) return } @@ -199,6 +200,20 @@ func (r *ModelPricingResolver) applyTokenOverrides(chPricing *ChannelModelPricin resolved.BasePricing.ImageOutputPricePerToken = 0 } resolved.BasePricing.ImageOutputPriceExplicit = true + applyChannelImageInputPrice(chPricing, resolved.BasePricing) +} + +// applyChannelImageInputPrice 应用渠道图片输入价：显式配置则用配置值； +// 未配置时归零，使 computeTokenBreakdown 回退到文本输入价（向后兼容， +// 避免 commit 引入的 LiteLLM 图片输入价泄漏进渠道自定义定价）。 +// 与 image_output 不同，此处不设 Explicit 标志——图片输入未配置应回退文本价， +// 而非硬置 0。 +func applyChannelImageInputPrice(chPricing *ChannelModelPricing, pricing *ModelPricing) { + if chPricing != nil && chPricing.ImageInputPrice != nil { + pricing.ImageInputPricePerToken = *chPricing.ImageInputPrice + } else { + pricing.ImageInputPricePerToken = 0 + } } // applyRequestTierOverrides 应用按次/图片模式的渠道覆盖 @@ -262,12 +277,14 @@ func intervalToModelPricing(iv *PricingInterval, supportsCacheBreakdown bool, ch pricing.CacheReadPricePerToken = *iv.CacheReadPrice pricing.CacheReadPricePerTokenPriority = *iv.CacheReadPrice } - // 渠道定价存在时，ImageOutputPrice 显式覆盖 + // 渠道定价存在时，ImageOutputPrice 显式覆盖；图片输入价用渠道级配置 + // （区间不携带图片输入价，与 image_output 一致）。 if chPricing != nil { pricing.ImageOutputPriceExplicit = true if chPricing.ImageOutputPrice != nil { pricing.ImageOutputPricePerToken = *chPricing.ImageOutputPrice } + applyChannelImageInputPrice(chPricing, pricing) } return pricing } diff --git a/backend/internal/service/openai_account_model_transient.go b/backend/internal/service/openai_account_model_transient.go new file mode 100644 index 0000000000..6ef710d23c --- /dev/null +++ b/backend/internal/service/openai_account_model_transient.go @@ -0,0 +1,177 @@ +package service + +import ( + "strings" + "sync" + "time" +) + +const ( + openAIModelTransientFailureWindow = time.Minute + openAIModelTransientShortCooldown = 10 * time.Second + openAIModelTransientLongCooldown = 45 * time.Second + openAIModelTransientDefaultMax = 4096 + openAIModelTransientMaxModelBytes = 512 +) + +type openAIAccountModelKey struct { + AccountID int64 + Model string +} + +type openAIAccountModelTransientEntry struct { + failureStreak int + lastFailure time.Time + blockUntil time.Time + lastTouched time.Time +} + +type openAIAccountModelTransientDecision struct { + FailureStreak int + Cooldown time.Duration + BlockUntil time.Time +} + +type openAIAccountModelTransientState struct { + mu sync.Mutex + entries map[openAIAccountModelKey]openAIAccountModelTransientEntry + maxEntries int +} + +func newOpenAIAccountModelTransientState(maxEntries int) *openAIAccountModelTransientState { + if maxEntries <= 0 { + maxEntries = openAIModelTransientDefaultMax + } + return &openAIAccountModelTransientState{ + entries: make(map[openAIAccountModelKey]openAIAccountModelTransientEntry), + maxEntries: maxEntries, + } +} + +func normalizeOpenAIAccountModelTransientModel(model string) string { + model = strings.TrimSpace(model) + if len(model) > openAIModelTransientMaxModelBytes { + return "" + } + return strings.ToLower(model) +} + +func openAIAccountModelTransientKey(accountID int64, model string) (openAIAccountModelKey, bool) { + model = normalizeOpenAIAccountModelTransientModel(model) + if accountID <= 0 || model == "" { + return openAIAccountModelKey{}, false + } + return openAIAccountModelKey{AccountID: accountID, Model: model}, true +} + +func (s *openAIAccountModelTransientState) recordFailure(accountID int64, model string, now time.Time) openAIAccountModelTransientDecision { + key, ok := openAIAccountModelTransientKey(accountID, model) + if s == nil || !ok { + return openAIAccountModelTransientDecision{} + } + if now.IsZero() { + now = time.Now() + } + + s.mu.Lock() + defer s.mu.Unlock() + if s.entries == nil { + s.entries = make(map[openAIAccountModelKey]openAIAccountModelTransientEntry) + } + if s.maxEntries <= 0 { + s.maxEntries = openAIModelTransientDefaultMax + } + + entry, exists := s.entries[key] + if !exists { + s.evictOldestLocked() + } + if !exists || entry.lastFailure.IsZero() || now.Sub(entry.lastFailure) > openAIModelTransientFailureWindow || now.Before(entry.lastFailure) { + entry.failureStreak = 0 + entry.blockUntil = time.Time{} + } + entry.failureStreak++ + entry.lastFailure = now + entry.lastTouched = now + + cooldown := time.Duration(0) + switch { + case entry.failureStreak >= 3: + cooldown = openAIModelTransientLongCooldown + case entry.failureStreak == 2: + cooldown = openAIModelTransientShortCooldown + } + if cooldown > 0 { + entry.blockUntil = now.Add(cooldown) + } else { + entry.blockUntil = time.Time{} + } + s.entries[key] = entry + return openAIAccountModelTransientDecision{ + FailureStreak: entry.failureStreak, + Cooldown: cooldown, + BlockUntil: entry.blockUntil, + } +} + +func (s *openAIAccountModelTransientState) recordSuccess(accountID int64, model string) { + key, ok := openAIAccountModelTransientKey(accountID, model) + if s == nil || !ok { + return + } + s.mu.Lock() + delete(s.entries, key) + s.mu.Unlock() +} + +func (s *openAIAccountModelTransientState) isBlocked(accountID int64, model string, now time.Time) bool { + key, ok := openAIAccountModelTransientKey(accountID, model) + if s == nil || !ok { + return false + } + if now.IsZero() { + now = time.Now() + } + + s.mu.Lock() + defer s.mu.Unlock() + entry, exists := s.entries[key] + if !exists { + return false + } + if !entry.lastFailure.IsZero() && now.Sub(entry.lastFailure) > openAIModelTransientFailureWindow { + delete(s.entries, key) + return false + } + entry.lastTouched = now + s.entries[key] = entry + return !entry.blockUntil.IsZero() && now.Before(entry.blockUntil) +} + +func (s *openAIAccountModelTransientState) size() int { + if s == nil { + return 0 + } + s.mu.Lock() + defer s.mu.Unlock() + return len(s.entries) +} + +func (s *openAIAccountModelTransientState) evictOldestLocked() { + if len(s.entries) < s.maxEntries { + return + } + var oldestKey openAIAccountModelKey + var oldestTime time.Time + found := false + for key, entry := range s.entries { + if !found || entry.lastTouched.Before(oldestTime) { + oldestKey = key + oldestTime = entry.lastTouched + found = true + } + } + if found { + delete(s.entries, oldestKey) + } +} diff --git a/backend/internal/service/openai_account_model_transient_test.go b/backend/internal/service/openai_account_model_transient_test.go new file mode 100644 index 0000000000..bd8d92f4d2 --- /dev/null +++ b/backend/internal/service/openai_account_model_transient_test.go @@ -0,0 +1,130 @@ +package service + +import ( + "fmt" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestOpenAIModelTransient_FirstFailureDoesNotCreateLongBlock(t *testing.T) { + state := newOpenAIAccountModelTransientState(128) + now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC) + + decision := state.recordFailure(35, "gpt-5.5", now) + + assert.Equal(t, 1, decision.FailureStreak) + assert.Zero(t, decision.Cooldown) + assert.False(t, state.isBlocked(35, "gpt-5.5", now)) +} + +func TestOpenAIModelTransient_SecondFailureCreatesShortModelBlock(t *testing.T) { + state := newOpenAIAccountModelTransientState(128) + now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC) + state.recordFailure(35, "gpt-5.5", now) + + decision := state.recordFailure(35, "gpt-5.5", now.Add(time.Second)) + + assert.Equal(t, 2, decision.FailureStreak) + assert.Equal(t, openAIModelTransientShortCooldown, decision.Cooldown) + assert.True(t, state.isBlocked(35, "gpt-5.5", now.Add(2*time.Second))) + assert.False(t, state.isBlocked(35, "gpt-5.5", now.Add(openAIModelTransientShortCooldown+2*time.Second))) +} + +func TestOpenAIModelTransient_ThirdFailureCreatesFortyFiveSecondModelBlock(t *testing.T) { + state := newOpenAIAccountModelTransientState(128) + now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC) + state.recordFailure(35, "gpt-5.5", now) + state.recordFailure(35, "gpt-5.5", now.Add(time.Second)) + + decision := state.recordFailure(35, "gpt-5.5", now.Add(2*time.Second)) + + assert.Equal(t, 3, decision.FailureStreak) + assert.Equal(t, 45*time.Second, decision.Cooldown) + assert.True(t, state.isBlocked(35, "gpt-5.5", now.Add(40*time.Second))) + assert.False(t, state.isBlocked(35, "gpt-5.5", now.Add(48*time.Second))) +} + +func TestOpenAIModelTransient_BlockIsIsolatedByModel(t *testing.T) { + state := newOpenAIAccountModelTransientState(128) + now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC) + state.recordFailure(35, "gpt-5.6-terra", now) + state.recordFailure(35, "GPT-5.6-TERRA", now.Add(time.Second)) + + assert.True(t, state.isBlocked(35, "gpt-5.6-terra", now.Add(2*time.Second))) + assert.False(t, state.isBlocked(35, "gpt-5.5", now.Add(2*time.Second))) + assert.False(t, state.isBlocked(47, "gpt-5.6-terra", now.Add(2*time.Second))) +} + +func TestOpenAIModelTransient_SuccessClearsStreakAndBlock(t *testing.T) { + state := newOpenAIAccountModelTransientState(128) + now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC) + state.recordFailure(35, "gpt-5.5", now) + state.recordFailure(35, "gpt-5.5", now.Add(time.Second)) + require.True(t, state.isBlocked(35, "gpt-5.5", now.Add(2*time.Second))) + + state.recordSuccess(35, "gpt-5.5") + + assert.False(t, state.isBlocked(35, "gpt-5.5", now.Add(2*time.Second))) + decision := state.recordFailure(35, "gpt-5.5", now.Add(3*time.Second)) + assert.Equal(t, 1, decision.FailureStreak) + assert.Zero(t, decision.Cooldown) +} + +func TestOpenAIModelTransient_StaleStreakExpires(t *testing.T) { + state := newOpenAIAccountModelTransientState(128) + now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC) + state.recordFailure(35, "gpt-5.5", now) + + decision := state.recordFailure(35, "gpt-5.5", now.Add(openAIModelTransientFailureWindow+time.Second)) + + assert.Equal(t, 1, decision.FailureStreak) + assert.Zero(t, decision.Cooldown) +} + +func TestOpenAIModelTransient_IgnoresInvalidKeys(t *testing.T) { + state := newOpenAIAccountModelTransientState(128) + now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC) + + assert.Zero(t, state.recordFailure(0, "gpt-5.5", now).FailureStreak) + assert.Zero(t, state.recordFailure(35, " ", now).FailureStreak) + assert.False(t, state.isBlocked(0, "gpt-5.5", now)) + assert.False(t, state.isBlocked(35, "", now)) + assert.Equal(t, 0, state.size()) +} + +func TestOpenAIModelTransient_IgnoresOversizedModelKey(t *testing.T) { + state := newOpenAIAccountModelTransientState(128) + now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC) + model := strings.Repeat("m", openAIModelTransientMaxModelBytes+1) + + decision := state.recordFailure(35, model, now) + + assert.Zero(t, decision.FailureStreak) + assert.False(t, state.isBlocked(35, model, now)) + assert.Equal(t, 0, state.size()) +} + +func TestOpenAIModelTransient_StateIsBoundedAndConcurrencySafe(t *testing.T) { + const maxEntries = 16 + state := newOpenAIAccountModelTransientState(maxEntries) + now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC) + var wg sync.WaitGroup + + for i := 0; i < 128; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + model := fmt.Sprintf("gpt-test-%d", i) + state.recordFailure(int64(i+1), model, now.Add(time.Duration(i)*time.Millisecond)) + _ = state.isBlocked(int64(i+1), model, now.Add(time.Second)) + }(i) + } + wg.Wait() + + assert.LessOrEqual(t, state.size(), maxEntries) +} diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index 00d971ba22..7b3df522f1 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -2,7 +2,9 @@ package service import ( "context" + "log/slog" "net/http" + "strings" "sync" "time" ) @@ -43,7 +45,9 @@ func isOpenAIAccount(account *Account) bool { return account != nil && (account.Platform == PlatformOpenAI || account.Platform == PlatformGrok) } -func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, requestedModel ...string) bool { +// handleOpenAIAccountUpstreamError expects canonicalModel to be the model used +// for scheduling after applying account mapping exactly once. +func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, canonicalModel ...string) bool { stateCtx, cancel := openAIAccountStateContext(ctx) defer cancel() @@ -64,16 +68,43 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont if s == nil || account == nil || s.rateLimitService == nil { return false } - if len(requestedModel) > 0 && s.rateLimitService.HandleUpstreamModelNotFound(stateCtx, account, requestedModel[0], statusCode, responseBody) { + if len(canonicalModel) > 0 && s.rateLimitService.HandleUpstreamModelNotFound(stateCtx, account, canonicalModel[0], statusCode, responseBody) { return true } shouldDisable := s.rateLimitService.HandleUpstreamError(stateCtx, account, statusCode, headers, responseBody) if shouldDisable { s.BlockAccountScheduling(account, time.Time{}, "upstream_disable") } + if !shouldDisable && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey && shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) { + model := "" + if len(canonicalModel) > 0 { + model = canonicalModel[0] + } + decision := s.recordOpenAIAccountModelTransientFailure(account, model, time.Now()) + if decision.FailureStreak > 0 { + slog.Warn("openai_model_transient_state", + "account_id", account.ID, + "model", openAIAccountModelTransientModel(model), + "failure_streak", decision.FailureStreak, + "cooldown_ms", decision.Cooldown.Milliseconds(), + "block_scope", "account_model", + ) + } + } return shouldDisable } +func shouldCooldownOpenAITransientUpstreamError(statusCode int, responseBody []byte) bool { + switch statusCode { + case http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout, 520, 521, 522, 523, 524: + return true + case http.StatusBadRequest: + return isOpenAITransientProcessingError(statusCode, "", responseBody) + default: + return false + } +} + func (s *OpenAIGatewayService) markOpenAIOAuth429RateLimited(ctx context.Context, account *Account, headers http.Header, responseBody []byte) { if s == nil || !isOpenAIOAuthAccount(account) { return @@ -191,6 +222,68 @@ func (s *OpenAIGatewayService) isOpenAIAccountRuntimeBlocked(account *Account) b return false } +func (s *OpenAIGatewayService) getOpenAIAccountModelTransientState() *openAIAccountModelTransientState { + if s == nil { + return nil + } + s.openaiModelTransientOnce.Do(func() { + if s.openaiModelTransient == nil { + s.openaiModelTransient = newOpenAIAccountModelTransientState(openAIModelTransientDefaultMax) + } + }) + return s.openaiModelTransient +} + +func canonicalOpenAIAccountSchedulingModel(account *Account, requestedModel string) string { + model := strings.TrimSpace(requestedModel) + if account == nil || model == "" { + return model + } + if mapped := strings.TrimSpace(account.GetMappedModel(model)); mapped != "" { + return mapped + } + return model +} + +func openAIAccountModelTransientModel(canonicalModel string) string { + return normalizeOpenAIAccountModelTransientModel(canonicalModel) +} + +func (s *OpenAIGatewayService) recordOpenAIAccountModelTransientFailure(account *Account, canonicalModel string, now time.Time) openAIAccountModelTransientDecision { + if s == nil || account == nil { + return openAIAccountModelTransientDecision{} + } + state := s.getOpenAIAccountModelTransientState() + if state == nil { + return openAIAccountModelTransientDecision{} + } + return state.recordFailure(account.ID, openAIAccountModelTransientModel(canonicalModel), now) +} + +func (s *OpenAIGatewayService) clearOpenAIAccountModelTransientState(accountID int64, model string) { + state := s.getOpenAIAccountModelTransientState() + if state == nil { + return + } + state.recordSuccess(accountID, model) +} + +func (s *OpenAIGatewayService) isOpenAIAccountModelRuntimeBlocked(account *Account, requestedModel string) bool { + if s == nil || account == nil { + return false + } + state := s.getOpenAIAccountModelTransientState() + if state == nil { + return false + } + canonicalModel := canonicalOpenAIAccountSchedulingModel(account, requestedModel) + return state.isBlocked(account.ID, openAIAccountModelTransientModel(canonicalModel), time.Now()) +} + +func (s *OpenAIGatewayService) isOpenAIAccountRequestRuntimeBlocked(account *Account, requestedModel string) bool { + return s != nil && (s.isOpenAIAccountRuntimeBlocked(account) || s.isOpenAIAccountModelRuntimeBlocked(account, requestedModel)) +} + func (s *OpenAIGatewayService) recordOpenAIOAuth429() { if s == nil { return diff --git a/backend/internal/service/openai_account_runtime_transient_test.go b/backend/internal/service/openai_account_runtime_transient_test.go new file mode 100644 index 0000000000..140e15d490 --- /dev/null +++ b/backend/internal/service/openai_account_runtime_transient_test.go @@ -0,0 +1,116 @@ +package service + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type transientCooldownAccountRepo struct { + AccountRepository +} + +func (transientCooldownAccountRepo) SetOverloaded(context.Context, int64, time.Time) error { + return nil +} + +func TestHandleOpenAITransientError_BlocksOnlyRequestedModel(t *testing.T) { + svc := &OpenAIGatewayService{} + svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil) + account := &Account{ + ID: 5105, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + } + + firstShouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadGateway, http.Header{}, []byte(`{"error":{"message":"Upstream request failed","type":"upstream_error"}}`), "gpt-5.5") + secondShouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadGateway, http.Header{}, []byte(`{"error":{"message":"Upstream request failed","type":"upstream_error"}}`), "gpt-5.5") + + require.False(t, firstShouldDisable) + require.False(t, secondShouldDisable) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5")) + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.6-terra")) +} + +func TestHandleOpenAITransientError_TransientStatusesUseModelScope(t *testing.T) { + for _, statusCode := range []int{http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout, 520, 521, 522, 523, 524} { + t.Run(http.StatusText(statusCode), func(t *testing.T) { + svc := &OpenAIGatewayService{} + svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil) + account := &Account{ + ID: int64(5100 + statusCode), + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + } + + firstShouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, statusCode, http.Header{}, []byte(`{"error":{"message":"temporary upstream failure"}}`), "gpt-5.5") + secondShouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, statusCode, http.Header{}, []byte(`{"error":{"message":"temporary upstream failure"}}`), "gpt-5.5") + + require.False(t, firstShouldDisable) + require.False(t, secondShouldDisable) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account), "status %d must not block the whole account", statusCode) + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5"), "status %d should block the failing model", statusCode) + }) + } +} + +func TestHandleOpenAITransientError_529RemainsOverloadOnly(t *testing.T) { + require.False(t, shouldCooldownOpenAITransientUpstreamError(529, []byte(`{"error":{"message":"overloaded"}}`))) +} + +func TestHandleOpenAITransientError_CanonicalModelIsNotMappedTwice(t *testing.T) { + svc := &OpenAIGatewayService{} + svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil) + account := &Account{ + ID: 5107, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "model_mapping": map[string]any{ + "public-alias": "upstream-a", + "upstream-a": "upstream-b", + }, + }, + } + canonicalModel := account.GetMappedModel("public-alias") + require.Equal(t, "upstream-a", canonicalModel) + + for range 2 { + svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadGateway, http.Header{}, []byte(`{"error":{"message":"temporary upstream failure"}}`), canonicalModel) + } + + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "public-alias")) + svc.ReportOpenAIAccountScheduleResult(account.ID, canonicalModel, true, nil) + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "public-alias")) +} + +func TestHandleOpenAITransientError_DoesNotBlockParameter400(t *testing.T) { + svc := &OpenAIGatewayService{} + svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil) + account := &Account{ + ID: 5103, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + } + + shouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadRequest, http.Header{}, []byte(`{"error":{"message":"Invalid type for input[0].arguments"}}`), "gpt-5.5") + + require.False(t, shouldDisable) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5")) +} + +func TestHandleOpenAITransientError_HardDisableStillBlocksWholeAccount(t *testing.T) { + svc := &OpenAIGatewayService{} + account := &Account{ID: 5106, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + + svc.BlockAccountScheduling(account, time.Now().Add(time.Minute), "upstream_disable") + + require.True(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.5")) + require.True(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.6-sol")) +} diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 4da30819e3..9b6f6015dc 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -1301,7 +1301,7 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance( if !account.IsSchedulable() || account.Platform != normalizeOpenAICompatiblePlatform(req.Platform) || !account.IsOpenAICompatible() { continue } - if s.service.isOpenAIAccountRuntimeBlocked(account) { + if s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) { continue } // require_privacy_set: 跳过 privacy 未设置的账号并标记异常 @@ -1607,7 +1607,7 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatible(ctx context.C if account == nil { return false } - if s != nil && s.service != nil && s.service.isOpenAIAccountRuntimeBlocked(account) { + if s != nil && s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) { return false } // Quota auto-pause must be evaluated during the initial filter too. Without it the @@ -2111,7 +2111,10 @@ func (s *OpenAIGatewayService) isOpenAIAccountTransportCompatible(account *Accou return s.getOpenAIWSProtocolResolver().Resolve(account).Transport == requiredTransport } -func (s *OpenAIGatewayService) ReportOpenAIAccountScheduleResult(accountID int64, success bool, firstTokenMs *int) { +func (s *OpenAIGatewayService) ReportOpenAIAccountScheduleResult(accountID int64, model string, success bool, firstTokenMs *int) { + if success { + s.clearOpenAIAccountModelTransientState(accountID, normalizeOpenAIAccountModelTransientModel(model)) + } scheduler := s.getOpenAIAccountScheduler(context.Background()) if scheduler == nil { return diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index 70329077e8..931b8f9108 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -1098,7 +1098,7 @@ func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics_DisabledNoOp(t *test svc := &OpenAIGatewayService{} ttft := 120 - svc.ReportOpenAIAccountScheduleResult(10, true, &ttft) + svc.ReportOpenAIAccountScheduleResult(10, "", true, &ttft) svc.RecordOpenAIAccountSwitch() snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics() @@ -2239,6 +2239,30 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_UsesAccountPriorityWith } } +func TestOpenAIAccountScheduler_SkipsAccountBlockedForRequestedModel(t *testing.T) { + now := time.Now() + account := &Account{ID: 21633, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(128)} + svc.openaiModelTransient.recordFailure(account.ID, "gpt-5.5", now) + svc.openaiModelTransient.recordFailure(account.ID, "gpt-5.5", now.Add(time.Millisecond)) + scheduler := &defaultOpenAIAccountScheduler{service: svc} + + require.False(t, scheduler.isAccountRequestCompatible(context.Background(), account, OpenAIAccountScheduleRequest{RequestedModel: "gpt-5.5"})) + require.True(t, scheduler.isAccountRequestCompatible(context.Background(), account, OpenAIAccountScheduleRequest{RequestedModel: "gpt-5.6-sol"})) +} + +func TestReportOpenAIAccountScheduleResult_SuccessClearsModelTransientState(t *testing.T) { + svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(128)} + now := time.Now() + svc.openaiModelTransient.recordFailure(21636, "gpt-5.5", now) + svc.openaiModelTransient.recordFailure(21636, "gpt-5.5", now.Add(time.Millisecond)) + require.True(t, svc.openaiModelTransient.isBlocked(21636, "gpt-5.5", now.Add(2*time.Millisecond))) + + svc.ReportOpenAIAccountScheduleResult(21636, "gpt-5.5", true, nil) + + require.False(t, svc.openaiModelTransient.isBlocked(21636, "gpt-5.5", now.Add(2*time.Millisecond))) +} + func TestDefaultOpenAIAccountScheduler_ShouldEscapeStickyAccount_ThresholdBoundary(t *testing.T) { stats := newOpenAIAccountRuntimeStats() accountID := int64(21501) @@ -2694,7 +2718,7 @@ func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics(t *testing.T) { selection, _, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_metrics", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false) require.NoError(t, err) require.NotNil(t, selection) - svc.ReportOpenAIAccountScheduleResult(account.ID, true, intPtrForTest(120)) + svc.ReportOpenAIAccountScheduleResult(account.ID, "", true, intPtrForTest(120)) svc.RecordOpenAIAccountSwitch() snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics() @@ -3031,7 +3055,7 @@ func TestOpenAIGatewayService_SchedulerWrappersAndDefaults(t *testing.T) { svc := &OpenAIGatewayService{} ttft := 120 - svc.ReportOpenAIAccountScheduleResult(10, true, &ttft) + svc.ReportOpenAIAccountScheduleResult(10, "", true, &ttft) svc.RecordOpenAIAccountSwitch() snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics() require.Equal(t, OpenAIAccountSchedulerMetricsSnapshot{}, snapshot) diff --git a/backend/internal/service/openai_alpha_search.go b/backend/internal/service/openai_alpha_search.go index 50d37f95bb..d220cb7a1e 100644 --- a/backend/internal/service/openai_alpha_search.go +++ b/backend/internal/service/openai_alpha_search.go @@ -3,6 +3,7 @@ package service import ( "bytes" "context" + "encoding/json" "fmt" "io" "net/http" @@ -39,21 +40,38 @@ func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Co if upstreamModel != "" && upstreamModel != requestedModel { body = ReplaceModelInBody(body, upstreamModel) } + sanitizedBody, err := sanitizeOpenAIAlphaSearchBody(body) + if err != nil { + return nil, fmt.Errorf("sanitize alpha search request body: %w", err) + } + body = sanitizedBody token, _, err := s.GetAccessToken(ctx, account) if err != nil { return nil, err } - req, err := s.buildOpenAIAlphaSearchRequest(ctx, c, account, body, token) - if err != nil { - return nil, err - } - proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { proxyURL = account.Proxy.URL() } + if err := s.ensureOpenAIAlphaSearchAuthMetadata(ctx, account, token, proxyURL); err != nil { + return nil, err + } + + // Codex Personal Access Token（at-...）目前可访问 ChatGPT Codex + // /responses，但会被 standalone /alpha/search 的 access enforcement + // 拒绝为 no_matching_rule。对 PAT 账号使用等价的 hosted web_search + // Responses 路径兜底，避免把可用账号误判为搜索不可用。 + if account.IsOpenAIPersonalAccessToken() { + return s.forwardAlphaSearchViaResponsesWebSearch(ctx, c, account, body, token, proxyURL, requestedModel, upstreamModel) + } + + req, err := s.buildOpenAIAlphaSearchRequest(ctx, c, account, body, token) + if err != nil { + return nil, err + } + upstreamStart := time.Now() resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) @@ -71,7 +89,14 @@ func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Co upstreamMessage := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMessage, respBody) { resp.Body = io.NopCloser(bytes.NewReader(respBody)) - s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + // alpha/search 是独立的工具端点，单次 401 不能证明账号的模型调用 + // 凭据全局失效。若沿用通用 401 逻辑，PAT 会因没有 refresh_token + // 被永久标记为 error；历史导入且缺少 auth_mode 标记的 at- token 也会 + // 漏过 PAT 类型判断。这里仍允许本次请求换号，但不修改任何账号状态； + // 真正的凭据失效由普通 Responses 请求或 whoami 校验判定。 + if shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(resp.StatusCode) { + s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + } return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, @@ -102,16 +127,212 @@ func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Co }, nil } -func (s *OpenAIGatewayService) buildOpenAIAlphaSearchRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token string) (*http.Request, error) { - clientBeta := "" - if c != nil { - clientBeta = c.GetHeader("OpenAI-Beta") +func (s *OpenAIGatewayService) forwardAlphaSearchViaResponsesWebSearch( + ctx context.Context, + c *gin.Context, + account *Account, + alphaBody []byte, + token string, + proxyURL string, + requestedModel string, + upstreamModel string, +) (*OpenAIForwardResult, error) { + if upstreamModel == "" { + upstreamModel = requestedModel } - req, err := s.buildUpstreamRequestOpenAIPassthrough(ctx, c, account, body, token) + responsesBody, err := buildOpenAIAlphaSearchResponsesWebSearchBody(alphaBody, upstreamModel) if err != nil { return nil, err } + req, err := s.buildOpenAIAlphaSearchResponsesWebSearchRequest(ctx, c, account, alphaBody, responsesBody, token) + if err != nil { + return nil, err + } + SetActualOpenAIUpstreamEndpoint(c, "/v1/responses") + upstreamStart := time.Now() + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + if err != nil { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) + if err != nil { + return nil, fmt.Errorf("read alpha search responses fallback response: %w", err) + } + + if resp.StatusCode >= http.StatusBadRequest { + upstreamMessage := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) + if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMessage, respBody) { + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + // 仍按 alpha/search 工具请求处理：PAT 的工具链路失败不能直接永久置错。 + if shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(resp.StatusCode) { + s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + } + return nil, &UpstreamFailoverError{ + StatusCode: resp.StatusCode, + ResponseBody: respBody, + RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + } + } + } + + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + writeOpenAIPassthroughResponseHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) + contentType := resp.Header.Get("Content-Type") + if contentType == "" { + contentType = "application/json" + } + c.Data(resp.StatusCode, contentType, respBody) + return nil, nil + } + + if !account.IsShadow() { + s.UpdateCodexUsageSnapshotFromHeaders(ctx, account.ID, resp.Header) + } + alphaRespBody, err := openAIAlphaSearchResponseFromResponsesSSE(respBody) + if err != nil { + return nil, err + } + c.Data(http.StatusOK, "application/json", alphaRespBody) + return &OpenAIForwardResult{ + RequestID: strings.TrimSpace(resp.Header.Get("x-request-id")), + Model: requestedModel, + UpstreamModel: upstreamModel, + UpstreamEndpoint: "/v1/responses", + ResponseHeaders: resp.Header.Clone(), + Duration: time.Since(upstreamStart), + WebSearchCalls: 1, + }, nil +} + +func (s *OpenAIGatewayService) buildOpenAIAlphaSearchResponsesWebSearchRequest(ctx context.Context, c *gin.Context, account *Account, alphaBody []byte, body []byte, token string) (*http.Request, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, chatgptCodexURL, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI)) + + authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token) + if err != nil { + return nil, fmt.Errorf("build openai authentication headers: %w", err) + } + for key, values := range authHeaders { + for _, value := range values { + req.Header.Add(key, value) + } + } + req.Host = "chatgpt.com" + if err := resolveAndSetOpenAIChatGPTAccountHeaders(ctx, s.accountRepo, req.Header, account); err != nil { + return nil, fmt.Errorf("resolve chatgpt account headers: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("OpenAI-Beta", "responses=experimental") + if turnMetadata := openAIAlphaSearchInboundHeader(c, "X-Codex-Turn-Metadata"); turnMetadata != "" { + req.Header.Set("X-Codex-Turn-Metadata", turnMetadata) + } + if version := openAIAlphaSearchInboundHeader(c, "Version"); version != "" { + req.Header.Set("Version", version) + } else { + req.Header.Set("Version", codexCLIVersion) + } + if originator := openAIAlphaSearchInboundHeader(c, "Originator"); originator != "" { + req.Header.Set("Originator", originator) + } else { + req.Header.Set("Originator", "codex_cli_rs") + } + if customUA := account.GetOpenAIUserAgent(); customUA != "" { + req.Header.Set("User-Agent", customUA) + } else if userAgent := openAIAlphaSearchInboundHeader(c, "User-Agent"); userAgent != "" { + req.Header.Set("User-Agent", userAgent) + } else { + req.Header.Set("User-Agent", codexCLIUserAgent) + } + if s.cfg != nil && s.cfg.Gateway.ForceCodexCLI { + req.Header.Set("User-Agent", codexCLIUserAgent) + } + apiKeyID := getAPIKeyIDFromContext(c) + if sessionID := strings.TrimSpace(gjson.GetBytes(alphaBody, "id").String()); sessionID != "" { + isolated := isolateOpenAISessionID(apiKeyID, sessionID) + req.Header.Set("Session_ID", isolated) + req.Header.Set("Conversation_ID", isolated) + } + s.overrideBrowserUserAgent(ctx, account, req) + enforceCodexIdentityHeaders(req.Header) + account.ApplyHeaderOverrides(req.Header) + return req, nil +} + +func buildOpenAIAlphaSearchResponsesWebSearchBody(alphaBody []byte, model string) ([]byte, error) { + if strings.TrimSpace(model) == "" { + return nil, fmt.Errorf("model is required") + } + tool := map[string]any{"type": "web_search"} + if contextSize := strings.TrimSpace(gjson.GetBytes(alphaBody, "settings.search_context_size").String()); contextSize != "" { + tool["search_context_size"] = contextSize + } + if userLocation := gjson.GetBytes(alphaBody, "settings.user_location"); userLocation.IsObject() { + var loc map[string]any + if err := json.Unmarshal([]byte(userLocation.Raw), &loc); err == nil && len(loc) > 0 { + tool["user_location"] = loc + } + } + payload := map[string]any{ + "model": model, + "stream": true, + "store": false, + "input": []any{ + map[string]any{ + "role": "user", + "content": []any{ + map[string]any{ + "type": "input_text", + "text": openAIAlphaSearchResponsesWebSearchPrompt(alphaBody), + }, + }, + }, + }, + "tools": []any{tool}, + } + return json.Marshal(payload) +} + +func openAIAlphaSearchResponsesWebSearchPrompt(alphaBody []byte) string { + var b strings.Builder + _, _ = b.WriteString("Execute this Codex standalone web.run request for another model.\n") + _, _ = b.WriteString("Use the hosted web_search tool when web/current information is needed.\n") + _, _ = b.WriteString("Return concise source-backed results. Include titles, URLs, dates, and direct answers when available.\n") + if commands := strings.TrimSpace(gjson.GetBytes(alphaBody, "commands").Raw); commands != "" { + _, _ = b.WriteString("\nCommands JSON:\n") + _, _ = b.WriteString(truncateOpenAIAlphaSearchPromptJSON(commands, 12000)) + } + if settings := strings.TrimSpace(gjson.GetBytes(alphaBody, "settings").Raw); settings != "" { + _, _ = b.WriteString("\n\nSearch settings JSON:\n") + _, _ = b.WriteString(truncateOpenAIAlphaSearchPromptJSON(settings, 4000)) + } + if input := strings.TrimSpace(gjson.GetBytes(alphaBody, "input").Raw); input != "" { + _, _ = b.WriteString("\n\nRecent conversation/input JSON:\n") + _, _ = b.WriteString(truncateOpenAIAlphaSearchPromptJSON(input, 8000)) + } + if b.Len() == 0 { + return "Execute the requested web search and return concise source-backed results." + } + return b.String() +} + +func truncateOpenAIAlphaSearchPromptJSON(value string, limit int) string { + value = strings.TrimSpace(value) + if limit <= 0 || len(value) <= limit { + return value + } + return value[:limit] + "\n..." +} + +func (s *OpenAIGatewayService) buildOpenAIAlphaSearchRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token string) (*http.Request, error) { targetURL, err := s.openAIAlphaSearchURL(account) if err != nil { return nil, err @@ -129,19 +350,279 @@ func (s *OpenAIGatewayService) buildOpenAIAlphaSearchRequest(ctx context.Context } parsedURL.RawQuery = query.Encode() } - req.URL = parsedURL + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, parsedURL.String(), bytes.NewReader(body)) + if err != nil { + return nil, err + } + req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI)) + + authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token) + if err != nil { + return nil, fmt.Errorf("build openai authentication headers: %w", err) + } + for key, values := range authHeaders { + for _, value := range values { + req.Header.Add(key, value) + } + } + + req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json") - if clientBeta == "" { - req.Header.Del("OpenAI-Beta") - } - if version := strings.TrimSpace(c.GetHeader("Version")); version != "" { - req.Header.Set("Version", version) - } else if account.Type == AccountTypeOAuth { - req.Header.Set("Version", codexCLIVersion) + + if account.Type == AccountTypeOAuth { + req.Host = "chatgpt.com" + if err := resolveAndSetOpenAIChatGPTAccountHeaders(ctx, s.accountRepo, req.Header, account); err != nil { + return nil, fmt.Errorf("resolve chatgpt account headers: %w", err) + } + + if turnMetadata := openAIAlphaSearchInboundHeader(c, "X-Codex-Turn-Metadata"); turnMetadata != "" { + req.Header.Set("X-Codex-Turn-Metadata", turnMetadata) + } + if version := openAIAlphaSearchInboundHeader(c, "Version"); version != "" { + req.Header.Set("Version", version) + } else { + req.Header.Set("Version", codexCLIVersion) + } + if originator := openAIAlphaSearchInboundHeader(c, "Originator"); originator != "" { + req.Header.Set("Originator", originator) + } else { + req.Header.Set("Originator", "codex_cli_rs") + } + if customUA := account.GetOpenAIUserAgent(); customUA != "" { + req.Header.Set("User-Agent", customUA) + } else if userAgent := openAIAlphaSearchInboundHeader(c, "User-Agent"); userAgent != "" { + req.Header.Set("User-Agent", userAgent) + } else { + req.Header.Set("User-Agent", codexCLIUserAgent) + } + if s.cfg != nil && s.cfg.Gateway.ForceCodexCLI { + req.Header.Set("User-Agent", codexCLIUserAgent) + } + s.overrideBrowserUserAgent(ctx, account, req) + enforceCodexIdentityHeaders(req.Header) } + + account.ApplyHeaderOverrides(req.Header) + stripOpenAIAlphaSearchResponsesHeaders(req.Header) return req, nil } +// stripOpenAIAlphaSearchResponsesHeaders 让独立搜索请求与官方 Codex +// SearchClient 的线协议保持一致。alpha/search 不是 /responses 的子请求：官方 +// 客户端仅在 Provider/Auth 基础头之外附加 x-codex-turn-metadata，不发送 +// OpenAI-Beta、会话隔离或 Responses Lite 状态头。originator 与 User-Agent +// 属于官方默认客户端头，必须保留。 +// +// alpha/search 使用专用构造器生成官方 SearchClient 的最小线协议形态； +// 该函数作为最后一道防线，避免账号 header 覆写或后续改动重新带入 +// Responses 专用头，使 PAT 的 alpha/search 被上游按错误认证路径处理。 +func stripOpenAIAlphaSearchResponsesHeaders(headers http.Header) { + if headers == nil { + return + } + for _, key := range []string{ + "OpenAI-Beta", + "Session_ID", + "Conversation_ID", + "X-Codex-Beta-Features", + "X-Codex-Turn-State", + responsesLiteHeaderKey, + } { + headers.Del(key) + } +} + +func openAIAlphaSearchInboundHeader(c *gin.Context, key string) string { + if c == nil { + return "" + } + return strings.TrimSpace(c.GetHeader(key)) +} + +var openAIAlphaSearchUnsupportedBodyFields = [...]string{ + // Codex alpha/search 是 SearchRequest 独立协议，不是 /responses 子请求。 + // 新版 Codex/第三方代理可能把 Responses 公共字段误带到搜索请求里；ChatGPT + // alpha/search 会对这些字段返回 Unknown parameter（例如 prompt_cache_key）。 + "prompt_cache_key", + "prompt_cache_retention", +} + +func sanitizeOpenAIAlphaSearchBody(body []byte) ([]byte, error) { + if len(body) == 0 { + return body, nil + } + var obj map[string]json.RawMessage + if err := json.Unmarshal(body, &obj); err != nil || obj == nil { + return body, nil + } + changed := false + for _, field := range openAIAlphaSearchUnsupportedBodyFields { + if _, ok := obj[field]; ok { + delete(obj, field) + changed = true + } + } + if !changed { + return body, nil + } + out, err := json.Marshal(obj) + if err != nil { + return nil, err + } + return out, nil +} + +func (s *OpenAIGatewayService) ensureOpenAIAlphaSearchAuthMetadata(ctx context.Context, account *Account, token string, proxyURL string) error { + if s == nil || account == nil || !account.IsOpenAIPersonalAccessToken() { + return nil + } + if strings.TrimSpace(account.GetChatGPTAccountID()) != "" { + return nil + } + var oauthService *OpenAIOAuthService + if s.openAITokenProvider != nil { + oauthService = s.openAITokenProvider.openAIOAuthService + } + if oauthService == nil { + return nil + } + tokenInfo, err := oauthService.ValidateCodexPersonalAccessToken(ctx, token, proxyURL) + if err != nil { + return fmt.Errorf("validate Codex PAT metadata for alpha/search: %w", err) + } + credentials := shallowCopyMap(account.Credentials) + for key, value := range oauthService.BuildAccountCredentials(tokenInfo) { + credentials[key] = value + } + credentials = NormalizeOpenAIPersonalAccessTokenCredentials(account, tokenInfo, credentials) + account.Credentials = shallowCopyMap(credentials) + if s.accountRepo != nil { + if err := persistAccountCredentials(ctx, s.accountRepo, account, credentials); err != nil { + return fmt.Errorf("persist Codex PAT metadata for alpha/search: %w", err) + } + } + return nil +} + +func shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(statusCode int) bool { + return statusCode != http.StatusUnauthorized +} + +func openAIAlphaSearchResponseFromResponsesSSE(body []byte) ([]byte, error) { + output, results := parseOpenAIResponsesSSEForAlphaSearch(body) + resp := map[string]any{ + "output": output, + } + if len(results) > 0 { + resp["results"] = results + } + return json.Marshal(resp) +} + +func parseOpenAIResponsesSSEForAlphaSearch(body []byte) (string, []any) { + text := strings.ReplaceAll(string(body), "\r\n", "\n") + var output strings.Builder + var completedResponse any + results := make([]any, 0) + seenURLs := make(map[string]struct{}) + + for _, block := range strings.Split(text, "\n\n") { + data := openAIAlphaSearchSSEData(block) + if data == "" || data == "[DONE]" { + continue + } + var event map[string]any + if err := json.Unmarshal([]byte(data), &event); err != nil { + continue + } + if delta, _ := event["delta"].(string); delta != "" && event["type"] == "response.output_text.delta" { + _, _ = output.WriteString(delta) + } + if event["type"] == "response.completed" { + completedResponse = event["response"] + } + collectOpenAIAlphaSearchURLCitations(event, &results, seenURLs) + } + + out := output.String() + if strings.TrimSpace(out) == "" && completedResponse != nil { + out = extractOpenAIResponsesCompletedText(completedResponse) + collectOpenAIAlphaSearchURLCitations(completedResponse, &results, seenURLs) + } + return out, results +} + +func openAIAlphaSearchSSEData(block string) string { + var lines []string + for _, line := range strings.Split(block, "\n") { + line = strings.TrimRight(line, "\r") + if !strings.HasPrefix(line, "data:") { + continue + } + lines = append(lines, strings.TrimSpace(strings.TrimPrefix(line, "data:"))) + } + return strings.TrimSpace(strings.Join(lines, "\n")) +} + +func extractOpenAIResponsesCompletedText(response any) string { + resp, ok := response.(map[string]any) + if !ok { + return "" + } + outputItems, _ := resp["output"].([]any) + var b strings.Builder + for _, item := range outputItems { + itemMap, ok := item.(map[string]any) + if !ok || itemMap["type"] != "message" { + continue + } + contentItems, _ := itemMap["content"].([]any) + for _, content := range contentItems { + contentMap, ok := content.(map[string]any) + if !ok { + continue + } + if contentMap["type"] == "output_text" { + if text, _ := contentMap["text"].(string); text != "" { + _, _ = b.WriteString(text) + } + } + } + } + return b.String() +} + +func collectOpenAIAlphaSearchURLCitations(value any, results *[]any, seen map[string]struct{}) { + switch typed := value.(type) { + case map[string]any: + if typed["type"] == "url_citation" { + if urlValue, _ := typed["url"].(string); strings.TrimSpace(urlValue) != "" { + urlValue = strings.TrimSpace(urlValue) + if _, exists := seen[urlValue]; !exists { + seen[urlValue] = struct{}{} + result := map[string]any{ + "type": "text_result", + "ref_id": fmt.Sprintf("turn0search%d", len(*results)), + "url": urlValue, + } + if title, _ := typed["title"].(string); strings.TrimSpace(title) != "" { + result["title"] = strings.TrimSpace(title) + } + *results = append(*results, result) + } + } + } + for _, child := range typed { + collectOpenAIAlphaSearchURLCitations(child, results, seen) + } + case []any: + for _, child := range typed { + collectOpenAIAlphaSearchURLCitations(child, results, seen) + } + } +} + func (s *OpenAIGatewayService) openAIAlphaSearchURL(account *Account) (string, error) { if account == nil { return "", fmt.Errorf("account is required") diff --git a/backend/internal/service/openai_alpha_search_test.go b/backend/internal/service/openai_alpha_search_test.go index 458dbe9b5b..6b01356016 100644 --- a/backend/internal/service/openai_alpha_search_test.go +++ b/backend/internal/service/openai_alpha_search_test.go @@ -6,7 +6,9 @@ import ( "io" "net/http" "net/http/httptest" + "strconv" "strings" + "sync/atomic" "testing" "github.com/Wei-Shaw/sub2api/internal/config" @@ -15,6 +17,27 @@ import ( "github.com/tidwall/gjson" ) +type alphaSearchAccountStateRepo struct { + AccountRepository + setErrorCalls int + lastError string +} + +func (r *alphaSearchAccountStateRepo) SetError(_ context.Context, _ int64, errorMsg string) error { + r.setErrorCalls++ + r.lastError = errorMsg + return nil +} + +func alphaSearchResponsesSSE(output string) string { + return "event: response.output_text.delta\n" + + `data: {"type":"response.output_text.delta","delta":` + strconv.Quote(output) + `}` + "\n\n" + + "event: response.output_text.annotation.added\n" + + `data: {"type":"response.output_text.annotation.added","annotation":{"type":"url_citation","url":"https://example.com/news","title":"Example News"}}` + "\n\n" + + "event: response.completed\n" + + `data: {"type":"response.completed","response":{"output":[{"type":"message","content":[{"type":"output_text","text":` + strconv.Quote(output) + `}]}]}}` + "\n\n" +} + func TestForwardAlphaSearchOAuthPreservesWire(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{ @@ -70,6 +93,143 @@ func TestForwardAlphaSearchOAuthPreservesWire(t *testing.T) { require.JSONEq(t, string(body), string(upstream.lastBody)) } +func TestForwardAlphaSearchPATUsesResponsesWebSearchFallback(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{ + "id":"search-session", + "model":"gpt-5.6-sol", + "commands":{"search_query":[{"q":"OpenAI news"}]}, + "prompt_cache_key":"responses-cache-key", + "prompt_cache_retention":"24h" + }`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set("User-Agent", codexCLIUserAgent) + c.Request.Header.Set("Originator", "codex_cli_rs") + c.Request.Header.Set("Version", "0.144.1") + c.Request.Header.Set("OpenAI-Beta", "responses=experimental") + c.Request.Header.Set("Accept-Language", "zh-CN") + c.Request.Header.Set("Authorization", "Bearer client-token") + c.Request.Header.Set("Session_ID", "session-client") + c.Request.Header.Set("Conversation_ID", "conversation-client") + c.Request.Header.Set("X-Codex-Beta-Features", "feature-a") + c.Request.Header.Set("X-Codex-Turn-State", "turn-state") + c.Request.Header.Set(responsesLiteHeaderKey, "true") + c.Request.Header.Set("X-Codex-Turn-Metadata", `{"turn_id":"turn-1"}`) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"req-search"}}, + Body: io.NopCloser(strings.NewReader(alphaSearchResponsesSSE("search result"))), + }} + service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 43, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "at-test-token", + "auth_mode": OpenAIAuthModePersonalAccessToken, + "chatgpt_account_id": "chatgpt-account", + "chatgpt_account_is_fedramp": true, + }, + } + + result, err := service.ForwardAlphaSearch(context.Background(), c, account, body) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, result.WebSearchCalls) + require.Equal(t, "/v1/responses", result.UpstreamEndpoint) + require.Equal(t, http.StatusOK, recorder.Code) + require.JSONEq(t, `{"output":"search result","results":[{"type":"text_result","ref_id":"turn0search0","url":"https://example.com/news","title":"Example News"}]}`, recorder.Body.String()) + require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String()) + require.Equal(t, "Bearer at-test-token", upstream.lastReq.Header.Get("Authorization")) + require.Equal(t, "chatgpt-account", upstream.lastReq.Header.Get("ChatGPT-Account-ID")) + require.Equal(t, "true", upstream.lastReq.Header.Get("X-OpenAI-Fedramp")) + require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type")) + require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept")) + require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta")) + require.Equal(t, "0.144.1", upstream.lastReq.Header.Get("Version")) + require.Equal(t, `{"turn_id":"turn-1"}`, upstream.lastReq.Header.Get("X-Codex-Turn-Metadata")) + require.Equal(t, "codex_cli_rs", upstream.lastReq.Header.Get("Originator")) + require.Empty(t, upstream.lastReq.Header.Get("X-Codex-Beta-Features")) + require.Empty(t, upstream.lastReq.Header.Get("X-Codex-Turn-State")) + require.Empty(t, upstream.lastReq.Header.Get(responsesLiteHeaderKey)) + require.Empty(t, upstream.lastReq.Header.Get("Accept-Language")) + require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_retention").Exists()) + require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(upstream.lastBody, "model").String()) + require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) + require.False(t, gjson.GetBytes(upstream.lastBody, "store").Bool()) + require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String()) + require.Contains(t, gjson.GetBytes(upstream.lastBody, "input.0.content.0.text").String(), `"search_query"`) +} + +func TestForwardAlphaSearchPATBackfillsMissingChatGPTAccountMetadata(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"OpenAI news"}]}}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + var whoamiCalls int32 + whoamiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&whoamiCalls, 1) + require.Equal(t, "Bearer at-test-token", r.Header.Get("Authorization")) + require.Equal(t, "application/json", r.Header.Get("Accept")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "email":"pat@example.com", + "chatgpt_user_id":"user-123", + "chatgpt_account_id":"acct-123", + "chatgpt_plan_type":"plus", + "chatgpt_account_is_fedramp":true + }`)) + })) + defer whoamiServer.Close() + oldWhoamiURL := openAICodexPATWhoamiURL + openAICodexPATWhoamiURL = whoamiServer.URL + defer func() { openAICodexPATWhoamiURL = oldWhoamiURL }() + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"output":"search result"}`)), + }} + oauthService := NewOpenAIOAuthService(nil, nil) + service := &OpenAIGatewayService{ + cfg: &config.Config{}, + httpUpstream: upstream, + openAITokenProvider: NewOpenAITokenProvider(nil, nil, oauthService), + } + account := &Account{ + ID: 45, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "at-test-token", + "auth_mode": OpenAIAuthModePersonalAccessToken, + }, + } + + result, err := service.ForwardAlphaSearch(context.Background(), c, account, body) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, int32(1), atomic.LoadInt32(&whoamiCalls)) + require.Equal(t, "acct-123", upstream.lastReq.Header.Get("ChatGPT-Account-ID")) + require.Equal(t, "true", upstream.lastReq.Header.Get("X-OpenAI-Fedramp")) + require.Equal(t, "acct-123", account.Credentials["chatgpt_account_id"]) + require.Equal(t, "user-123", account.Credentials["chatgpt_user_id"]) + require.Equal(t, OpenAIAuthModePersonalAccessToken, account.Credentials["auth_mode"]) +} + func TestForwardAlphaSearchAPIKeyMapsModelAndPassesThroughError(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`) @@ -143,3 +303,98 @@ func TestForwardAlphaSearchReturnsFailoverBeforeWriting(t *testing.T) { require.False(t, c.Writer.Written()) require.Empty(t, recorder.Body.String()) } + +func TestForwardAlphaSearchUnauthorizedDoesNotMarkAccountError(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body)) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"detail":"Unauthorized"}`)), + }} + repo := &alphaSearchAccountStateRepo{} + cfg := &config.Config{} + service := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: upstream, + accountRepo: repo, + rateLimitService: NewRateLimitService(repo, nil, cfg, nil, nil), + } + account := &Account{ + ID: 44, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + // 刻意不设置 auth_mode：覆盖历史上把 at- token 当普通 OAuth 导入的账号。 + "access_token": "at-test-token", + "chatgpt_account_id": "chatgpt-account", + }, + } + + result, err := service.ForwardAlphaSearch(context.Background(), c, account, body) + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusUnauthorized, failoverErr.StatusCode) + require.Zero(t, repo.setErrorCalls) + require.Empty(t, repo.lastError) + require.False(t, c.Writer.Written()) +} + +func TestForwardAlphaSearchPATResponsesFallbackUnauthorizedDoesNotMarkAccountError(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body)) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"detail":"Unauthorized"}`)), + }} + repo := &alphaSearchAccountStateRepo{} + cfg := &config.Config{} + service := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: upstream, + accountRepo: repo, + rateLimitService: NewRateLimitService(repo, nil, cfg, nil, nil), + } + account := &Account{ + ID: 46, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "at-test-token", + "auth_mode": OpenAIAuthModePersonalAccessToken, + "chatgpt_account_id": "chatgpt-account", + }, + } + + result, err := service.ForwardAlphaSearch(context.Background(), c, account, body) + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusUnauthorized, failoverErr.StatusCode) + require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String()) + require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept")) + require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta")) + require.Zero(t, repo.setErrorCalls) + require.Empty(t, repo.lastError) + require.False(t, c.Writer.Written()) +} + +func TestShouldApplyOpenAIAlphaSearchAccountErrorSideEffects(t *testing.T) { + require.False(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusUnauthorized)) + require.True(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusForbidden)) + require.True(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusTooManyRequests)) +} diff --git a/backend/internal/service/openai_gateway_cc_pipeline.go b/backend/internal/service/openai_gateway_cc_pipeline.go index b38de924a3..52658b2729 100644 --- a/backend/internal/service/openai_gateway_cc_pipeline.go +++ b/backend/internal/service/openai_gateway_cc_pipeline.go @@ -115,12 +115,13 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( if account.Platform != PlatformGrok { s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) } - return &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - ResponseHeaders: resp.Header.Clone(), - RetryableOnSameAccount: account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), - } + return newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), + ) } // openAIChatCompletionsTargetURL 解析账号的（非 Grok）Chat Completions 上游端点。 diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 5259aa843c..bd8e31ad64 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -391,9 +391,8 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco if maxOutputTokens.Exists() { switch account.Platform { case PlatformOpenAI: - if account.Type == AccountTypeAPIKey { - markPatchDelete("max_output_tokens") - } + // Preserve Responses-native output limits unless the selected upstream + // explicitly rejects the field in the bounded HTTP retry loop below. case PlatformAnthropic: decoded, decodeErr := ensureReqBody() if decodeErr != nil { @@ -744,6 +743,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco httpInvalidEncryptedContentRetryTried := false agentTaskRecoveryTried := false + rejectedFieldRetryState := newOpenAIResponsesRejectedFieldRetryState(body) for { // Build upstream request upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) @@ -830,11 +830,21 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco return nil, fmt.Errorf("serialize invalid_encrypted_content retry body: %w", err) } httpInvalidEncryptedContentRetryTried = true + rejectedFieldRetryState.remember(body) logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Retrying non-WSv2 request once after invalid_encrypted_content (account: %s)", account.Name) continue } logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Skip non-WSv2 invalid_encrypted_content retry because encrypted reasoning items are missing (account: %s)", account.Name) } + if retryBody, reason, changed, retryErr := normalizeOpenAIResponsesRejectedFieldRetryBody(resp.StatusCode, body, respBody); retryErr != nil { + return nil, fmt.Errorf("normalize rejected Responses field retry body: %w", retryErr) + } else if changed && rejectedFieldRetryState.Allow(retryBody) { + body = retryBody + requestView = newOpenAIRequestView(body) + reqBody = nil + logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Retrying non-WSv2 request after %s (account: %s)", reason, account.Name) + continue + } if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMsg, respBody) { upstreamDetail := "" if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { @@ -856,11 +866,13 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco }) s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), - } + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), + ) } return s.handleErrorResponse(ctx, resp, c, account, body, billingModel) } diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index bbf80c97c9..a17d050afe 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -461,6 +461,9 @@ func shouldFailoverOpenAIPassthroughResponse(account *Account, statusCode int, r if isOpenAIContextWindowError("", responseBody) { return false } + if isOpenAIRequestBodyTooLargeError(statusCode, "", responseBody) { + return true + } switch statusCode { case http.StatusTooManyRequests, 529: return true @@ -576,7 +579,8 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) logOpenAIInstructionsRequiredDebug(ctx, c, account, resp.StatusCode, upstreamMsg, requestBody, body) reqModel, _, _ := extractOpenAIRequestMetaFromBody(requestBody) - _ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, reqModel) + canonicalModel := canonicalOpenAIAccountSchedulingModel(account, reqModel) + _ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, canonicalModel) appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, AccountID: account.ID, @@ -589,12 +593,13 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( Detail: upstreamDetail, UpstreamResponseBody: upstreamDetail, }) - return &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: body, - ResponseHeaders: resp.Header.Clone(), - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), - } + return newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + body, + upstreamMsg, + account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + ) } func (s *OpenAIGatewayService) handleErrorResponsePassthrough( @@ -637,7 +642,8 @@ func (s *OpenAIGatewayService) handleErrorResponsePassthrough( // 刚被限流的账号。cyber 例外：不冷却账号。 if !cyberHit { reqModel, _, _ := extractOpenAIRequestMetaFromBody(requestBody) - _ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, reqModel) + canonicalModel := canonicalOpenAIAccountSchedulingModel(account, reqModel) + _ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, canonicalModel) } appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index 436e431163..8857bdd527 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -1015,8 +1015,16 @@ func openAIUsageFromGJSON(value gjson.Result) (OpenAIUsage, bool) { if imageOutputTokens == 0 { imageOutputTokens = value.Get("completion_tokens_details.image_tokens").Int() } + // 图片输入 token（如 gpt-image-2 的 /v1/images/edits 带图请求）， + // 上游在 input_tokens_details.image_tokens 单独回传，用于图/文输入分价计费。 + // 普通文本请求该字段为 0，走原路径行为不变。 + imageInputTokens := firstPositiveGJSONInt( + value.Get("input_tokens_details.image_tokens"), + value.Get("prompt_tokens_details.image_tokens"), + ) return OpenAIUsage{ InputTokens: int(inputTokens), + ImageInputTokens: imageInputTokens, OutputTokens: int(outputTokens), CacheCreationInputTokens: cacheCreationTokens, CacheReadInputTokens: cacheReadTokens, diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index 0457b3a560..a8ed5fd2b5 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -659,7 +659,7 @@ func (s *OpenAIGatewayService) tryStickySessionHit(ctx context.Context, groupID _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash) return nil } - if s.isOpenAIAccountRuntimeBlocked(account) { + if s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) { _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash) return nil } @@ -863,7 +863,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash) } else if !s.openAIAccountMatchesSchedulingGroup(account, groupID) { _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash) - } else if s.isOpenAIAccountRuntimeBlocked(account) { + } else if s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) { _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash) } else if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel, requireCompact) { _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash) @@ -926,7 +926,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex if !parentHealthyForShadow(acc, parentLookupL2) { continue } - if s.isOpenAIAccountRuntimeBlocked(acc) { + if s.isOpenAIAccountRequestRuntimeBlocked(acc, requestedModel) { continue } if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, acc, requestedModel, requireCompact) { @@ -1180,7 +1180,7 @@ func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccount(ctx context. if !parentHealthyForShadow(fresh, s.parentAccountLookup(ctx)) { return nil } - if s.isOpenAIAccountRuntimeBlocked(fresh) { + if s.isOpenAIAccountRequestRuntimeBlocked(fresh, requestedModel) { return nil } return fresh @@ -1228,7 +1228,7 @@ func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDB(ctx context.Co if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) { return nil } - if s.isOpenAIAccountRuntimeBlocked(latest) { + if s.isOpenAIAccountRequestRuntimeBlocked(latest, requestedModel) { return nil } return latest diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 956f167d1b..a7fc33b5e1 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -236,22 +236,25 @@ type OpenAIForwardResult struct { ServiceTier *string // ReasoningEffort is extracted from request body (reasoning.effort) or derived from model suffix. // Stored for usage records display; nil means not provided / not applicable. - ReasoningEffort *string - Stream bool - OpenAIWSMode bool - ResponseHeaders http.Header - Duration time.Duration - FirstTokenMs *int - ClientDisconnect bool - ImageCount int - ImageSize string - ImageInputSize string - ImageOutputSize string - ImageOutputSizes []string - ImageSizeSource string - ImageSizeBreakdown map[string]int - VideoCount int - VideoResolution string + ReasoningEffort *string + Stream bool + OpenAIWSMode bool + // UpstreamTerminalEvent is the normalized terminal event observed on an + // upstream Responses WebSocket turn. Empty preserves legacy/non-WS success. + UpstreamTerminalEvent string + ResponseHeaders http.Header + Duration time.Duration + FirstTokenMs *int + ClientDisconnect bool + ImageCount int + ImageSize string + ImageInputSize string + ImageOutputSize string + ImageOutputSizes []string + ImageSizeSource string + ImageSizeBreakdown map[string]int + VideoCount int + VideoResolution string // VideoDurationSeconds 是提交时请求的生成时长（xAI 按输出秒数计费），已归一化到 1-15 秒。 VideoDurationSeconds int // WebSearchCalls 是 Codex alpha/search 网页搜索调用次数（每次成功请求为 1）。 @@ -262,6 +265,21 @@ type OpenAIForwardResult struct { wsReplayInputExists bool } +// SucceededForScheduling reports whether this result is an upstream success +// that may clear model-scoped transient state. The zero value remains a success +// for existing non-WS callers. +func (r *OpenAIForwardResult) SucceededForScheduling() bool { + if r == nil || !r.OpenAIWSMode || r.UpstreamTerminalEvent == "" { + return true + } + switch r.UpstreamTerminalEvent { + case "response.completed", "response.done": + return true + default: + return false + } +} + // SetActualOpenAIUpstreamEndpoint records the endpoint selected by the current // forwarding attempt. It covers error paths where no OpenAIForwardResult is // available for usage and operations logging. @@ -393,12 +411,14 @@ type OpenAIGatewayService struct { openaiWSStateStoreOnce sync.Once openaiSchedulerOnce sync.Once openaiWSPassthroughDialerOnce sync.Once + openaiModelTransientOnce sync.Once agentIdentityTaskMu sync.Mutex openaiWSPool *openAIWSConnPool openaiWSStateStore OpenAIWSStateStore openaiScheduler OpenAIAccountScheduler openaiWSPassthroughDialer openAIWSClientDialer openaiAccountStats *openAIAccountRuntimeStats + openaiModelTransient *openAIAccountModelTransientState openaiWSFallbackUntil sync.Map // key: int64(accountID), value: time.Time openaiAccountRuntimeBlockUntil sync.Map // key: int64(accountID), value: time.Time @@ -475,6 +495,7 @@ func NewOpenAIGatewayService( userPlatformQuotaRepo: userPlatformQuotaRepo, responseHeaderFilter: compileResponseHeaderFilter(cfg), codexSnapshotThrottle: newAccountWriteThrottle(openAICodexSnapshotPersistMinInterval), + openaiModelTransient: newOpenAIAccountModelTransientState(openAIModelTransientDefaultMax), } if rateLimitService != nil { rateLimitService.SetAccountRuntimeBlocker(svc) diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index 507fb9d5c3..49ef5b5f7e 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -259,6 +259,31 @@ func TestExtractOpenAIResponseIDFromJSONBytes(t *testing.T) { require.Empty(t, extractOpenAIResponseIDFromJSONBytes([]byte(`not-json`))) } +// 复现 #4386：gpt-image-2 /v1/images/edits 的 usage 携带 input_tokens_details.image_tokens， +// 提取器须将图片输入 token 单独填入 ImageInputTokens（此前被丢弃并入 InputTokens 按文本价计费）。 +func TestExtractOpenAIUsage_CapturesImageInputTokens(t *testing.T) { + body := []byte(`{"usage":{"input_tokens":371,"input_tokens_details":{"image_tokens":352,"text_tokens":19},"output_tokens":439,"output_tokens_details":{"image_tokens":439,"text_tokens":0},"total_tokens":810}}`) + usage, ok := extractOpenAIUsageFromJSONBytes(body) + require.True(t, ok) + require.Equal(t, 371, usage.InputTokens) + require.Equal(t, 352, usage.ImageInputTokens) + require.Equal(t, 439, usage.OutputTokens) + require.Equal(t, 439, usage.ImageOutputTokens) + + // prompt_tokens_details 回退路径（部分上游用 prompt_tokens 口径）。 + promptStyle := []byte(`{"usage":{"prompt_tokens":100,"prompt_tokens_details":{"image_tokens":80}}}`) + pu, ok := extractOpenAIUsageFromJSONBytes(promptStyle) + require.True(t, ok) + require.Equal(t, 100, pu.InputTokens) + require.Equal(t, 80, pu.ImageInputTokens) + + // 纯文本请求：无 image_tokens 时 ImageInputTokens 为 0，行为不变。 + textOnly := []byte(`{"usage":{"input_tokens":50,"output_tokens":10}}`) + tu, ok := extractOpenAIUsageFromJSONBytes(textOnly) + require.True(t, ok) + require.Zero(t, tu.ImageInputTokens) +} + func TestOpenAIGatewayService_BindHTTPResponseAccount(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() diff --git a/backend/internal/service/openai_gateway_upstream_errors.go b/backend/internal/service/openai_gateway_upstream_errors.go index d8823fe4ed..a21775308d 100644 --- a/backend/internal/service/openai_gateway_upstream_errors.go +++ b/backend/internal/service/openai_gateway_upstream_errors.go @@ -222,12 +222,55 @@ func (s *OpenAIGatewayService) shouldFailoverOpenAIUpstreamResponse(statusCode i if isOpenAIContextWindowError(upstreamMsg, upstreamBody) { return false } + if isOpenAIRequestBodyTooLargeError(statusCode, upstreamMsg, upstreamBody) { + return true + } if s.shouldFailoverUpstreamError(statusCode) { return true } return isOpenAITransientProcessingError(statusCode, upstreamMsg, upstreamBody) } +// OpenAIRequestBodyTooLargeClientMessage is the fixed downstream message used +// after all account-specific request body limit failovers are exhausted. +const OpenAIRequestBodyTooLargeClientMessage = "Request payload is too large" + +const openAIRequestBodyTooLargeReason = GatewayFailureReason("openai_request_body_too_large") + +func isOpenAIRequestBodyTooLargeError(statusCode int, upstreamMsg string, upstreamBody []byte) bool { + return statusCode == http.StatusRequestEntityTooLarge && !isOpenAIContextWindowError(upstreamMsg, upstreamBody) +} + +func newOpenAIUpstreamFailoverError( + statusCode int, + responseHeaders http.Header, + responseBody []byte, + upstreamMsg string, + retryableOnSameAccount bool, +) *UpstreamFailoverError { + failoverErr := &UpstreamFailoverError{ + StatusCode: statusCode, + ResponseBody: responseBody, + ResponseHeaders: responseHeaders.Clone(), + RetryableOnSameAccount: retryableOnSameAccount, + } + if isOpenAIRequestBodyTooLargeError(statusCode, upstreamMsg, responseBody) { + failoverErr.RetryableOnSameAccount = false + failoverErr.Scope = GatewayFailureScopeAccount + failoverErr.Reason = openAIRequestBodyTooLargeReason + failoverErr.NextAccountAction = NextAccountRetry + failoverErr.ClientStatusCode = http.StatusRequestEntityTooLarge + failoverErr.ClientMessage = OpenAIRequestBodyTooLargeClientMessage + } + return failoverErr +} + +// IsOpenAIRequestBodyTooLarge reports whether another account may accept the +// same request even though the selected account rejected its serialized size. +func (e *UpstreamFailoverError) IsOpenAIRequestBodyTooLarge() bool { + return e != nil && e.Reason == openAIRequestBodyTooLargeReason +} + func marshalOpenAIUpstreamJSON(v any) ([]byte, error) { var buf bytes.Buffer enc := json.NewEncoder(&buf) @@ -262,9 +305,9 @@ func (s *OpenAIGatewayService) readUpstreamErrorBody(resp *http.Response) []byte return body } -func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, responseBody []byte, requestedModel ...string) { - if len(requestedModel) > 0 { - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody, requestedModel[0]) +func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, responseBody []byte, canonicalModel ...string) { + if len(canonicalModel) > 0 { + s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody, canonicalModel[0]) return } s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody) @@ -328,6 +371,27 @@ func (s *OpenAIGatewayService) handleErrorResponse( ) } + if isOpenAIRequestBodyTooLargeError(resp.StatusCode, upstreamMsg, body) { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "failover", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel...) + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + body, + upstreamMsg, + false, + ) + } + if status, errType, errMsg, matched := applyErrorPassthroughRule( c, PlatformOpenAI, @@ -385,6 +449,7 @@ func (s *OpenAIGatewayService) handleErrorResponse( } if reqModel == "" { reqModel, _, _ = extractOpenAIRequestMetaFromBody(requestBody) + reqModel = canonicalOpenAIAccountSchedulingModel(account, reqModel) } shouldDisable := s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, reqModel) kind := "http_error" diff --git a/backend/internal/service/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index 410b4d944d..773b31d625 100644 --- a/backend/internal/service/openai_gateway_usage.go +++ b/backend/internal/service/openai_gateway_usage.go @@ -98,6 +98,18 @@ func (s *OpenAIGatewayService) RecordCyberPolicyUsageLog(ctx context.Context, in } } +// ResolveUserGroupRateMultiplier resolves the same cached multiplier used by OpenAI usage billing. +func (s *OpenAIGatewayService) ResolveUserGroupRateMultiplier(ctx context.Context, userID, groupID int64, groupDefaultMultiplier float64) float64 { + if s == nil { + return groupDefaultMultiplier + } + resolver := s.userGroupRateResolver + if resolver == nil { + resolver = newUserGroupRateResolver(nil, nil, resolveUserGroupRateCacheTTL(s.cfg), nil, "service.openai_gateway") + } + return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier) +} + // RecordUsage records usage and deducts balance func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRecordUsageInput) error { if input == nil { @@ -142,11 +154,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec multiplier = s.cfg.Default.RateMultiplier } if apiKey.GroupID != nil && apiKey.Group != nil { - resolver := s.userGroupRateResolver - if resolver == nil { - resolver = newUserGroupRateResolver(nil, nil, resolveUserGroupRateCacheTTL(s.cfg), nil, "service.openai_gateway") - } - multiplier = resolver.Resolve(ctx, user.ID, *apiKey.GroupID, apiKey.Group.RateMultiplier) + multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, apiKey.Group.RateMultiplier) } // token 倍率叠加高峰因子（token 计费含图片 token，图片按次倍率不受影响）。高峰因子按请求时刻现算， // 不并入上面的 Resolve，以免污染 user:group 倍率缓存。 @@ -254,6 +262,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec OutputTokens: result.Usage.OutputTokens, CacheCreationTokens: result.Usage.CacheCreationInputTokens, CacheReadTokens: result.Usage.CacheReadInputTokens, + ImageInputTokens: result.Usage.ImageInputTokens, ImageOutputTokens: result.Usage.ImageOutputTokens, ImageCount: result.ImageCount, ImageSize: optionalTrimmedStringPtr(result.ImageSize), @@ -271,6 +280,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec } if cost != nil { usageLog.InputCost = cost.InputCost + usageLog.ImageInputCost = cost.ImageInputCost usageLog.OutputCost = cost.OutputCost usageLog.ImageOutputCost = cost.ImageOutputCost usageLog.CacheCreationCost = cost.CacheCreationCost diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index 081b5e741b..fa0b7f8b69 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -455,8 +455,13 @@ func applyOpenAIImagesDefaults(req *OpenAIImagesRequest) { } func isOpenAIImageGenerationModel(model string) bool { + return IsGPTImageGenerationModel(model) || isGrokImageGenerationModel(model) +} + +// IsGPTImageGenerationModel identifies the GPT native image-generation model family. +func IsGPTImageGenerationModel(model string) bool { model = strings.ToLower(strings.TrimSpace(model)) - return strings.HasPrefix(model, "gpt-image-") || isGrokImageGenerationModel(model) + return strings.HasPrefix(model, "gpt-image-") } func isGrokImageGenerationModel(model string) bool { @@ -1146,6 +1151,9 @@ func mergeOpenAIUsage(dst *OpenAIUsage, body []byte) { if parsed.CacheReadInputTokens > 0 { dst.CacheReadInputTokens = parsed.CacheReadInputTokens } + if parsed.ImageInputTokens > 0 { + dst.ImageInputTokens = parsed.ImageInputTokens + } if parsed.ImageOutputTokens > 0 { dst.ImageOutputTokens = parsed.ImageOutputTokens } diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 151de2c4b1..f5b99ac9ca 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -1072,6 +1072,165 @@ func (s *OpenAIGatewayService) tryWriteOpenAIImagesStreamEvent( return true } +func (s *OpenAIGatewayService) parseOpenAIImagesSSEUsageBytes(data []byte, usage *OpenAIUsage) { + s.parseSSEUsageBytes(data, usage) + if usage == nil || !gjson.ValidBytes(data) || gjson.GetBytes(data, "type").String() != "response.completed" { + return + } + if toolUsage, ok := openAIImagesToolUsageFromGJSON(gjson.GetBytes(data, "response.tool_usage.image_gen")); ok { + *usage = toolUsage + } +} + +func openAIImagesToolUsageFromGJSON(value gjson.Result) (OpenAIUsage, bool) { + if !value.Exists() || !value.IsObject() { + return OpenAIUsage{}, false + } + inputTokens, inputOK := boundedJSONNonNegativeInt(value.Get("input_tokens")) + outputTokens, outputOK := boundedJSONNonNegativeInt(value.Get("output_tokens")) + imageOutputTokens, imageOutputOK := boundedJSONNonNegativeInt(value.Get("output_tokens_details.image_tokens")) + if !inputOK || !outputOK || !imageOutputOK { + return OpenAIUsage{}, false + } + return OpenAIUsage{ + InputTokens: inputTokens, + OutputTokens: outputTokens, + ImageOutputTokens: imageOutputTokens, + }, true +} + +// boundedJSONNonNegativeInt parses integral JSON exponent notation without +// invoking an arbitrary-precision parser on an upstream-controlled exponent. +func boundedJSONNonNegativeInt(value gjson.Result) (int, bool) { + if !value.Exists() || value.Type != gjson.Number { + return 0, false + } + raw := value.Raw + if len(raw) == 0 || len(raw) > 64 || raw[0] == '-' { + return 0, false + } + + mantissaEnd := len(raw) + for i, c := range raw { + if c != 'e' && c != 'E' { + continue + } + mantissaEnd = i + break + } + + digits := raw[:mantissaEnd] + fractionDigits := 0 + digitCount := 0 + dotSeen := false + mantissaIsZero := true + for _, c := range digits { + switch { + case c == '.' && !dotSeen: + dotSeen = true + case c >= '0' && c <= '9': + digitCount++ + mantissaIsZero = mantissaIsZero && c == '0' + if dotSeen { + fractionDigits++ + } + default: + return 0, false + } + } + + exponent := 0 + if mantissaEnd < len(raw) { + exponentRaw := raw[mantissaEnd+1:] + negative := false + if len(exponentRaw) > 0 && (exponentRaw[0] == '+' || exponentRaw[0] == '-') { + negative = exponentRaw[0] == '-' + exponentRaw = exponentRaw[1:] + } + if len(exponentRaw) == 0 { + return 0, false + } + for len(exponentRaw) > 1 && exponentRaw[0] == '0' { + exponentRaw = exponentRaw[1:] + } + for _, digit := range exponentRaw { + if digit < '0' || digit > '9' { + return 0, false + } + } + if mantissaIsZero { + return 0, true + } + if len(exponentRaw) > 3 { + return 0, false + } + for _, digit := range exponentRaw { + exponent = exponent*10 + int(digit-'0') + } + if exponent > 100 { + return 0, false + } + if negative { + exponent = -exponent + } + } + + trailingZeros := exponent - fractionDigits + scaleReduction := 0 + if trailingZeros < 0 { + scaleReduction = -trailingZeros + remaining := scaleReduction + allZeros := true + for i := len(digits) - 1; i >= 0; i-- { + if digits[i] == '.' { + continue + } + if digits[i] != '0' { + allZeros = false + if remaining > 0 { + return 0, false + } + } + if remaining > 0 { + remaining-- + } + } + if remaining > 0 { + if allZeros { + return 0, true + } + return 0, false + } + } + + maxInt := int(^uint(0) >> 1) + parsed := 0 + digitsToAccumulate := digitCount - scaleReduction + for _, c := range digits { + if c == '.' { + continue + } + if digitsToAccumulate <= 0 { + break + } + if parsed > (maxInt-int(c-'0'))/10 { + return 0, false + } + parsed = parsed*10 + int(c-'0') + digitsToAccumulate-- + } + if trailingZeros < 0 { + return parsed, true + } + for ; trailingZeros > 0; trailingZeros-- { + if parsed > maxInt/10 { + return 0, false + } + parsed *= 10 + } + return parsed, true +} + func (s *OpenAIGatewayService) handleOpenAIImagesOAuthNonStreamingResponse( resp *http.Response, c *gin.Context, @@ -1085,7 +1244,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthNonStreamingResponse( var usage OpenAIUsage forEachOpenAISSEDataPayload(string(body), func(data []byte) { - s.parseSSEUsageBytes(data, &usage) + s.parseOpenAIImagesSSEUsageBytes(data, &usage) }) results, createdAt, usageRaw, firstMeta, _, err := collectOpenAIImagesFromResponsesBody(body) if err != nil { @@ -1192,7 +1351,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( ms := int(time.Since(startTime).Milliseconds()) firstTokenMs = &ms } - s.parseSSEUsageBytes(dataBytes, &usage) + s.parseOpenAIImagesSSEUsageBytes(dataBytes, &usage) if !gjson.ValidBytes(dataBytes) { return } diff --git a/backend/internal/service/openai_images_test.go b/backend/internal/service/openai_images_test.go index 0bcd68a386..2d43dc7c3d 100644 --- a/backend/internal/service/openai_images_test.go +++ b/backend/internal/service/openai_images_test.go @@ -492,9 +492,24 @@ func TestAccountSupportsOpenAIEndpointCapability(t *testing.T) { } require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions)) + require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch)) require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings)) }) + t.Run("alpha search 仅允许 OpenAI OAuth/PAT 类账号", func(t *testing.T) { + apiKey := &Account{ + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + } + oauth := &Account{ + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + } + + require.False(t, apiKey.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch)) + require.True(t, oauth.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch)) + }) + t.Run("显式列表支持同时声明 chat 和 embeddings", func(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, @@ -518,9 +533,22 @@ func TestAccountSupportsOpenAIEndpointCapability(t *testing.T) { } require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions)) + require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch)) require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings)) }) + t.Run("OAuth 显式列表沿用 chat 能力放行 alpha search", func(t *testing.T) { + account := &Account{ + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "openai_capabilities": []any{"chat_completions"}, + }, + } + + require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch)) + }) + t.Run("显式 map 支持单独关闭 chat 并开启 embeddings", func(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, @@ -631,7 +659,7 @@ func TestOpenAIGatewayServiceForwardImages_OAuthPassesNAndReturnsAllImages(t *te "X-Request-Id": []string{"req_img_123"}, }, Body: io.NopCloser(strings.NewReader( - "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000000,\"usage\":{\"input_tokens\":11,\"output_tokens\":22,\"input_tokens_details\":{\"cached_tokens\":3},\"output_tokens_details\":{\"image_tokens\":7}},\"tool_usage\":{\"image_gen\":{\"images\":3}},\"output\":[{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2UtMQ==\",\"revised_prompt\":\"draw a cat 1\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"},{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2UtMg==\",\"revised_prompt\":\"draw a cat 2\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"},{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2UtMw==\",\"revised_prompt\":\"draw a cat 3\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"}]}}\n\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000000,\"usage\":{\"input_tokens\":11,\"output_tokens\":22,\"input_tokens_details\":{\"cached_tokens\":3},\"output_tokens_details\":{\"image_tokens\":7}},\"tool_usage\":{\"image_gen\":{\"input_tokens\":46,\"output_tokens\":2459,\"output_tokens_details\":{\"image_tokens\":2459},\"images\":3}},\"output\":[{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2UtMQ==\",\"revised_prompt\":\"draw a cat 1\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"},{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2UtMg==\",\"revised_prompt\":\"draw a cat 2\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"},{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2UtMw==\",\"revised_prompt\":\"draw a cat 3\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"}]}}\n\n" + "data: [DONE]\n\n", )), }, @@ -655,9 +683,9 @@ func TestOpenAIGatewayServiceForwardImages_OAuthPassesNAndReturnsAllImages(t *te require.Equal(t, "gpt-image-2", result.Model) require.Equal(t, "gpt-image-2", result.UpstreamModel) require.Equal(t, 3, result.ImageCount) - require.Equal(t, 11, result.Usage.InputTokens) - require.Equal(t, 22, result.Usage.OutputTokens) - require.Equal(t, 7, result.Usage.ImageOutputTokens) + require.Equal(t, 46, result.Usage.InputTokens) + require.Equal(t, 2459, result.Usage.OutputTokens) + require.Equal(t, 2459, result.Usage.ImageOutputTokens) require.NotNil(t, upstream.lastReq) require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String()) @@ -688,6 +716,81 @@ func TestOpenAIGatewayServiceForwardImages_OAuthPassesNAndReturnsAllImages(t *te require.Equal(t, "draw a cat 3", gjson.Get(rec.Body.String(), "data.2.revised_prompt").String()) } +func TestParseOpenAIImagesSSEUsageBytes_ToolUsagePrecedenceAndFallback(t *testing.T) { + svc := &OpenAIGatewayService{} + fallback := OpenAIUsage{InputTokens: 3, OutputTokens: 4, ImageOutputTokens: 2} + tests := []struct { + name string + toolUsage string + want OpenAIUsage + }{ + { + name: "valid tool usage takes atomic precedence", + toolUsage: `{"input_tokens":4.6e1,"output_tokens":2459e0,"output_tokens_details":{"image_tokens":24590e-1}}`, + want: OpenAIUsage{InputTokens: 46, OutputTokens: 2459, ImageOutputTokens: 2459}, + }, + {name: "absent", want: fallback}, + {name: "malformed field", toolUsage: `{"input_tokens":"46","output_tokens":2459,"output_tokens_details":{"image_tokens":2459}}`, want: fallback}, + {name: "fractional field", toolUsage: `{"input_tokens":46,"output_tokens":2459.5,"output_tokens_details":{"image_tokens":2459}}`, want: fallback}, + {name: "negative field", toolUsage: `{"input_tokens":46,"output_tokens":2459,"output_tokens_details":{"image_tokens":-1}}`, want: fallback}, + {name: "overflow field", toolUsage: `{"input_tokens":46,"output_tokens":9223372036854775808,"output_tokens_details":{"image_tokens":2459}}`, want: fallback}, + {name: "incomplete object", toolUsage: `{"input_tokens":46,"output_tokens":2459}`, want: fallback}, + {name: "hostile huge exponent", toolUsage: `{"input_tokens":1e1000000000,"output_tokens":2459,"output_tokens_details":{"image_tokens":2459}}`, want: fallback}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + toolUsageField := "" + if tt.toolUsage != "" { + toolUsageField = `,"tool_usage":{"image_gen":` + tt.toolUsage + `}` + } + payload := []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":3,"output_tokens":4,"output_tokens_details":{"image_tokens":2}}` + toolUsageField + `}}`) + var got OpenAIUsage + svc.parseOpenAIImagesSSEUsageBytes(payload, &got) + require.Equal(t, tt.want, got) + }) + } +} + +func TestParseOpenAIImagesSSEUsageBytes_MalformedCompletedDoesNotOverrideUsage(t *testing.T) { + svc := &OpenAIGatewayService{} + var usage OpenAIUsage + + svc.parseOpenAIImagesSSEUsageBytes([]byte(`{"type":"response.output_item.done","item":{"type":"image_generation_call","result":"aW1hZ2U="}}`), &usage) + svc.parseOpenAIImagesSSEUsageBytes([]byte(`{"type":"response.completed","response":{"usage":{"input_tokens":3,"output_tokens":4,"output_tokens_details":{"image_tokens":2}}}}`), &usage) + svc.parseOpenAIImagesSSEUsageBytes([]byte(`{"type":"response.completed","response":{"tool_usage":{"image_gen":{"input_tokens":46,"output_tokens":2459,"output_tokens_details":{"image_tokens":2459}}}}} trailing`), &usage) + + require.Equal(t, OpenAIUsage{InputTokens: 3, OutputTokens: 4, ImageOutputTokens: 2}, usage) +} + +func TestBoundedJSONNonNegativeInt(t *testing.T) { + tests := []struct { + name string + raw string + want int + ok bool + }{ + {name: "scale reduction before accumulation", raw: `10000000000000000000e-19`, want: 1, ok: true}, + {name: "decimal scale reduction", raw: `10000000000000000000.0e-19`, want: 1, ok: true}, + {name: "fractional after scale reduction", raw: `10000000000000000001e-19`, ok: false}, + {name: "overflow after scale reduction", raw: `92233720368547758080e-1`, ok: false}, + {name: "zero with negative exponent", raw: `0e-100`, want: 0, ok: true}, + {name: "zero beyond exponent bound", raw: `0e101`, want: 0, ok: true}, + {name: "zero padded decimal beyond exponent bound", raw: `0.000000e+000000000000000000000000000000000000000000000000101`, want: 0, ok: true}, + {name: "zero padded exponent", raw: `1e0000`, want: 1, ok: true}, + {name: "negative zero syntax", raw: `-0e101`, ok: false}, + {name: "hostile exponent", raw: `1e-1000`, ok: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := boundedJSONNonNegativeInt(gjson.Parse(tt.raw)) + require.Equal(t, tt.ok, ok) + require.Equal(t, tt.want, got) + }) + } +} + func TestOpenAIGatewayServiceForwardImages_OAuthUpstreamHTTPErrorSurfacesRealError(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","response_format":"b64_json"}`) @@ -1244,7 +1347,7 @@ func TestOpenAIGatewayServiceForwardImages_OAuthStreamingTransformsEvents(t *tes Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000001,\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"background\":\"auto\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"}]}}\n\n" + "data: {\"type\":\"response.image_generation_call.partial_image\",\"partial_image_b64\":\"cGFydGlhbA==\",\"partial_image_index\":0,\"output_format\":\"png\",\"background\":\"auto\"}\n\n" + - "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000001,\"usage\":{\"input_tokens\":5,\"output_tokens\":9,\"output_tokens_details\":{\"image_tokens\":4}},\"tool_usage\":{\"image_gen\":{\"images\":1}},\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"background\":\"auto\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"}],\"output\":[{\"type\":\"image_generation_call\",\"result\":\"ZmluYWw=\",\"output_format\":\"png\"}]}}\n\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000001,\"usage\":{\"input_tokens\":5,\"output_tokens\":9,\"output_tokens_details\":{\"image_tokens\":4}},\"tool_usage\":{\"image_gen\":{\"input_tokens\":46,\"output_tokens\":2459,\"output_tokens_details\":{\"image_tokens\":2459},\"images\":1}},\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"background\":\"auto\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"}],\"output\":[{\"type\":\"image_generation_call\",\"result\":\"ZmluYWw=\",\"output_format\":\"png\"}]}}\n\n" + "data: [DONE]\n\n", )), }, @@ -1266,6 +1369,7 @@ func TestOpenAIGatewayServiceForwardImages_OAuthStreamingTransformsEvents(t *tes require.NotNil(t, result) require.True(t, result.Stream) require.Equal(t, 1, result.ImageCount) + require.Equal(t, OpenAIUsage{InputTokens: 46, OutputTokens: 2459, ImageOutputTokens: 2459}, result.Usage) events := parseOpenAIImageTestSSEEvents(rec.Body.String()) partial, ok := findOpenAIImageTestSSEEvent(events, "image_generation.partial_image") require.True(t, ok) @@ -1290,7 +1394,7 @@ func TestOpenAIGatewayServiceForwardImages_OAuthStreamingTransformsEvents(t *tes require.Equal(t, "high", gjson.Get(completed.Data, "quality").String()) require.Equal(t, "1024x1024", gjson.Get(completed.Data, "size").String()) require.Equal(t, "auto", gjson.Get(completed.Data, "background").String()) - require.JSONEq(t, `{"images":1}`, gjson.Get(completed.Data, "usage").Raw) + require.JSONEq(t, `{"input_tokens":46,"output_tokens":2459,"output_tokens_details":{"image_tokens":2459},"images":1}`, gjson.Get(completed.Data, "usage").Raw) require.False(t, gjson.Get(completed.Data, "revised_prompt").Exists()) } diff --git a/backend/internal/service/openai_request_body_limit_failover_test.go b/backend/internal/service/openai_request_body_limit_failover_test.go new file mode 100644 index 0000000000..6d6b5cb63e --- /dev/null +++ b/backend/internal/service/openai_request_body_limit_failover_test.go @@ -0,0 +1,135 @@ +package service + +import ( + "bytes" + "context" + "errors" + "fmt" + "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 TestOpenAIRequestBodyLimitFailover_HTTP413SwitchesAccountsBeforeWrite(t *testing.T) { + gin.SetMode(gin.TestMode) + requestBody := []byte(`{"model":"gpt-5.2","stream":false,"input":"hello"}`) + + for _, passthrough := range []bool{false, true} { + name := "native_responses" + if passthrough { + name = "api_key_passthrough" + } + t.Run(name, func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + + const upstreamBody = `{"error":{"message":"request body exceeds this account's 16MB proxy limit; secret=must-not-leak","type":"invalid_request_error"}}` + body := &passthroughCloseTrackingReadCloser{Reader: strings.NewReader(upstreamBody)} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusRequestEntityTooLarge, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "X-Request-Id": []string{"rid-body-limit"}, + }, + Body: body, + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{ForceCodexCLI: false}}, + httpUpstream: upstream, + } + account := &Account{ + ID: 161, + Name: name, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://api.example.test", + "pool_mode": true, + "pool_mode_retry_status_codes": []any{ + float64(http.StatusRequestEntityTooLarge), + }, + }, + Extra: map[string]any{ + "openai_passthrough": passthrough, + "openai_responses_supported": true, + }, + Status: StatusActive, + Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, requestBody) + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusRequestEntityTooLarge, failoverErr.StatusCode) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, GatewayFailureReason("openai_request_body_too_large"), failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Equal(t, http.StatusRequestEntityTooLarge, failoverErr.ClientStatusCode) + require.Equal(t, "Request payload is too large", failoverErr.ClientMessage) + require.False(t, failoverErr.RetryableOnSameAccount, "a body limit requires another account, not another attempt on the same account") + require.False(t, c.Writer.Written(), "account failover must happen before downstream output is committed") + require.Empty(t, rec.Body.String()) + require.True(t, body.closed) + if passthrough { + require.Equal(t, requestBody, upstream.lastBody) + } else { + require.Equal(t, "gpt-5.2", gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "hello", gjson.GetBytes(upstream.lastBody, "input").String()) + } + }) + } +} + +func TestOpenAIRequestBodyLimitFailover_ContextWindow413DoesNotSwitchAccounts(t *testing.T) { + gin.SetMode(gin.TestMode) + requestBody := []byte(`{"model":"gpt-5.2","stream":false,"input":"hello"}`) + + for _, passthrough := range []bool{false, true} { + t.Run(fmt.Sprintf("passthrough_%t", passthrough), func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + + const upstreamBody = `{"error":{"message":"Your input exceeds the context window of this model. Please adjust your input and try again.","type":"invalid_request_error"}}` + body := &passthroughCloseTrackingReadCloser{Reader: strings.NewReader(upstreamBody)} + svc := &OpenAIGatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{ForceCodexCLI: false}}, + httpUpstream: &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusRequestEntityTooLarge, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: body, + }}, + } + account := &Account{ + ID: 162, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test", "base_url": "https://api.example.test"}, + Extra: map[string]any{ + "openai_passthrough": passthrough, + "openai_responses_supported": true, + }, + Status: StatusActive, Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, requestBody) + + require.Nil(t, result) + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr), "context-window failures are deterministic request errors") + require.True(t, c.Writer.Written()) + require.Contains(t, rec.Body.String(), "exceeds the context window") + require.True(t, body.closed) + }) + } +} diff --git a/backend/internal/service/openai_responses_lite_tools.go b/backend/internal/service/openai_responses_lite_tools.go index a15bc498bc..c5c16ecaa1 100644 --- a/backend/internal/service/openai_responses_lite_tools.go +++ b/backend/internal/service/openai_responses_lite_tools.go @@ -7,18 +7,23 @@ import ( "strings" ) -// normalizeOpenAIResponsesLiteTools converts private namespace declarations -// into the input.additional_tools carrier required by Responses Lite. Other -// top-level tools must belong to the small set accepted by the Lite endpoint; -// rejecting unsupported hosted tools is intentional because silently dropping -// them would change the client's requested behavior. +// normalizeOpenAIResponsesLiteTools applies the Responses Lite request +// contract: reasoning must cover all turns, and private namespace declarations +// use the input.additional_tools carrier. Other top-level tools must belong to +// the small set accepted by the Lite endpoint; rejecting unsupported hosted +// tools is intentional because silently dropping them would change behavior. func normalizeOpenAIResponsesLiteTools(reqBody map[string]any) (bool, error) { if reqBody == nil { return false, nil } + if rawReasoning, exists := reqBody["reasoning"]; exists && rawReasoning != nil { + if _, ok := rawReasoning.(map[string]any); !ok { + return false, fmt.Errorf("responses Lite requires reasoning to be an object") + } + } rawTools, exists := reqBody["tools"] if !exists || rawTools == nil { - return false, nil + return ensureOpenAIResponsesLiteReasoningContext(reqBody) } tools, ok := rawTools.([]any) if !ok { @@ -52,13 +57,16 @@ func normalizeOpenAIResponsesLiteTools(reqBody map[string]any) (bool, error) { } } if len(namespaceTools) == 0 { - return false, nil + return ensureOpenAIResponsesLiteReasoningContext(reqBody) } input, err := appendOpenAIResponsesLiteAdditionalTools(reqBody["input"], namespaceTools) if err != nil { return false, err } + if _, err := ensureOpenAIResponsesLiteReasoningContext(reqBody); err != nil { + return false, err + } reqBody["input"] = input if len(topLevelTools) == 0 { delete(reqBody, "tools") @@ -68,6 +76,23 @@ func normalizeOpenAIResponsesLiteTools(reqBody map[string]any) (bool, error) { return true, nil } +func ensureOpenAIResponsesLiteReasoningContext(reqBody map[string]any) (bool, error) { + rawReasoning, exists := reqBody["reasoning"] + if !exists || rawReasoning == nil { + reqBody["reasoning"] = map[string]any{"context": "all_turns"} + return true, nil + } + reasoning, ok := rawReasoning.(map[string]any) + if !ok { + return false, fmt.Errorf("responses Lite requires reasoning to be an object") + } + if context, ok := reasoning["context"].(string); ok && context == "all_turns" { + return false, nil + } + reasoning["context"] = "all_turns" + return true, nil +} + func appendOpenAIResponsesLiteAdditionalTools(input any, namespaceTools []any) ([]any, error) { var items []any switch typed := input.(type) { diff --git a/backend/internal/service/openai_responses_lite_tools_test.go b/backend/internal/service/openai_responses_lite_tools_test.go index aa03954463..b76e97e272 100644 --- a/backend/internal/service/openai_responses_lite_tools_test.go +++ b/backend/internal/service/openai_responses_lite_tools_test.go @@ -135,6 +135,7 @@ func TestNormalizeOpenAIResponsesLiteTools_ConvertsStringInput(t *testing.T) { func TestNormalizeOpenAIResponsesLiteTools_KeepsSupportedTopLevelTools(t *testing.T) { reqBody := map[string]any{ + "reasoning": map[string]any{"context": "all_turns"}, "tools": []any{ map[string]any{"type": "function", "name": "shell"}, map[string]any{"type": "custom", "name": "exec"}, @@ -150,6 +151,46 @@ func TestNormalizeOpenAIResponsesLiteTools_KeepsSupportedTopLevelTools(t *testin require.Len(t, reqBody["tools"], 4) } +func TestNormalizeOpenAIResponsesLiteTools_EnsuresReasoningContext(t *testing.T) { + tests := []struct { + name string + reasoning any + }{ + {name: "missing"}, + {name: "missing context", reasoning: map[string]any{"effort": "high"}}, + {name: "wrong context", reasoning: map[string]any{"effort": "medium", "context": "current_turn"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + reqBody := map[string]any{"input": "hello"} + if tt.reasoning != nil { + reqBody["reasoning"] = tt.reasoning + } + + changed, err := normalizeOpenAIResponsesLiteTools(reqBody) + + require.NoError(t, err) + require.True(t, changed) + reasoning := reqBody["reasoning"].(map[string]any) + require.Equal(t, "all_turns", reasoning["context"]) + if tt.name != "missing" { + require.Equal(t, tt.reasoning.(map[string]any)["effort"], reasoning["effort"]) + } + }) + } +} + +func TestNormalizeOpenAIResponsesLiteTools_RejectsNonObjectReasoning(t *testing.T) { + reqBody := map[string]any{"reasoning": "high"} + + changed, err := normalizeOpenAIResponsesLiteTools(reqBody) + + require.ErrorContains(t, err, "reasoning to be an object") + require.False(t, changed) + require.Equal(t, "high", reqBody["reasoning"]) +} + func TestNormalizeOpenAIResponsesLiteTools_RejectsUnsupportedTools(t *testing.T) { tests := []struct { name string @@ -240,6 +281,7 @@ func TestOpenAIGatewayServiceForward_NormalizesResponsesLiteToolsForOAuth(t *tes } body := []byte(`{ "model":"gpt-5.6-terra","stream":true,"instructions":"test", + "reasoning":{"effort":"high","context":"current_turn"}, "tools":[ {"type":"function","name":"shell","parameters":{"type":"object"}}, {"type":"custom","name":"exec"}, @@ -255,6 +297,8 @@ func TestOpenAIGatewayServiceForward_NormalizesResponsesLiteToolsForOAuth(t *tes require.NoError(t, err) require.NotNil(t, result) require.Equal(t, "true", upstream.lastReq.Header.Get(responsesLiteHeader)) + require.Equal(t, "high", gjson.GetBytes(upstream.lastBody, "reasoning.effort").String()) + require.Equal(t, "all_turns", gjson.GetBytes(upstream.lastBody, "reasoning.context").String()) require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="namespace")`).Exists()) require.Equal(t, "shell", gjson.GetBytes(upstream.lastBody, `tools.#(type=="function").name`).String()) require.Equal(t, "exec", gjson.GetBytes(upstream.lastBody, `tools.#(type=="custom").name`).String()) diff --git a/backend/internal/service/openai_responses_rejected_field_retry.go b/backend/internal/service/openai_responses_rejected_field_retry.go new file mode 100644 index 0000000000..d8d2bb02bf --- /dev/null +++ b/backend/internal/service/openai_responses_rejected_field_retry.go @@ -0,0 +1,133 @@ +package service + +import ( + "crypto/sha256" + "fmt" + "net/http" + "regexp" + "strconv" + "strings" + + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +const maxOpenAIResponsesRejectedFieldRetries = 6 + +var ( + openAIResponsesRejectedNamespaceParamPattern = regexp.MustCompile(`(?i)^input\[(\d+)\]\.namespace$`) + openAIResponsesRejectedMessageParamPattern = regexp.MustCompile(`(?i)(?:unknown|unsupported)[ _-]+parameter\s*(?::|=|is)?\s*["']?(max_output_tokens|input\[\d+\]\.namespace)(?:["']|\b)`) +) + +type openAIResponsesRejectedFieldRetryState struct { + attempts int + seenBodyHashes map[[sha256.Size]byte]struct{} +} + +func newOpenAIResponsesRejectedFieldRetryState(initialBody []byte) *openAIResponsesRejectedFieldRetryState { + state := &openAIResponsesRejectedFieldRetryState{ + seenBodyHashes: make(map[[sha256.Size]byte]struct{}, maxOpenAIResponsesRejectedFieldRetries+1), + } + state.remember(initialBody) + return state +} + +func (s *openAIResponsesRejectedFieldRetryState) Allow(nextBody []byte) bool { + if s == nil || len(nextBody) == 0 || s.attempts >= maxOpenAIResponsesRejectedFieldRetries { + return false + } + bodyHash := sha256.Sum256(nextBody) + if _, seen := s.seenBodyHashes[bodyHash]; seen { + return false + } + s.seenBodyHashes[bodyHash] = struct{}{} + s.attempts++ + return true +} + +func (s *openAIResponsesRejectedFieldRetryState) remember(body []byte) { + if s == nil || len(body) == 0 { + return + } + if s.seenBodyHashes == nil { + s.seenBodyHashes = make(map[[sha256.Size]byte]struct{}, maxOpenAIResponsesRejectedFieldRetries+1) + } + s.seenBodyHashes[sha256.Sum256(body)] = struct{}{} +} + +func normalizeOpenAIResponsesRejectedFieldRetryBody(statusCode int, body, responseBody []byte) ([]byte, string, bool, error) { + if statusCode != http.StatusBadRequest || len(body) == 0 || len(responseBody) == 0 { + return nil, "", false, nil + } + + code := strings.ToLower(strings.TrimSpace(extractUpstreamErrorCode(responseBody))) + message := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(responseBody))) + if !isExplicitOpenAIResponsesFieldRejection(code, message) { + return nil, "", false, nil + } + + param := strings.ToLower(strings.TrimSpace(gjson.GetBytes(responseBody, "error.param").String())) + if param == "" { + param = openAIResponsesRejectedParamFromMessage(message) + } + if index, ok := openAIResponsesRejectedNamespaceIndex(param); ok { + return removeOpenAIResponsesRejectedNamespaceAtIndex(body, index) + } + if param == "max_output_tokens" && gjson.GetBytes(body, "max_output_tokens").Exists() { + retryBody, err := sjson.DeleteBytes(body, "max_output_tokens") + if err != nil { + return nil, "", false, fmt.Errorf("delete rejected max_output_tokens: %w", err) + } + return retryBody, "max_output_tokens parameter rejection", true, nil + } + return nil, "", false, nil +} + +func isExplicitOpenAIResponsesFieldRejection(code, message string) bool { + switch strings.TrimSpace(code) { + case "unknown_parameter", "unsupported_parameter": + return true + } + return strings.Contains(message, "unknown parameter") || + strings.Contains(message, "unsupported parameter") +} + +func openAIResponsesRejectedParamFromMessage(message string) string { + match := openAIResponsesRejectedMessageParamPattern.FindStringSubmatch(strings.TrimSpace(message)) + if len(match) != 2 { + return "" + } + return strings.ToLower(strings.TrimSpace(match[1])) +} + +func openAIResponsesRejectedNamespaceIndex(param string) (int, bool) { + match := openAIResponsesRejectedNamespaceParamPattern.FindStringSubmatch(strings.TrimSpace(param)) + if len(match) != 2 { + return 0, false + } + index, err := strconv.Atoi(match[1]) + if err == nil && index >= 0 { + return index, true + } + return 0, false +} + +func removeOpenAIResponsesRejectedNamespaceAtIndex(body []byte, index int) ([]byte, string, bool, error) { + itemPath := fmt.Sprintf("input.%d", index) + itemType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, itemPath+".type").String())) + switch itemType { + case "function_call", "tool_call", "custom_tool_call", "mcp_tool_call": + default: + return nil, "", false, nil + } + + namespacePath := itemPath + ".namespace" + if !gjson.GetBytes(body, namespacePath).Exists() { + return nil, "", false, nil + } + retryBody, err := sjson.DeleteBytes(body, namespacePath) + if err != nil { + return nil, "", false, fmt.Errorf("delete rejected namespace at input[%d]: %w", index, err) + } + return retryBody, "indexed namespace parameter rejection", true, nil +} diff --git a/backend/internal/service/openai_responses_rejected_field_retry_test.go b/backend/internal/service/openai_responses_rejected_field_retry_test.go new file mode 100644 index 0000000000..96c20de508 --- /dev/null +++ b/backend/internal/service/openai_responses_rejected_field_retry_test.go @@ -0,0 +1,230 @@ +package service + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestOpenAIResponsesRejectedFieldRetryStateRejectsDuplicateBodyAndCap(t *testing.T) { + initialBody := []byte(`{"model":"gpt-5.5"}`) + state := newOpenAIResponsesRejectedFieldRetryState(initialBody) + + require.False(t, state.Allow(initialBody)) + for attempt := 0; attempt < maxOpenAIResponsesRejectedFieldRetries; attempt++ { + nextBody := []byte(fmt.Sprintf(`{"model":"gpt-5.5","variant":%d}`, attempt)) + require.True(t, state.Allow(nextBody)) + require.False(t, state.Allow(nextBody)) + } + require.False(t, state.Allow([]byte(`{"model":"gpt-5.5","variant":"overflow"}`))) +} + +func TestNormalizeOpenAIResponsesRejectedFieldRetryBodyRejectsAmbiguousErrors(t *testing.T) { + tests := []struct { + name string + body []byte + responseBody []byte + }{ + { + name: "namespace belongs to message", + body: []byte(`{"input":[{"type":"message","namespace":"keep"}]}`), + responseBody: []byte(`{"error":{"code":"unknown_parameter","message":"Unknown parameter: 'input[0].namespace'.","param":"input[0].namespace"}}`), + }, + { + name: "max output tokens only mentioned", + body: []byte(`{"max_output_tokens":4096}`), + responseBody: []byte(`{"error":{"code":"invalid_request_error","message":"max_output_tokens must be positive","param":"max_output_tokens"}}`), + }, + { + name: "structured param overrides namespace mention", + body: []byte(`{"input":[{"type":"function_call","namespace":"keep","arguments":"{}"}]}`), + responseBody: []byte(`{"error":{"code":"unknown_parameter","message":"Unknown parameter: 'input[0].namespace'.","param":"tools"}}`), + }, + { + name: "nested max output tokens param is not top level", + body: []byte(`{"max_output_tokens":4096,"input":[{"type":"message","content":{"max_output_tokens":"keep"}}]}`), + responseBody: []byte(`{"error":{"code":"unknown_parameter","message":"Unknown parameter: input[0].content.max_output_tokens","param":"input[0].content.max_output_tokens"}}`), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + retryBody, _, changed, err := normalizeOpenAIResponsesRejectedFieldRetryBody(http.StatusBadRequest, tt.body, tt.responseBody) + require.NoError(t, err) + require.False(t, changed) + require.Nil(t, retryBody) + }) + } +} + +func TestNormalizeOpenAIResponsesRejectedFieldRetryBodyFindsNamespacePathInMessage(t *testing.T) { + body := []byte(`{"input":[{"type":"function_call","namespace":"keep","arguments":"{}"},{"type":"function_call","namespace":"remove","arguments":"{}"}]}`) + responseBody := []byte(`{"error":{"code":"unknown_parameter","message":"input[0] was accepted; Unknown parameter: 'input[1].namespace'."}}`) + + retryBody, _, changed, err := normalizeOpenAIResponsesRejectedFieldRetryBody(http.StatusBadRequest, body, responseBody) + + require.NoError(t, err) + require.True(t, changed) + require.Equal(t, "keep", gjson.GetBytes(retryBody, "input.0.namespace").String()) + require.False(t, gjson.GetBytes(retryBody, "input.1.namespace").Exists()) +} + +func TestNormalizeOpenAIResponsesRejectedFieldRetryBodyBindsNamespacePathToRejectionPhrase(t *testing.T) { + body := []byte(`{"input":[{"type":"function_call","namespace":"keep","arguments":"{}"},{"type":"function_call","namespace":"remove","arguments":"{}"}]}`) + responseBody := []byte(`{"error":{"code":"unknown_parameter","message":"input[0].namespace is supported; Unknown parameter: input[1].namespace."}}`) + + retryBody, _, changed, err := normalizeOpenAIResponsesRejectedFieldRetryBody(http.StatusBadRequest, body, responseBody) + + require.NoError(t, err) + require.True(t, changed) + require.Equal(t, "keep", gjson.GetBytes(retryBody, "input.0.namespace").String()) + require.False(t, gjson.GetBytes(retryBody, "input.1.namespace").Exists()) +} + +func TestNormalizeOpenAIResponsesRejectedFieldRetryBodyDoesNotTreatMaxOutputTokensSuggestionAsRejection(t *testing.T) { + body := []byte(`{"max_tokens":4096,"max_output_tokens":2048}`) + responseBody := []byte(`{"error":{"code":"unknown_parameter","message":"Unknown parameter: max_tokens. Use max_output_tokens instead."}}`) + + retryBody, _, changed, err := normalizeOpenAIResponsesRejectedFieldRetryBody(http.StatusBadRequest, body, responseBody) + + require.NoError(t, err) + require.False(t, changed) + require.Nil(t, retryBody) +} + +func TestNormalizeOpenAIResponsesRejectedFieldRetryBodyBindsMaxOutputTokensToRejectionPhrase(t *testing.T) { + body := []byte(`{"max_output_tokens":2048}`) + responseBody := []byte(`{"error":{"code":"unsupported_parameter","message":"Unsupported parameter: max_output_tokens."}}`) + + retryBody, _, changed, err := normalizeOpenAIResponsesRejectedFieldRetryBody(http.StatusBadRequest, body, responseBody) + + require.NoError(t, err) + require.True(t, changed) + require.False(t, gjson.GetBytes(retryBody, "max_output_tokens").Exists()) +} + +func TestOpenAIGatewayService_RetriesRejectedIndexedNamespaceField(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"function_call","name":"first","namespace":"keep","arguments":"{}"},{"type":"custom_tool_call","name":"second","namespace":"remove","input":"{}"}]}`) + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + newOpenAIRejectedFieldTestResponse(http.StatusBadRequest, `{"error":{"code":"unknown_parameter","message":"Unknown parameter: 'input[1].namespace'.","param":"input[1].namespace","type":"invalid_request_error"}}`), + newOpenAIRejectedFieldTestResponse(http.StatusOK, `{"output":[],"usage":{"input_tokens":1,"output_tokens":1,"input_tokens_details":{"cached_tokens":0}}}`), + }} + + result, err := newOpenAIRejectedFieldTestService(upstream).Forward( + context.Background(), + newOpenAIRejectedFieldTestContext(body), + newOpenAIRejectedFieldTestAccount(), + body, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.bodies, 2) + require.Equal(t, "keep", gjson.GetBytes(upstream.bodies[1], "input.0.namespace").String()) + require.False(t, gjson.GetBytes(upstream.bodies[1], "input.1.namespace").Exists()) +} + +func TestOpenAIGatewayService_RetriesExplicitMaxOutputTokensRejection(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","stream":false,"max_output_tokens":4096,"input":[{"type":"message","role":"user","content":{"max_output_tokens":"keep"}}]}`) + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + newOpenAIRejectedFieldTestResponse(http.StatusBadRequest, `{"error":{"code":"unsupported_parameter","message":"Unsupported parameter: max_output_tokens","param":"max_output_tokens","type":"invalid_request_error"}}`), + newOpenAIRejectedFieldTestResponse(http.StatusOK, `{"output":[],"usage":{"input_tokens":1,"output_tokens":1,"input_tokens_details":{"cached_tokens":0}}}`), + }} + + result, err := newOpenAIRejectedFieldTestService(upstream).Forward( + context.Background(), + newOpenAIRejectedFieldTestContext(body), + newOpenAIRejectedFieldTestAccount(), + body, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.bodies, 2) + require.Equal(t, int64(4096), gjson.GetBytes(upstream.bodies[0], "max_output_tokens").Int()) + require.False(t, gjson.GetBytes(upstream.bodies[1], "max_output_tokens").Exists()) + require.Equal(t, "keep", gjson.GetBytes(upstream.bodies[1], "input.0.content.max_output_tokens").String()) +} + +func TestOpenAIGatewayService_ComposesDistinctRejectedFieldRetries(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","stream":false,"max_output_tokens":2048,"input":[{"type":"function_call","name":"first","namespace":"keep","arguments":"{}"},{"type":"custom_tool_call","name":"second","namespace":"remove","input":"{}"}]}`) + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + newOpenAIRejectedFieldTestResponse(http.StatusBadRequest, `{"error":{"code":"unknown_parameter","message":"Unknown parameter: 'input[1].namespace'.","param":"input[1].namespace"}}`), + newOpenAIRejectedFieldTestResponse(http.StatusBadRequest, `{"error":{"code":"unsupported_parameter","message":"Unsupported parameter: max_output_tokens","param":"max_output_tokens"}}`), + newOpenAIRejectedFieldTestResponse(http.StatusOK, `{"output":[],"usage":{"input_tokens":1,"output_tokens":1,"input_tokens_details":{"cached_tokens":0}}}`), + }} + + result, err := newOpenAIRejectedFieldTestService(upstream).Forward( + context.Background(), + newOpenAIRejectedFieldTestContext(body), + newOpenAIRejectedFieldTestAccount(), + body, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.bodies, 3) + require.True(t, gjson.GetBytes(upstream.bodies[0], "input.1.namespace").Exists()) + require.False(t, gjson.GetBytes(upstream.bodies[1], "input.1.namespace").Exists()) + require.Equal(t, int64(2048), gjson.GetBytes(upstream.bodies[1], "max_output_tokens").Int()) + require.False(t, gjson.GetBytes(upstream.bodies[2], "input.1.namespace").Exists()) + require.False(t, gjson.GetBytes(upstream.bodies[2], "max_output_tokens").Exists()) +} + +func newOpenAIRejectedFieldTestService(upstream *httpUpstreamRecorder) *OpenAIGatewayService { + return &OpenAIGatewayService{ + cfg: &config.Config{Security: config.SecurityConfig{ + URLAllowlist: config.URLAllowlistConfig{Enabled: false}, + }}, + httpUpstream: upstream, + } +} + +func newOpenAIRejectedFieldTestContext(body []byte) *gin.Context { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set("User-Agent", "curl/8.0") + return c +} + +func newOpenAIRejectedFieldTestAccount() *Account { + return &Account{ + ID: 5107, + Name: "responses-compatible", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://compat.example", + }, + Extra: map[string]any{ + openai_compat.ExtraKeyResponsesMode: string(openai_compat.ResponsesSupportModeAuto), + openai_compat.ExtraKeyResponsesSupported: true, + }, + Status: StatusActive, + Schedulable: true, + } +} + +func newOpenAIRejectedFieldTestResponse(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} diff --git a/backend/internal/service/openai_sse_concatenated_json_test.go b/backend/internal/service/openai_sse_concatenated_json_test.go index 9b06ef9a39..27c057157e 100644 --- a/backend/internal/service/openai_sse_concatenated_json_test.go +++ b/backend/internal/service/openai_sse_concatenated_json_test.go @@ -90,6 +90,138 @@ func TestOpenAIWSv2StreamingRepairsConcatenatedJSONDocumentsInSingleMessage(t *t }) } +func TestOpenAIWSv2RejectsMalformedTypedEventBeforeWritingDownstream(t *testing.T) { + largeInProgress, _, _ := openAIConcatenatedJSONTestEvents(t) + testOpenAIWSv2RejectsMalformedEventBeforeWritingDownstream(t, []byte(largeInProgress+"unexpected-tail")) +} + +func TestOpenAIWSv2RejectsMalformedUntypedMessageBeforeWritingDownstream(t *testing.T) { + testOpenAIWSv2RejectsMalformedEventBeforeWritingDownstream(t, []byte("not-json")) +} + +func TestOpenAIWSv2RejectsMalformedEventAfterWritingDownstream(t *testing.T) { + gin.SetMode(gin.TestMode) + + outputTextDelta := `{"type":"response.output_text.delta","delta":"ok","sequence_number":1}` + malformedMessage := `{"type":"response.in_progress"}unexpected-tail` + captureConn := &openAIWSCaptureConn{events: [][]byte{ + []byte(outputTextDelta), + []byte(malformedMessage), + }} + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn}) + svc := &OpenAIGatewayService{ + cfg: cfg, + cache: &stubGatewayCache{}, + httpUpstream: &httpUpstreamRecorder{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + openaiWSPool: pool, + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 5, + Name: "ws-malformed-event-after-output", + 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}, + } + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + groupID := int64(1) + c.Set("api_key", &APIKey{GroupID: &groupID}) + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.6-sol","stream":true,"input":"hello"}`)) + require.Error(t, err) + require.Contains(t, err.Error(), "after downstream output") + require.Nil(t, result) + require.True(t, captureConn.closed) + require.Contains(t, recorder.Body.String(), `"delta":"ok"`) + require.NotContains(t, recorder.Body.String(), "unexpected-tail") + require.NotContains(t, recorder.Body.String(), "response.in_progress") + assertOpenAISSEFrames(t, recorder.Body.String(), []string{"response.output_text.delta"}) +} + +func testOpenAIWSv2RejectsMalformedEventBeforeWritingDownstream(t *testing.T, malformedMessage []byte) { + t.Helper() + gin.SetMode(gin.TestMode) + + _, _, completed := openAIConcatenatedJSONTestEvents(t) + outputTextDelta := `{"type":"response.output_text.delta","delta":"ok","sequence_number":3}` + captureConn := &openAIWSCaptureConn{events: [][]byte{ + malformedMessage, + []byte(outputTextDelta), + []byte(completed), + }} + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn}) + svc := &OpenAIGatewayService{ + cfg: cfg, + cache: &stubGatewayCache{}, + httpUpstream: &httpUpstreamRecorder{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + openaiWSPool: pool, + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 4, + Name: "ws-malformed-event", + 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}, + } + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + groupID := int64(1) + c.Set("api_key", &APIKey{GroupID: &groupID}) + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.6-sol","stream":true,"input":"hello"}`)) + require.Error(t, err) + var fallbackErr *openAIWSFallbackError + require.ErrorAs(t, err, &fallbackErr) + require.Equal(t, "invalid_event_json", fallbackErr.Reason) + require.Nil(t, result) + require.Empty(t, recorder.Body.String()) + require.True(t, captureConn.closed) +} + func TestSplitOpenAIConcatenatedJSONDocumentsRejectsPayloadOverRepairLimit(t *testing.T) { first := `{"type":"response.in_progress","padding":"` + strings.Repeat("x", 16*1024*1024) + `"}` second := `{"type":"response.completed"}` diff --git a/backend/internal/service/openai_ws_client_read.go b/backend/internal/service/openai_ws_client_read.go new file mode 100644 index 0000000000..da8d650be4 --- /dev/null +++ b/backend/internal/service/openai_ws_client_read.go @@ -0,0 +1,70 @@ +package service + +import ( + "context" + "errors" + "time" + + coderws "github.com/coder/websocket" +) + +type openAIWSClientReadResult struct { + messageType coderws.MessageType + payload []byte + err error +} + +// ReadOpenAIWSClientMessage keeps one reader alive while control events send +// their close frame, then closes the transport and joins that reader. +func ReadOpenAIWSClientMessage( + controlCtx context.Context, + conn *coderws.Conn, + timeout time.Duration, + timeoutStatus coderws.StatusCode, + timeoutReason string, +) (coderws.MessageType, []byte, error) { + if conn == nil { + return 0, nil, errors.New("openai websocket client connection is nil") + } + if controlCtx == nil { + controlCtx = context.Background() + } + + readDone := make(chan openAIWSClientReadResult, 1) + go func() { + messageType, payload, err := conn.Read(context.Background()) + readDone <- openAIWSClientReadResult{messageType: messageType, payload: payload, err: err} + }() + + var timeoutCh <-chan time.Time + var timer *time.Timer + if timeout > 0 { + timer = time.NewTimer(timeout) + timeoutCh = timer.C + defer timer.Stop() + } + + closeAndJoin := func(status coderws.StatusCode, reason string, cause error) (coderws.MessageType, []byte, error) { + _ = conn.Close(status, reason) + _ = conn.CloseNow() + <-readDone + return 0, nil, NewOpenAIWSClientCloseError(status, reason, cause) + } + + select { + case result := <-readDone: + return result.messageType, result.payload, result.err + case <-timeoutCh: + return closeAndJoin(timeoutStatus, timeoutReason, context.DeadlineExceeded) + case <-controlCtx.Done(): + cause := context.Cause(controlCtx) + if errors.Is(cause, ErrOpenAIWSIngressLeaseLost) { + return closeAndJoin( + coderws.StatusTryAgainLater, + "websocket ingress capacity lease lost; please reconnect", + cause, + ) + } + return closeAndJoin(coderws.StatusGoingAway, "websocket request canceled", cause) + } +} diff --git a/backend/internal/service/openai_ws_client_read_test.go b/backend/internal/service/openai_ws_client_read_test.go new file mode 100644 index 0000000000..82df19849b --- /dev/null +++ b/backend/internal/service/openai_ws_client_read_test.go @@ -0,0 +1,143 @@ +package service + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + coderws "github.com/coder/websocket" + "github.com/stretchr/testify/require" +) + +func TestReadOpenAIWSClientMessage_ControlCloseFrames(t *testing.T) { + tests := []struct { + name string + timeout time.Duration + timeoutStatus coderws.StatusCode + timeoutReason string + cancelCause error + wantStatus coderws.StatusCode + wantReason string + }{ + { + name: "inter-turn idle sends normal close", + timeout: 25 * time.Millisecond, + timeoutStatus: coderws.StatusNormalClosure, + timeoutReason: "websocket idle timeout", + wantStatus: coderws.StatusNormalClosure, + wantReason: "websocket idle timeout", + }, + { + name: "first message timeout sends policy close", + timeout: 25 * time.Millisecond, + timeoutStatus: coderws.StatusPolicyViolation, + timeoutReason: "missing first response.create message", + wantStatus: coderws.StatusPolicyViolation, + wantReason: "missing first response.create message", + }, + { + name: "lease loss sends retry close", + cancelCause: ErrOpenAIWSIngressLeaseLost, + wantStatus: coderws.StatusTryAgainLater, + wantReason: "websocket ingress capacity lease lost; please reconnect", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + serverResult := make(chan error, 1) + readStarted := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, nil) + if err != nil { + serverResult <- err + return + } + defer func() { _ = conn.CloseNow() }() + close(readStarted) + _, _, err = ReadOpenAIWSClientMessage( + controlCtx, + conn, + tt.timeout, + tt.timeoutStatus, + tt.timeoutReason, + ) + serverResult <- err + })) + defer server.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(server.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + <-readStarted + if tt.cancelCause != nil { + cancelControl(tt.cancelCause) + } + + readCtx, cancelRead := context.WithTimeout(context.Background(), time.Second) + _, _, err = clientConn.Read(readCtx) + cancelRead() + var clientClose coderws.CloseError + require.ErrorAs(t, err, &clientClose) + require.Equal(t, tt.wantStatus, clientClose.Code) + require.Equal(t, tt.wantReason, clientClose.Reason) + + select { + case serverErr := <-serverResult: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, serverErr, &closeErr) + require.Equal(t, tt.wantStatus, closeErr.StatusCode()) + require.Equal(t, tt.wantReason, closeErr.Reason()) + case <-time.After(time.Second): + t.Fatal("server read goroutine did not exit after close handshake") + } + }) + } +} + +func TestReadOpenAIWSClientMessage_ParentCancellationStillJoinsRead(t *testing.T) { + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + serverResult := make(chan error, 1) + readStarted := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, nil) + if err != nil { + serverResult <- err + return + } + defer func() { _ = conn.CloseNow() }() + close(readStarted) + _, _, err = ReadOpenAIWSClientMessage(controlCtx, conn, 0, 0, "") + serverResult <- err + })) + defer server.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(server.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + <-readStarted + cancelControl(errors.New("server shutting down")) + readCtx, cancelRead := context.WithTimeout(context.Background(), time.Second) + _, _, err = clientConn.Read(readCtx) + cancelRead() + var clientClose coderws.CloseError + require.ErrorAs(t, err, &clientClose) + require.Equal(t, coderws.StatusGoingAway, clientClose.Code) + require.Equal(t, "websocket request canceled", clientClose.Reason) + + select { + case <-serverResult: + case <-time.After(time.Second): + t.Fatal("server read goroutine leaked after parent cancellation") + } +} diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 51aace6419..f6a7dca250 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -305,7 +305,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( normalized = stripped logOpenAIWSModeInfo("ingress_ws_codex_spark_image_tool_stripped account_id=%d", account.ID) } - imageIntent := IsImageGenerationIntent(openAIResponsesEndpoint, originalModel, normalized) + imageIntent := IsImageGenerationIntentForPlatform(openAIResponsesEndpoint, originalModel, normalized, account.Platform) if imageIntent && !imageGenerationAllowed { return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, ImageGenerationPermissionMessage(), nil) } @@ -382,22 +382,18 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } readClientMessage := func() ([]byte, error) { - readCtx := ctx idleTimeout := s.openAIWSIngressInterTurnIdleTimeout() - cancelRead := func() {} - if idleTimeout > 0 { - readCtx, cancelRead = context.WithTimeout(ctx, idleTimeout) - } - msgType, payload, readErr := clientConn.Read(readCtx) - cancelRead() + msgType, payload, readErr := ReadOpenAIWSClientMessage( + ctx, + clientConn, + idleTimeout, + coderws.StatusNormalClosure, + "websocket idle timeout", + ) if readErr != nil { - if idleTimeout > 0 && errors.Is(readErr, context.DeadlineExceeded) && ctx.Err() == nil { + var closeErr *OpenAIWSClientCloseError + if errors.As(readErr, &closeErr) && closeErr.StatusCode() == coderws.StatusNormalClosure { logOpenAIWSModeInfo("ingress_ws_inter_turn_idle_timeout account_id=%d timeout_seconds=%d", account.ID, int(idleTimeout.Seconds())) - return nil, NewOpenAIWSClientCloseError( - coderws.StatusNormalClosure, - "websocket idle timeout", - readErr, - ) } return nil, readErr } @@ -674,6 +670,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( return acquireTurnLease(turn, preferred, forcePreferredConn) } if acquireErr != nil { + canonicalModel := canonicalOpenAIAccountSchedulingModel(account, ingressSessionOriginalModel) + s.handleOpenAIWSDialTransientFailure(ctx, account, canonicalModel, acquireErr) dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(acquireErr) logOpenAIWSModeInfo( "ingress_ws_upstream_acquire_fail account_id=%d turn=%d reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_preferred_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v", @@ -818,6 +816,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( lastEventType = eventType } if eventType == "error" { + canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel) + s.handleOpenAIWSErrorEventTransientFailure(ctx, account, canonicalModel, lease.HandshakeHeaders(), upstreamMessage) errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw) fallbackReason, _ := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) @@ -950,6 +950,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } } if isTerminalEvent { + canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel) + terminalEvent := s.handleOpenAIWSTerminalTransientFailure(ctx, account, canonicalModel, lease.HandshakeHeaders(), upstreamMessage) // 客户端已断连时，上游连接的 session 状态不可信，标记 broken 避免回池复用。 if clientDisconnected { lease.MarkBroken() @@ -977,17 +979,18 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } imageCount := imageCounter.Count() result := &OpenAIForwardResult{ - RequestID: responseID, - Usage: usage, - Model: originalModel, - UpstreamModel: mappedModel, - ServiceTier: extractOpenAIServiceTierFromBody(payload), - ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, mappedModel, originalModel), payload, mappedModel), - Stream: reqStream, - OpenAIWSMode: true, - ResponseHeaders: lease.HandshakeHeaders(), - Duration: time.Since(turnStart), - FirstTokenMs: firstTokenMs, + RequestID: responseID, + Usage: usage, + Model: originalModel, + UpstreamModel: mappedModel, + ServiceTier: extractOpenAIServiceTierFromBody(payload), + ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, mappedModel, originalModel), payload, mappedModel), + Stream: reqStream, + OpenAIWSMode: true, + UpstreamTerminalEvent: terminalEvent, + ResponseHeaders: lease.HandshakeHeaders(), + Duration: time.Since(turnStart), + FirstTokenMs: firstTokenMs, } if replayInput := replayCollector.Items(); len(replayInput) > 0 { result.wsReplayInput = replayInput 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 3752c1bd0e..8fbf168f66 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -72,11 +72,11 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT } serverErrCh := make(chan error, 1) - turnWSModeCh := make(chan bool, 2) + turnTerminalCh := make(chan string, 2) hooks := &OpenAIWSIngressHooks{ AfterTurn: func(_ int, result *OpenAIForwardResult, turnErr error) { if turnErr == nil && result != nil { - turnWSModeCh <- result.OpenAIWSMode + turnTerminalCh <- result.UpstreamTerminalEvent } }, } @@ -146,8 +146,8 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT secondTurnEvent := readMessage() require.Equal(t, "response.completed", gjson.GetBytes(secondTurnEvent, "type").String()) require.Equal(t, "resp_ingress_turn_2", gjson.GetBytes(secondTurnEvent, "response.id").String()) - require.True(t, <-turnWSModeCh, "首轮 turn 应标记为 WS 模式") - require.True(t, <-turnWSModeCh, "第二轮 turn 应标记为 WS 模式") + require.Equal(t, "response.completed", <-turnTerminalCh, "首轮 turn 应保留成功终态") + require.Equal(t, "response.completed", <-turnTerminalCh, "第二轮 turn 应保留成功终态") _ = clientConn.Close(coderws.StatusNormalClosure, "done") @@ -249,6 +249,14 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_IdleTimeoutRelea require.NoError(t, err) require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) + closeReadCtx, cancelCloseRead := context.WithTimeout(context.Background(), 3*time.Second) + _, _, err = clientConn.Read(closeReadCtx) + cancelCloseRead() + var clientClose coderws.CloseError + require.ErrorAs(t, err, &clientClose) + require.Equal(t, coderws.StatusNormalClosure, clientClose.Code) + require.Equal(t, "websocket idle timeout", clientClose.Reason) + select { case proxyErr := <-serverErrCh: var closeErr *OpenAIWSClientCloseError @@ -531,6 +539,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_CodexImageBridge "model":"gpt-5.5", "stream":false, "previous_response_id":"resp_codex_image_bridge", + "reasoning":{"effort":"high"}, "client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"}, "tools":[{"type":"namespace","name":"collaboration","tools":[{"type":"function","name":"spawn_agent"}]}], "input":[ @@ -583,6 +592,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_CodexImageBridge require.Equal(t, "png", gjson.Get(nonLitePayload, `tools.#(type=="image_generation").output_format`).String()) require.Equal(t, "auto", gjson.Get(nonLitePayload, "tool_choice").String()) require.Contains(t, gjson.Get(nonLitePayload, "instructions").String(), "image_generation") + require.False(t, gjson.Get(nonLitePayload, "reasoning.context").Exists()) litePayload := requestToJSONString(captureConn.writes[1]) require.False(t, gjson.Get(litePayload, `tools.#(type=="image_generation")`).Exists()) @@ -593,6 +603,8 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_CodexImageBridge require.Equal(t, "collaboration", gjson.Get(litePayload, `input.#(type=="additional_tools").tools.1.name`).String()) require.Equal(t, "namespace", gjson.Get(litePayload, "tool_choice.type").String()) require.Equal(t, "collaboration", gjson.Get(litePayload, "tool_choice.name").String()) + require.Equal(t, "high", gjson.Get(litePayload, "reasoning.effort").String()) + require.Equal(t, "all_turns", gjson.Get(litePayload, "reasoning.context").String()) functionPayload := requestToJSONString(captureConn.writes[2]) require.True(t, gjson.Get(functionPayload, `tools.#(name=="image_gen.imagegen")`).Exists()) @@ -976,6 +988,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughHeade "model":"gpt-5.1", "stream":false, "prompt_cache_key":"pcache_passthrough", + "reasoning":{"effort":"medium","context":"current_turn"}, "client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"}, "tools":[{"type":"namespace","name":"collaboration","tools":[{"type":"function","name":"spawn_agent"}]}], "input":[{"type":"message","role":"user","content":"hello"}], @@ -1009,6 +1022,8 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughHeade require.Equal(t, "collaboration", gjson.Get(forwarded, `input.#(type=="additional_tools").tools.0.name`).String()) require.Equal(t, "namespace", gjson.Get(forwarded, "tool_choice.type").String()) require.Equal(t, "collaboration", gjson.Get(forwarded, "tool_choice.name").String()) + require.Equal(t, "medium", gjson.Get(forwarded, "reasoning.effort").String()) + require.Equal(t, "all_turns", gjson.Get(forwarded, "reasoning.context").String()) } func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_HTTPBridgeModeRelaysHTTPStream(t *testing.T) { diff --git a/backend/internal/service/openai_ws_forwarder_success_test.go b/backend/internal/service/openai_ws_forwarder_success_test.go index bb4ac2242c..37c68a4126 100644 --- a/backend/internal/service/openai_ws_forwarder_success_test.go +++ b/backend/internal/service/openai_ws_forwarder_success_test.go @@ -464,10 +464,61 @@ func TestOpenAIGatewayService_Forward_WSv2_RewriteModelAndToolCallsOnCompletedEv require.NoError(t, err) require.NotNil(t, result) require.Equal(t, "resp_model_tool_1", result.RequestID) + require.Equal(t, "response.completed", result.UpstreamTerminalEvent) + require.True(t, result.SucceededForScheduling()) require.Equal(t, "custom-original-model", gjson.GetBytes(rec.Body.Bytes(), "model").String(), "响应模型应回写为原始请求模型") require.Equal(t, "edit", gjson.GetBytes(rec.Body.Bytes(), "tool_calls.0.function.name").String(), "工具名称应被修正为 OpenCode 规范") } +func TestOpenAIGatewayService_Forward_WSv2_ResponseFailedIsNotSchedulingSuccess(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0") + + cfg := newOpenAIWSV2TestConfig() + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + + captureConn := &openAIWSCaptureConn{events: [][]byte{ + []byte(`{"type":"response.failed","response":{"id":"resp_failed_1","model":"gpt-5.5","error":{"code":"server_error","message":"Internal error"}}}`), + }} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn}) + + svc := &OpenAIGatewayService{ + cfg: cfg, + rateLimitService: NewRateLimitService(transientCooldownAccountRepo{}, nil, cfg, nil, nil), + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + account := &Account{ + ID: 1302, + 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}, + } + svc.recordOpenAIAccountModelTransientFailure(account, "gpt-5.5", time.Now()) + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.5","stream":false,"input":"hello"}`)) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "response.failed", result.UpstreamTerminalEvent) + require.False(t, result.SucceededForScheduling()) + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5")) +} + func TestOpenAIWSPayloadString_OnlyAcceptsStringValues(t *testing.T) { payload := map[string]any{ "type": nil, diff --git a/backend/internal/service/openai_ws_forwarder_support.go b/backend/internal/service/openai_ws_forwarder_support.go index 31c63e214b..91b9ef5cb1 100644 --- a/backend/internal/service/openai_ws_forwarder_support.go +++ b/backend/internal/service/openai_ws_forwarder_support.go @@ -204,6 +204,98 @@ func isOpenAIWSTerminalEvent(eventType string) bool { } } +func normalizeOpenAIWSTerminalEvent(eventType string) string { + switch strings.TrimSpace(eventType) { + case "response.completed": + return "response.completed" + case "response.done": + return "response.done" + case "response.failed": + return "response.failed" + case "response.incomplete": + return "response.incomplete" + case "response.cancelled", "response.canceled": + return "response.cancelled" + default: + return "" + } +} + +func openAIWSPayloadTransientStatus(payload []byte) int { + if len(payload) == 0 { + return 0 + } + status := int(gjson.GetBytes(payload, "response.error.status_code").Int()) + if status == 0 { + status = int(gjson.GetBytes(payload, "response.error.status").Int()) + } + if status == 0 { + status = int(gjson.GetBytes(payload, "error.status_code").Int()) + } + if status == 0 { + status = int(gjson.GetBytes(payload, "error.status").Int()) + } + if shouldCooldownOpenAITransientUpstreamError(status, payload) { + return status + } + if status != 0 { + return 0 + } + code := strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "response.error.code").String())) + errType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "response.error.type").String())) + if code == "" { + code = strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "error.code").String())) + } + if errType == "" { + errType = strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "error.type").String())) + } + switch { + case code == "server_is_overloaded", code == "slow_down": + return http.StatusServiceUnavailable + case strings.Contains(code, "server_error"), + strings.Contains(code, "internal_error"), + strings.Contains(code, "upstream_error"), + strings.Contains(errType, "server_error"), + strings.Contains(errType, "internal_error"), + strings.Contains(errType, "upstream_error"): + return http.StatusInternalServerError + default: + return 0 + } +} + +func (s *OpenAIGatewayService) handleOpenAIWSTerminalTransientFailure(ctx context.Context, account *Account, canonicalModel string, headers http.Header, payload []byte) string { + eventType, _, _ := parseOpenAIWSEventEnvelope(payload) + terminalEvent := normalizeOpenAIWSTerminalEvent(eventType) + if terminalEvent != "response.failed" { + return terminalEvent + } + status := openAIWSPayloadTransientStatus(payload) + if status != 0 { + s.handleOpenAIAccountUpstreamError(ctx, account, status, headers, payload, canonicalModel) + } + return terminalEvent +} + +func (s *OpenAIGatewayService) handleOpenAIWSErrorEventTransientFailure(ctx context.Context, account *Account, canonicalModel string, headers http.Header, payload []byte) { + eventType, _, _ := parseOpenAIWSEventEnvelope(payload) + if eventType != "error" { + return + } + status := openAIWSPayloadTransientStatus(payload) + if status != 0 { + s.handleOpenAIAccountUpstreamError(ctx, account, status, headers, payload, canonicalModel) + } +} + +func (s *OpenAIGatewayService) handleOpenAIWSDialTransientFailure(ctx context.Context, account *Account, canonicalModel string, err error) { + var dialErr *openAIWSDialError + if !errors.As(err, &dialErr) || dialErr == nil || !shouldCooldownOpenAITransientUpstreamError(dialErr.StatusCode, dialErr.ResponseBody) { + return + } + s.handleOpenAIAccountUpstreamError(ctx, account, dialErr.StatusCode, dialErr.ResponseHeaders, dialErr.ResponseBody, canonicalModel) +} + func isOpenAIWSTokenEvent(eventType string) bool { eventType = strings.TrimSpace(eventType) if eventType == "" { @@ -440,7 +532,7 @@ func (s *OpenAIGatewayService) resolveAccountByPreviousResponseIDForCapability( if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, latest); paused { return 0, nil, "", nil } - if s.isOpenAIAccountRuntimeBlocked(latest) { + if s.isOpenAIAccountRequestRuntimeBlocked(latest, requestedModel) { _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) return 0, nil, "", nil } diff --git a/backend/internal/service/openai_ws_forwarder_test.go b/backend/internal/service/openai_ws_forwarder_test.go index b9944e92fa..81c011f7a6 100644 --- a/backend/internal/service/openai_ws_forwarder_test.go +++ b/backend/internal/service/openai_ws_forwarder_test.go @@ -1,9 +1,13 @@ package service import ( + "context" + "net/http" "net/http/httptest" "testing" + "time" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -107,6 +111,81 @@ func TestOpenAIWSCyberPolicyMark_NonCyberPayload(t *testing.T) { require.False(t, hit, "detectOpenAICyberPolicy should return false for non-cyber_policy error code") } +func TestOpenAIForwardResultSucceededForScheduling_TerminalEvents(t *testing.T) { + tests := []struct { + name string + result *OpenAIForwardResult + expected bool + }{ + {name: "nil legacy result", result: nil, expected: true}, + {name: "non websocket zero value", result: &OpenAIForwardResult{}, expected: true}, + {name: "websocket legacy empty terminal", result: &OpenAIForwardResult{OpenAIWSMode: true}, expected: true}, + {name: "completed", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.completed"}, expected: true}, + {name: "done", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.done"}, expected: true}, + {name: "failed", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.failed"}, expected: false}, + {name: "incomplete", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.incomplete"}, expected: false}, + {name: "cancelled", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.cancelled"}, expected: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.expected, tt.result.SucceededForScheduling()) + }) + } +} + +func TestOpenAIWSTerminalEvent_ResponseFailedRecordsModelTransient(t *testing.T) { + svc := &OpenAIGatewayService{} + svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil) + account := &Account{ID: 5201, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + payload := []byte(`{"type":"response.failed","response":{"error":{"code":"server_error","message":"Internal error"}}}`) + + for range 2 { + terminalEvent := svc.handleOpenAIWSTerminalTransientFailure(context.Background(), account, "gpt-5.5", http.Header{}, payload) + require.Equal(t, "response.failed", terminalEvent) + } + + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5")) +} + +func TestOpenAIWSErrorEvent_ServerErrorRecordsModelTransient(t *testing.T) { + svc := &OpenAIGatewayService{} + svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil) + account := &Account{ID: 5203, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + payload := []byte(`{"type":"error","error":{"code":"server_error","type":"server_error","message":"Internal error"}}`) + + for range 2 { + svc.handleOpenAIWSErrorEventTransientFailure(context.Background(), account, "gpt-5.5", http.Header{}, payload) + } + + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5")) +} + +func TestOpenAIWSPayloadTransientStatus_Explicit529IsNotModelTransient(t *testing.T) { + payload := []byte(`{"type":"response.failed","response":{"error":{"status_code":529,"code":"server_error","message":"overloaded"}}}`) + + require.Zero(t, openAIWSPayloadTransientStatus(payload)) +} + +func TestOpenAIWSDial5xxRecordsModelTransient(t *testing.T) { + svc := &OpenAIGatewayService{} + svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil) + account := &Account{ID: 5202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + dialErr := &openAIWSDialError{ + StatusCode: http.StatusBadGateway, + ResponseHeaders: http.Header{"X-Request-Id": []string{"req-ws-502"}}, + ResponseBody: []byte(`{"error":{"message":"bad gateway"}}`), + } + + for range 2 { + svc.handleOpenAIWSDialTransientFailure(context.Background(), account, "gpt-5.5", dialErr) + } + + require.Eventually(t, func() bool { + return svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5") + }, time.Second, 10*time.Millisecond) +} + // TestIsOpenAIWSTokenEvent_DisjointWithTerminal 守护「token 事件集合与终止事件集合互斥」的不变量。 // firstTokenMs 的计算依赖于 isTokenEvent && !isTerminalEvent； // 若两者再次出现交集，则 issue #2651 描述的 latency 误报会重现。 diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 9054e60502..c406b6b337 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -3,6 +3,7 @@ package service import ( "bytes" "context" + "encoding/json" "errors" "fmt" "net/http" @@ -199,6 +200,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } return nil, &agentIdentityTaskRecoveredError{} } + s.handleOpenAIWSDialTransientFailure(ctx, account, mappedModel, err) dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(err) logOpenAIWSModeInfo( "acquire_fail account_id=%d account_type=%s transport=%s reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_new_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v", @@ -351,6 +353,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( flushedBufferedEventCount := 0 firstEventType := "" lastEventType := "" + upstreamTerminalEvent := "" var flusher http.Flusher if reqStream { @@ -454,6 +457,25 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } } } + if readErr == nil && !json.Valid(message) { + eventType, _, _ := parseOpenAIWSEventEnvelope(message) + if eventType == "" { + eventType = "unknown" + } + lease.MarkBroken() + logOpenAIWSModeInfo( + "invalid_event_json account_id=%d conn_id=%s event_type=%s bytes=%d wrote_downstream=%v", + account.ID, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen), + len(message), + wroteDownstream, + ) + if !wroteDownstream { + return nil, wrapOpenAIWSFallback("invalid_event_json", errors.New("upstream websocket returned malformed Responses event JSON")) + } + return nil, errors.New("upstream websocket returned malformed Responses event JSON after downstream output") + } if readErr != nil { lease.MarkBroken() closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) @@ -552,6 +574,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } if eventType == "error" { + s.handleOpenAIWSErrorEventTransientFailure(ctx, account, mappedModel, lease.HandshakeHeaders(), message) errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message) s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw) errMsg := strings.TrimSpace(errMsgRaw) @@ -650,6 +673,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } if isTerminalEvent { + upstreamTerminalEvent = s.handleOpenAIWSTerminalTransientFailure(ctx, account, mappedModel, lease.HandshakeHeaders(), message) // A terminal event must be the final JSON document in its WS message. // Ignore any tail for the completed client turn, but never reuse the // ambiguous upstream connection for another request. @@ -721,19 +745,20 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( ) return &OpenAIForwardResult{ - RequestID: responseID, - Usage: *usage, - Model: originalModel, - UpstreamModel: mappedModel, - ImageCount: imageCounter.Count(), - ImageOutputSizes: imageCounter.Sizes(), - ServiceTier: extractOpenAIServiceTier(reqBody), - ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel), - Stream: reqStream, - OpenAIWSMode: true, - ResponseHeaders: lease.HandshakeHeaders(), - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, + RequestID: responseID, + Usage: *usage, + Model: originalModel, + UpstreamModel: mappedModel, + ImageCount: imageCounter.Count(), + ImageOutputSizes: imageCounter.Sizes(), + ServiceTier: extractOpenAIServiceTier(reqBody), + ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel), + Stream: reqStream, + OpenAIWSMode: true, + UpstreamTerminalEvent: upstreamTerminalEvent, + ResponseHeaders: lease.HandshakeHeaders(), + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, }, nil } diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index 091b3592d1..9140f61443 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -235,6 +235,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes)) if account.Platform == PlatformGrok { s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + } else if shouldCooldownOpenAITransientUpstreamError(resp.StatusCode, respBody) { + canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel) + s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, canonicalModel) } upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) if upstreamMsg == "" { @@ -258,6 +261,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( replayCollector := &openAIWSToolCallReplayCollector{} firstEventType := "" lastEventType := "" + upstreamTerminalEvent := "" sawDone := false wroteDownstream := false clientDisconnected := false @@ -275,17 +279,18 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( resultWithUsage := func() *OpenAIForwardResult { imageCount := imageCounter.Count() result := &OpenAIForwardResult{ - RequestID: responseID, - Usage: usage, - Model: originalModel, - UpstreamModel: mappedModel, - ServiceTier: extractOpenAIServiceTierFromBody(body), - ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, mappedModel, originalModel), body, mappedModel), - Stream: reqStream, - OpenAIWSMode: true, - ResponseHeaders: cloneHeader(resp.Header), - Duration: time.Since(turnStart), - FirstTokenMs: firstTokenMs, + RequestID: responseID, + Usage: usage, + Model: originalModel, + UpstreamModel: mappedModel, + ServiceTier: extractOpenAIServiceTierFromBody(body), + ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, mappedModel, originalModel), body, mappedModel), + Stream: reqStream, + OpenAIWSMode: true, + UpstreamTerminalEvent: upstreamTerminalEvent, + ResponseHeaders: cloneHeader(resp.Header), + Duration: time.Since(turnStart), + FirstTokenMs: firstTokenMs, } if replayInput := replayCollector.Items(); len(replayInput) > 0 { result.wsReplayInput = replayInput @@ -384,6 +389,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( } if eventType == "error" { + s.handleOpenAIWSErrorEventTransientFailure(ctx, account, canonicalOpenAIAccountSchedulingModel(account, originalModel), resp.Header, upstreamMessage) errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) s.persistOpenAIWSRateLimitSignal(ctx, account, resp.Header, upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw) errMessage := strings.TrimSpace(errMsgRaw) @@ -393,6 +399,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( return resultWithUsage(), errors.New(errMessage) } if isOpenAIWSTerminalEvent(eventType) { + upstreamTerminalEvent = s.handleOpenAIWSTerminalTransientFailure(ctx, account, canonicalOpenAIAccountSchedulingModel(account, originalModel), resp.Header, upstreamMessage) terminalEventCount++ firstTokenMsValue := -1 if firstTokenMs != nil { diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index cbc11abcc2..808a63c6b1 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -775,6 +775,14 @@ func TestOpenAIWSHTTPBridge_IdleTimeoutClosesClientSession(t *testing.T) { require.NoError(t, err) require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) + closeReadCtx, cancelCloseRead := context.WithTimeout(context.Background(), 3*time.Second) + _, _, err = clientConn.Read(closeReadCtx) + cancelCloseRead() + var clientClose coderws.CloseError + require.ErrorAs(t, err, &clientClose) + require.Equal(t, coderws.StatusNormalClosure, clientClose.Code) + require.Equal(t, "websocket idle timeout", clientClose.Reason) + select { case proxyErr := <-errCh: var closeErr *OpenAIWSClientCloseError diff --git a/backend/internal/service/openai_ws_ratelimit_signal_test.go b/backend/internal/service/openai_ws_ratelimit_signal_test.go index a3673d7458..9d1fdaf765 100644 --- a/backend/internal/service/openai_ws_ratelimit_signal_test.go +++ b/backend/internal/service/openai_ws_ratelimit_signal_test.go @@ -239,6 +239,52 @@ func TestOpenAIGatewayService_Forward_WSv2Handshake429PersistsRateLimit(t *testi require.Contains(t, repo.updateExtra[0], "codex_usage_updated_at") } +func TestOpenAIGatewayService_Forward_WSv2Handshake502RecordsModelTransient(t *testing.T) { + gin.SetMode(gin.TestMode) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("x-request-id", "req-ws-502") + w.WriteHeader(http.StatusBadGateway) + _, _ = w.Write([]byte(`{"error":{"type":"server_error","message":"bad gateway"}}`)) + })) + defer server.Close() + + cfg := newOpenAIWSV2TestConfig() + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + account := Account{ + ID: 504, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test", "base_url": server.URL}, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + svc := &OpenAIGatewayService{ + cfg: cfg, + rateLimitService: NewRateLimitService(transientCooldownAccountRepo{}, nil, cfg, nil, nil), + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + } + body := []byte(`{"model":"gpt-5.5","stream":false,"input":"hello"}`) + + for range 2 { + 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") + result, err := svc.Forward(context.Background(), c, &account, body) + require.Error(t, err) + require.Nil(t, result) + } + + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(&account, "gpt-5.5")) +} + func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_ErrorEventUsageLimitPersistsRateLimit(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index d2cf26a8ce..4505ffb665 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -386,7 +386,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( if errors.As(err, &handshakeErr) && handshakeErr != nil { responseBody = handshakeErr.Body } - dialErr := &openAIWSDialError{StatusCode: statusCode, ResponseBody: responseBody, Err: err} + dialErr := &openAIWSDialError{StatusCode: statusCode, ResponseHeaders: cloneHeader(handshakeHeaders), ResponseBody: responseBody, Err: err} if s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried { agentTaskRecoveryTried = true if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil { @@ -400,6 +400,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( statusCode, truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), ) + s.handleOpenAIWSDialTransientFailure(ctx, account, capturedSessionModel, dialErr) if statusCode == http.StatusTooManyRequests { s.persistOpenAIWSRateLimitSignal(ctx, account, handshakeHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error())) return &UpstreamFailoverError{ @@ -571,14 +572,15 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( CacheReadInputTokens: turn.Usage.CacheReadInputTokens, ImageOutputTokens: turn.Usage.ImageOutputTokens, }, - Model: turn.RequestModel, - ServiceTier: usageMeta.serviceTier.Load(), - ReasoningEffort: usageMeta.reasoningEffort.Load(), - Stream: true, - OpenAIWSMode: true, - ResponseHeaders: cloneHeader(handshakeHeaders), - Duration: turn.Duration, - FirstTokenMs: turn.FirstTokenMs, + Model: turn.RequestModel, + ServiceTier: usageMeta.serviceTier.Load(), + ReasoningEffort: usageMeta.reasoningEffort.Load(), + Stream: true, + OpenAIWSMode: true, + UpstreamTerminalEvent: normalizeOpenAIWSTerminalEvent(turn.TerminalEventType), + ResponseHeaders: cloneHeader(handshakeHeaders), + Duration: turn.Duration, + FirstTokenMs: turn.FirstTokenMs, } logOpenAIWSV2Passthrough( "relay_turn_completed account_id=%d turn=%d request_id=%s terminal_event=%s duration_ms=%d first_token_ms=%d input_tokens=%d output_tokens=%d cache_read_tokens=%d", @@ -597,10 +599,17 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( } }, BeforeWriteClient: func(msgType coderws.MessageType, payload []byte, wroteDownstream bool) error { - if msgType != coderws.MessageText || wroteDownstream { + if msgType != coderws.MessageText { return nil } - if eventType, _, _ := parseOpenAIWSEventEnvelope(payload); eventType != "error" { + eventType, _, _ := parseOpenAIWSEventEnvelope(payload) + if isOpenAIWSTerminalEvent(eventType) { + s.handleOpenAIWSTerminalTransientFailure(ctx, account, capturedSessionModel, handshakeHeaders, payload) + } + if eventType == "error" { + s.handleOpenAIWSErrorEventTransientFailure(ctx, account, capturedSessionModel, handshakeHeaders, payload) + } + if wroteDownstream || eventType != "error" { return nil } errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(payload) @@ -646,14 +655,15 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( CacheReadInputTokens: relayResult.Usage.CacheReadInputTokens, ImageOutputTokens: relayResult.Usage.ImageOutputTokens, }, - Model: relayResult.RequestModel, - ServiceTier: usageMeta.serviceTier.Load(), - ReasoningEffort: usageMeta.reasoningEffort.Load(), - Stream: true, - OpenAIWSMode: true, - ResponseHeaders: cloneHeader(handshakeHeaders), - Duration: relayResult.Duration, - FirstTokenMs: relayResult.FirstTokenMs, + Model: relayResult.RequestModel, + ServiceTier: usageMeta.serviceTier.Load(), + ReasoningEffort: usageMeta.reasoningEffort.Load(), + Stream: true, + OpenAIWSMode: true, + UpstreamTerminalEvent: normalizeOpenAIWSTerminalEvent(relayResult.TerminalEventType), + ResponseHeaders: cloneHeader(handshakeHeaders), + Duration: relayResult.Duration, + FirstTokenMs: relayResult.FirstTokenMs, } turnCount := int(completedTurns.Load()) diff --git a/backend/internal/service/pricing_service.go b/backend/internal/service/pricing_service.go index d0ae3c7d96..814d6ddaa1 100644 --- a/backend/internal/service/pricing_service.go +++ b/backend/internal/service/pricing_service.go @@ -125,6 +125,7 @@ type LiteLLMModelPricing struct { SupportsPromptCaching bool `json:"supports_prompt_caching"` OutputCostPerImage float64 `json:"output_cost_per_image"` // 图片生成模型每张图片价格 OutputCostPerImageToken float64 `json:"output_cost_per_image_token"` // 图片输出 token 价格 + InputCostPerImageToken float64 `json:"input_cost_per_image_token"` // 图片输入 token 价格（如 gpt-image-2 图片编辑） // TokenPricingAbsent 表示源数据中 input/output token 价格均缺失（仅有图片价）。 // 此类条目只可用于图片计费，token 计费必须回退到 fallback 或 fail-closed， @@ -158,6 +159,7 @@ type LiteLLMRawEntry struct { SupportsPromptCaching bool `json:"supports_prompt_caching"` OutputCostPerImage *float64 `json:"output_cost_per_image"` OutputCostPerImageToken *float64 `json:"output_cost_per_image_token"` + InputCostPerImageToken *float64 `json:"input_cost_per_image_token"` } // PricingService 动态价格服务 @@ -435,7 +437,7 @@ func (s *PricingService) parsePricingData(body []byte) (map[string]*LiteLLMModel } // 只保留有有效价格的条目 - if entry.InputCostPerToken == nil && entry.OutputCostPerToken == nil && entry.OutputCostPerImage == nil && entry.OutputCostPerImageToken == nil { + if entry.InputCostPerToken == nil && entry.OutputCostPerToken == nil && entry.OutputCostPerImage == nil && entry.OutputCostPerImageToken == nil && entry.InputCostPerImageToken == nil { continue } @@ -489,6 +491,9 @@ func (s *PricingService) parsePricingData(body []byte) (map[string]*LiteLLMModel if entry.OutputCostPerImageToken != nil { pricing.OutputCostPerImageToken = *entry.OutputCostPerImageToken } + if entry.InputCostPerImageToken != nil { + pricing.InputCostPerImageToken = *entry.InputCostPerImageToken + } result[modelName] = pricing } diff --git a/backend/internal/service/scheduler_snapshot_batch_query_test.go b/backend/internal/service/scheduler_snapshot_batch_query_test.go index a2e7ae8e85..221610bdbf 100644 --- a/backend/internal/service/scheduler_snapshot_batch_query_test.go +++ b/backend/internal/service/scheduler_snapshot_batch_query_test.go @@ -116,6 +116,72 @@ type batchSnapshotCache struct { beforeSet func() } +type batchSnapshotAccountIDCache struct { + *batchSnapshotCache + + reuseMu sync.Mutex + fullCalls map[SchedulerBucket]int + idOnlyCalls map[SchedulerBucket]int + idOnlyError map[SchedulerBucket]error + fullLateErr map[SchedulerBucket]error + returnEmpty bool +} + +func newBatchSnapshotAccountIDCache() *batchSnapshotAccountIDCache { + return &batchSnapshotAccountIDCache{ + batchSnapshotCache: newBatchSnapshotCache(), + fullCalls: make(map[SchedulerBucket]int), + idOnlyCalls: make(map[SchedulerBucket]int), + idOnlyError: make(map[SchedulerBucket]error), + fullLateErr: make(map[SchedulerBucket]error), + } +} + +func (c *batchSnapshotAccountIDCache) SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accounts []Account) ([]int64, error) { + c.reuseMu.Lock() + c.fullCalls[bucket]++ + c.reuseMu.Unlock() + if err := c.batchSnapshotCache.SetSnapshot(ctx, bucket, token, accounts); err != nil { + return nil, err + } + c.reuseMu.Lock() + lateErr := c.fullLateErr[bucket] + returnEmpty := c.returnEmpty + c.reuseMu.Unlock() + if lateErr != nil { + return nil, lateErr + } + if returnEmpty { + return []int64{}, nil + } + ids := make([]int64, 0, len(accounts)) + for _, account := range accounts { + ids = append(ids, account.ID) + } + return ids, nil +} + +func (c *batchSnapshotAccountIDCache) SetSnapshotByAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accountIDs []int64) error { + c.reuseMu.Lock() + c.idOnlyCalls[bucket]++ + err := c.idOnlyError[bucket] + c.reuseMu.Unlock() + if err != nil { + return err + } + accounts := make([]Account, 0, len(accountIDs)) + for _, id := range accountIDs { + accounts = append(accounts, Account{ID: id}) + } + return c.batchSnapshotCache.SetSnapshot(ctx, bucket, token, accounts) +} + +func (c *batchSnapshotAccountIDCache) reuseCounts(bucket SchedulerBucket) (full, idOnly int) { + c.reuseMu.Lock() + defer c.reuseMu.Unlock() + return c.fullCalls[bucket], c.idOnlyCalls[bucket] +} + func newBatchSnapshotCache() *batchSnapshotCache { return &batchSnapshotCache{ captured: make(map[SchedulerBucket]SchedulerBucketWriteToken), @@ -231,6 +297,186 @@ func TestSchedulerRebuildBatchReusesSingleForcedQueryAndKeepsSnapshotsIndependen } } +func TestSchedulerRebuildBatchReusesAccountPayloadForSingleForced(t *testing.T) { + const groupID int64 = 211 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotAccountIDCache() + repo := newBatchAccountQueryRepo() + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + for run := 1; run <= 2; run++ { + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "reuse")) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, run, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Equal(t, run, idOnly) + } + require.Equal(t, 2, repo.callCount(batchAccountQueryKey{groupID: groupID, platform: PlatformOpenAI}), "账号载荷不得跨重建批次复用") +} + +func TestSchedulerRebuildBatchDoesNotReuseAccountPayloadAfterFirstWriterFailure(t *testing.T) { + const groupID int64 = 212 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + wantErr := errors.New("snapshot write failed") + cache := newBatchSnapshotAccountIDCache() + cache.setErrors[single] = wantErr + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + err := svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "failure") + require.ErrorIs(t, err, wantErr) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, writes := cache.bucketState(forced) + require.Equal(t, 1, attempts, "首次完整写失败后，后续桶必须走原 SetSnapshot") + require.Len(t, writes, 1) +} + +func TestSchedulerRebuildBatchDoesNotReuseAccountPayloadAfterLateFirstWriterFailure(t *testing.T) { + const groupID int64 = 216 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + wantErr := errors.New("snapshot activation failed") + cache := newBatchSnapshotAccountIDCache() + cache.fullLateErr[single] = wantErr + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + err := svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "late-failure") + require.ErrorIs(t, err, wantErr) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, writes := cache.bucketState(forced) + require.Equal(t, 1, attempts, "首次激活失败后不得登记可复用 ID") + require.Len(t, writes, 1) +} + +func TestSchedulerRebuildBatchDoesNotReuseAccountPayloadAfterLockBusy(t *testing.T) { + const groupID int64 = 213 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotAccountIDCache() + cache.lockBusy[single] = true + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "busy")) + full, idOnly := cache.reuseCounts(single) + require.Zero(t, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, writes := cache.bucketState(forced) + require.Equal(t, 1, attempts) + require.Len(t, writes, 1) +} + +func TestSchedulerRebuildBatchKeepsMixedAndDifferentQueriesOnFullWrites(t *testing.T) { + const groupID int64 = 214 + openAISingle := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + openAIForced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + anthropicSingle := SchedulerBucket{GroupID: groupID, Platform: PlatformAnthropic, Mode: SchedulerModeSingle} + anthropicMixed := SchedulerBucket{GroupID: groupID, Platform: PlatformAnthropic, Mode: SchedulerModeMixed} + cache := newBatchSnapshotAccountIDCache() + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{openAISingle, openAIForced, anthropicSingle, anthropicMixed}, "scope")) + full, idOnly := cache.reuseCounts(openAISingle) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(openAIForced) + require.Zero(t, full) + require.Equal(t, 1, idOnly) + full, idOnly = cache.reuseCounts(anthropicSingle) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, writes := cache.bucketState(anthropicSingle) + require.Equal(t, 1, attempts) + require.Len(t, writes, 1) + full, idOnly = cache.reuseCounts(anthropicMixed) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, _ = cache.bucketState(anthropicMixed) + require.Equal(t, 1, attempts, "mixed 桶必须继续走原 SetSnapshot") +} + +func TestSchedulerRebuildBatchPropagatesAccountIDOnlyWriteFailure(t *testing.T) { + const groupID int64 = 215 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + wantErr := errors.New("id-only write failed") + cache := newBatchSnapshotAccountIDCache() + cache.idOnlyError[forced] = wantErr + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + err := svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "id-error") + require.ErrorIs(t, err, wantErr) + full, idOnly := cache.reuseCounts(forced) + require.Zero(t, full, "ID-only 失败不得静默回退为完整写") + require.Equal(t, 1, idOnly) +} + +func TestSchedulerRebuildBatchReusesSuccessfulEmptyAccountIDs(t *testing.T) { + const groupID int64 = 217 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotAccountIDCache() + cache.returnEmpty = true + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "empty")) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Equal(t, 1, idOnly, "已成功缓存的空 ID 集也必须通过 map presence 复用") + _, _, _, writes := cache.bucketState(forced) + require.Len(t, writes, 1) + require.Empty(t, writes[0].accounts) +} + +func TestSchedulerRebuildBatchReusesAccountPayloadForSimpleGroupZero(t *testing.T) { + single := SchedulerBucket{GroupID: 0, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: 0, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotAccountIDCache() + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeSimple) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "simple")) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Equal(t, 1, idOnly) +} + +func TestSchedulerAccountQueryCacheReleasesSnapshotAccountIDs(t *testing.T) { + single := schedulerBucketWriteTask{bucket: SchedulerBucket{GroupID: 218, Platform: PlatformOpenAI, Mode: SchedulerModeSingle}} + forced := schedulerBucketWriteTask{bucket: SchedulerBucket{GroupID: 218, Platform: PlatformOpenAI, Mode: SchedulerModeForced}} + queries := newSchedulerAccountQueryCache([]schedulerBucketWriteTask{single, forced}) + key, ok := schedulerAccountQueryKeyForBucket(single.bucket) + require.True(t, ok) + queries.snapshotAccountIDs[key] = []int64{1, 2} + + queries.release(single.bucket) + require.Contains(t, queries.snapshotAccountIDs, key) + queries.release(forced.bucket) + require.NotContains(t, queries.snapshotAccountIDs, key) + require.Empty(t, queries.remaining) + require.Empty(t, queries.accounts) +} + func TestSchedulerRebuildBatchKeepsMixedAndDifferentKeysIndependent(t *testing.T) { const groupID int64 = 202 buckets := []SchedulerBucket{ diff --git a/backend/internal/service/scheduler_snapshot_bulk_event_test.go b/backend/internal/service/scheduler_snapshot_bulk_event_test.go new file mode 100644 index 0000000000..00454a498d --- /dev/null +++ b/backend/internal/service/scheduler_snapshot_bulk_event_test.go @@ -0,0 +1,211 @@ +//go:build unit + +package service + +import ( + "context" + "sync" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type bulkEventAccountRepo struct { + *batchAccountQueryRepo + accounts []*Account +} + +func newBulkEventAccountRepo(accounts ...*Account) *bulkEventAccountRepo { + return &bulkEventAccountRepo{ + batchAccountQueryRepo: newBatchAccountQueryRepo(), + accounts: accounts, + } +} + +func (r *bulkEventAccountRepo) GetByIDs(context.Context, []int64) ([]*Account, error) { + return append([]*Account(nil), r.accounts...), nil +} + +type bulkEventSnapshotCache struct { + *batchSnapshotCache + + accountMu sync.Mutex + setAccountIDs []int64 + deleteAccountIDs []int64 +} + +func newBulkEventSnapshotCache() *bulkEventSnapshotCache { + return &bulkEventSnapshotCache{batchSnapshotCache: newBatchSnapshotCache()} +} + +func (c *bulkEventSnapshotCache) SetAccount(_ context.Context, account *Account) error { + c.accountMu.Lock() + defer c.accountMu.Unlock() + c.setAccountIDs = append(c.setAccountIDs, account.ID) + return nil +} + +func (c *bulkEventSnapshotCache) DeleteAccount(_ context.Context, accountID int64) error { + c.accountMu.Lock() + defer c.accountMu.Unlock() + c.deleteAccountIDs = append(c.deleteAccountIDs, accountID) + return nil +} + +func (c *bulkEventSnapshotCache) accountWrites() (set []int64, deleted []int64) { + c.accountMu.Lock() + defer c.accountMu.Unlock() + return append([]int64(nil), c.setAccountIDs...), append([]int64(nil), c.deleteAccountIDs...) +} + +func (c *bulkEventSnapshotCache) capturedBuckets() []SchedulerBucket { + c.mu.Lock() + defer c.mu.Unlock() + return append([]SchedulerBucket(nil), c.captures...) +} + +func newBulkEventTestService(cache SchedulerCache, accounts AccountRepository) *SchedulerSnapshotService { + return NewSchedulerSnapshotService(cache, nil, accounts, nil, &config.Config{RunMode: config.RunModeStandard}) +} + +func bulkEventPayload(accountIDs []int64, groupIDs []int64) map[string]any { + accountValues := make([]any, 0, len(accountIDs)) + for _, id := range accountIDs { + accountValues = append(accountValues, id) + } + groupValues := make([]any, 0, len(groupIDs)) + for _, id := range groupIDs { + groupValues = append(groupValues, id) + } + return map[string]any{ + "account_ids": accountValues, + "group_ids": groupValues, + } +} + +func schedulerBucketsForTest(groupIDs []int64, platforms ...string) []SchedulerBucket { + buckets := make([]SchedulerBucket, 0, len(groupIDs)*len(platforms)*3) + for _, platform := range platforms { + for _, groupID := range groupIDs { + buckets = append(buckets, + SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeSingle}, + SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeForced}, + ) + if platform == PlatformAnthropic || platform == PlatformGemini { + buckets = append(buckets, SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeMixed}) + } + } + } + return buckets +} + +func TestSchedulerBulkAccountEventScopesOpenAIRebuildToFreshPlatform(t *testing.T) { + cache := newBulkEventSnapshotCache() + repo := newBulkEventAccountRepo(&Account{ID: 1, Platform: PlatformOpenAI, GroupIDs: []int64{12}}) + svc := newBulkEventTestService(cache, repo) + + err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{1}, []int64{11}), make(map[batchSeenKey]struct{})) + + require.NoError(t, err) + require.ElementsMatch(t, schedulerBucketsForTest([]int64{11, 12}, PlatformOpenAI), cache.capturedBuckets()) + set, deleted := cache.accountWrites() + require.Equal(t, []int64{1}, set) + require.Empty(t, deleted) +} + +func TestSchedulerBulkAccountEventRebuildsOpenAIUngroupedBucket(t *testing.T) { + cache := newBulkEventSnapshotCache() + repo := newBulkEventAccountRepo(&Account{ID: 6, Platform: PlatformOpenAI}) + svc := newBulkEventTestService(cache, repo) + + err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{6}, nil), make(map[batchSeenKey]struct{})) + + require.NoError(t, err) + require.ElementsMatch(t, schedulerBucketsForTest([]int64{0}, PlatformOpenAI), cache.capturedBuckets()) +} + +func TestSchedulerBulkAccountEventKeepsGroupedAndUngroupedBuckets(t *testing.T) { + cache := newBulkEventSnapshotCache() + repo := newBulkEventAccountRepo( + &Account{ID: 7, Platform: PlatformOpenAI, GroupIDs: []int64{51}}, + &Account{ID: 8, Platform: PlatformOpenAI}, + ) + svc := newBulkEventTestService(cache, repo) + + err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{7, 8}, nil), make(map[batchSeenKey]struct{})) + + require.NoError(t, err) + require.ElementsMatch(t, schedulerBucketsForTest([]int64{0, 51}, PlatformOpenAI), cache.capturedBuckets()) +} + +func TestSchedulerBulkAccountEventDoesNotCrossCurrentGroupsBetweenPlatforms(t *testing.T) { + cache := newBulkEventSnapshotCache() + repo := newBulkEventAccountRepo( + &Account{ID: 9, Platform: PlatformOpenAI, GroupIDs: []int64{61}}, + &Account{ID: 10, Platform: PlatformGrok, GroupIDs: []int64{62}}, + ) + svc := newBulkEventTestService(cache, repo) + + err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{9, 10}, []int64{63}), make(map[batchSeenKey]struct{})) + + require.NoError(t, err) + want := append( + schedulerBucketsForTest([]int64{61, 63}, PlatformOpenAI), + schedulerBucketsForTest([]int64{62, 63}, PlatformGrok)..., + ) + require.ElementsMatch(t, want, cache.capturedBuckets()) +} + +func TestSchedulerBulkAccountEventUsesGroupZeroInSimpleMode(t *testing.T) { + cache := newBulkEventSnapshotCache() + repo := newBulkEventAccountRepo(&Account{ID: 11, Platform: PlatformOpenAI, GroupIDs: []int64{71}}) + svc := NewSchedulerSnapshotService(cache, nil, repo, nil, &config.Config{RunMode: config.RunModeSimple}) + + err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{11}, []int64{72}), make(map[batchSeenKey]struct{})) + + require.NoError(t, err) + require.ElementsMatch(t, schedulerBucketsForTest([]int64{0}, PlatformOpenAI), cache.capturedBuckets()) +} + +func TestSchedulerBulkAccountEventConservativelyExpandsAntigravityPlatforms(t *testing.T) { + cache := newBulkEventSnapshotCache() + // fresh 值可能已经关闭 mixed_scheduling，兼容平台仍要重建以清理旧快照。 + repo := newBulkEventAccountRepo(&Account{ID: 2, Platform: PlatformAntigravity, GroupIDs: []int64{22}}) + svc := newBulkEventTestService(cache, repo) + + err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{2}, []int64{21}), make(map[batchSeenKey]struct{})) + + require.NoError(t, err) + require.ElementsMatch(t, + schedulerBucketsForTest([]int64{21, 22}, PlatformAnthropic, PlatformGemini, PlatformAntigravity), + cache.capturedBuckets(), + ) +} + +func TestSchedulerBulkAccountEventMissingAccountFallsBackToAllPlatforms(t *testing.T) { + cache := newBulkEventSnapshotCache() + repo := newBulkEventAccountRepo(&Account{ID: 3, Platform: PlatformOpenAI, GroupIDs: []int64{32}}) + svc := newBulkEventTestService(cache, repo) + + err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{3, 4}, []int64{31}), make(map[batchSeenKey]struct{})) + + require.NoError(t, err) + platforms := schedulerSnapshotPlatforms() + require.ElementsMatch(t, schedulerBucketsForTest([]int64{31, 32}, platforms[:]...), cache.capturedBuckets()) + set, deleted := cache.accountWrites() + require.Equal(t, []int64{3}, set) + require.Equal(t, []int64{4}, deleted) +} + +func TestSchedulerBulkAccountEventUnknownPlatformFallsBackToAllPlatforms(t *testing.T) { + cache := newBulkEventSnapshotCache() + repo := newBulkEventAccountRepo(&Account{ID: 5, GroupIDs: []int64{42}}) + svc := newBulkEventTestService(cache, repo) + + err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{5}, []int64{41}), make(map[batchSeenKey]struct{})) + + require.NoError(t, err) + platforms := schedulerSnapshotPlatforms() + require.ElementsMatch(t, schedulerBucketsForTest([]int64{41, 42}, platforms[:]...), cache.capturedBuckets()) +} diff --git a/backend/internal/service/scheduler_snapshot_service.go b/backend/internal/service/scheduler_snapshot_service.go index a5bccdb265..18c1d523d4 100644 --- a/backend/internal/service/scheduler_snapshot_service.go +++ b/backend/internal/service/scheduler_snapshot_service.go @@ -55,14 +55,24 @@ type schedulerAccountQueryKey struct { // mixed 与历史模式保持独立。每个 task 都用 defer 消费 remaining，最后一个消费者会立即释放结果， // 避免把账号切片的生命周期扩大到整轮 full rebuild。 type schedulerAccountQueryCache struct { - remaining map[schedulerAccountQueryKey]int - accounts map[schedulerAccountQueryKey][]Account + remaining map[schedulerAccountQueryKey]int + accounts map[schedulerAccountQueryKey][]Account + snapshotAccountIDs map[schedulerAccountQueryKey][]int64 +} + +// schedulerSnapshotAccountIDWriter 是 SchedulerCache 的可选批次优化能力。 +// 首次完整发布成功后返回实际可编码账号 ID；同一查询结果的后续桶只需发布这些 ID， +// 避免重复序列化并覆盖全局账号缓存。未实现该接口的缓存继续走原 SetSnapshot 路径。 +type schedulerSnapshotAccountIDWriter interface { + SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accounts []Account) ([]int64, error) + SetSnapshotByAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accountIDs []int64) error } func newSchedulerAccountQueryCache(taskSets ...[]schedulerBucketWriteTask) *schedulerAccountQueryCache { queries := &schedulerAccountQueryCache{ - remaining: make(map[schedulerAccountQueryKey]int), - accounts: make(map[schedulerAccountQueryKey][]Account), + remaining: make(map[schedulerAccountQueryKey]int), + accounts: make(map[schedulerAccountQueryKey][]Account), + snapshotAccountIDs: make(map[schedulerAccountQueryKey][]int64), } for _, tasks := range taskSets { for _, task := range tasks { @@ -93,6 +103,7 @@ func (c *schedulerAccountQueryCache) release(bucket SchedulerBucket) { if remaining <= 0 { delete(c.remaining, key) delete(c.accounts, key) + delete(c.snapshotAccountIDs, key) return } c.remaining[key] = remaining @@ -529,11 +540,13 @@ func (s *SchedulerSnapshotService) handleBulkAccountEvent(ctx context.Context, p } } - if s.cache != nil { - for _, id := range ids { - if _, ok := found[id]; ok { - continue - } + allAccountsFound := true + for _, id := range ids { + if _, ok := found[id]; ok { + continue + } + allAccountsFound = false + if s.cache != nil { if err := s.cache.DeleteAccount(ctx, id); err != nil { return err } @@ -544,7 +557,67 @@ func (s *SchedulerSnapshotService) handleBulkAccountEvent(ctx context.Context, p for gid := range rebuildGroupSet { rebuildGroupIDs = append(rebuildGroupIDs, gid) } - return s.rebuildByGroupIDs(ctx, rebuildGroupIDs, "account_bulk_change", seen) + + // 缺失账户无法确定原平台，保留五平台重建以避免遗留旧快照。 + if !allAccountsFound { + return s.rebuildByGroupIDs(ctx, rebuildGroupIDs, "account_bulk_change", seen) + } + + platformGroupSets := make(map[string]map[int64]struct{}, len(accounts)) + addPlatformGroups := func(platform string, groupIDs []int64) { + groupSet := platformGroupSets[platform] + if groupSet == nil { + groupSet = make(map[int64]struct{}, len(groupIDs)) + platformGroupSets[platform] = groupSet + } + for _, groupID := range groupIDs { + groupSet[groupID] = struct{}{} + } + } + for _, account := range accounts { + if account == nil || account.ID <= 0 { + continue + } + accountGroupIDs := s.normalizeGroupIDs(account.GroupIDs) + switch account.Platform { + case PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformGrok: + addPlatformGroups(account.Platform, accountGroupIDs) + case PlatformAntigravity: + // 批量更新可能刚关闭 mixed_scheduling，仍需清理两个兼容平台的旧快照。 + addPlatformGroups(PlatformAntigravity, accountGroupIDs) + addPlatformGroups(PlatformAnthropic, accountGroupIDs) + addPlatformGroups(PlatformGemini, accountGroupIDs) + default: + return s.rebuildByGroupIDs(ctx, rebuildGroupIDs, "account_bulk_change", seen) + } + } + + // payload 携带更新前的组；只扩散到本事件实际涉及的平台，避免平台间交叉重建。 + if len(preloadGroupIDs) > 0 { + preloadGroupIDs = s.normalizeGroupIDs(preloadGroupIDs) + for platform := range platformGroupSets { + addPlatformGroups(platform, preloadGroupIDs) + } + } + + bucketCapacity := 0 + for _, groupSet := range platformGroupSets { + bucketCapacity += len(groupSet) * 3 + } + buckets := make([]SchedulerBucket, 0, bucketCapacity) + for _, platform := range schedulerSnapshotPlatforms() { + groupSet, ok := platformGroupSets[platform] + if !ok { + continue + } + platformGroupIDs := make([]int64, 0, len(groupSet)) + for groupID := range groupSet { + platformGroupIDs = append(platformGroupIDs, groupID) + } + sort.Slice(platformGroupIDs, func(i, j int) bool { return platformGroupIDs[i] < platformGroupIDs[j] }) + buckets = append(buckets, s.bucketsForPlatform(platform, platformGroupIDs, seen)...) + } + return s.rebuildBuckets(ctx, buckets, "account_bulk_change") } func (s *SchedulerSnapshotService) handleAccountEvent(ctx context.Context, accountID *int64, payload map[string]any, seen map[batchSeenKey]struct{}) error { @@ -869,7 +942,7 @@ func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicyAndQueryCache( logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild failed: bucket=%s reason=%s err=%v", bucket.String(), reason, err) return err } - if err := s.cache.SetSnapshot(rebuildCtx, bucket, task.token, accounts); err != nil { + if err := s.setRebuildSnapshot(rebuildCtx, task, accounts, queries); err != nil { if errors.Is(err, ErrSchedulerBucketRetired) || errors.Is(err, ErrSchedulerBucketWriteFenced) { slog.Debug("[Scheduler] rebuild fenced", "bucket", bucket.String(), "reason", reason) if strict { @@ -884,6 +957,38 @@ func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicyAndQueryCache( return nil } +func (s *SchedulerSnapshotService) setRebuildSnapshot( + ctx context.Context, + task schedulerBucketWriteTask, + accounts []Account, + queries *schedulerAccountQueryCache, +) error { + writer, ok := s.cache.(schedulerSnapshotAccountIDWriter) + key, reusable := schedulerAccountQueryKeyForBucket(task.bucket) + if !ok || queries == nil || !reusable { + return s.cache.SetSnapshot(ctx, task.bucket, task.token, accounts) + } + + if accountIDs, exists := queries.snapshotAccountIDs[key]; exists { + return writer.SetSnapshotByAccountIDs(ctx, task.bucket, task.token, accountIDs) + } + if queries.remaining[key] <= 1 { + return s.cache.SetSnapshot(ctx, task.bucket, task.token, accounts) + } + + accountIDs, err := writer.SetSnapshotAndReturnAccountIDs(ctx, task.bucket, task.token, accounts) + if err != nil { + return err + } + if queries.remaining[key] > 1 { + // 必须保存 writeAccounts 实际接受的有序 ID，不能从原账号切片重新推导； + // 否则不可编码账号会只出现在后续桶中，破坏两个快照的成员一致性。 + // 返回切片由当前批次独占，直接接管可避免 10k 账号场景再次复制。 + queries.snapshotAccountIDs[key] = accountIDs + } + return nil +} + func (s *SchedulerSnapshotService) triggerFullRebuild(reason string) error { if s.cache == nil { return ErrSchedulerCacheNotReady diff --git a/backend/internal/service/usage_log.go b/backend/internal/service/usage_log.go index 0adcc04a94..4e8df04aaf 100644 --- a/backend/internal/service/usage_log.go +++ b/backend/internal/service/usage_log.go @@ -139,6 +139,8 @@ type UsageLog struct { CacheCreation5mTokens int `gorm:"column:cache_creation_5m_tokens"` CacheCreation1hTokens int `gorm:"column:cache_creation_1h_tokens"` + ImageInputTokens int + ImageInputCost float64 ImageOutputTokens int ImageOutputCost float64 diff --git a/backend/migrations/178_channel_image_input_price.sql b/backend/migrations/178_channel_image_input_price.sql new file mode 100644 index 0000000000..34a8f885e9 --- /dev/null +++ b/backend/migrations/178_channel_image_input_price.sql @@ -0,0 +1,8 @@ +-- 178_channel_image_input_price.sql +-- 渠道自定义定价（token 模式）新增图片输入 token 单价列。 +-- 用于 gpt-image-2 等模型的图片编辑/图生图请求：上游 usage 的 +-- input_tokens_details.image_tokens 需按独立单价计费，区别于文本 input_price。 +-- 未配置（NULL）时计费回退到文本输入价，保持向后兼容。 + +ALTER TABLE channel_model_pricing + ADD COLUMN IF NOT EXISTS image_input_price NUMERIC(20,12); diff --git a/backend/migrations/179_usage_log_image_input_tokens.sql b/backend/migrations/179_usage_log_image_input_tokens.sql new file mode 100644 index 0000000000..99706da3f5 --- /dev/null +++ b/backend/migrations/179_usage_log_image_input_tokens.sql @@ -0,0 +1,6 @@ +-- 179_usage_log_image_input_tokens.sql +-- usage_logs 单独记录图片输入 token 数与费用，便于图片编辑/图生图场景对账。 +-- image_input_tokens 从 input_tokens 中拆出，image_input_cost 从 input_cost 中拆出， +-- total_cost 口径不变。 +ALTER TABLE usage_logs ADD COLUMN IF NOT EXISTS image_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE usage_logs ADD COLUMN IF NOT EXISTS image_input_cost DECIMAL(20, 10) NOT NULL DEFAULT 0; diff --git a/frontend/src/api/admin/channels.ts b/frontend/src/api/admin/channels.ts index afa43a2d7b..fdbeadf57a 100644 --- a/frontend/src/api/admin/channels.ts +++ b/frontend/src/api/admin/channels.ts @@ -30,6 +30,7 @@ export interface ChannelModelPricing { output_price: number | null cache_write_price: number | null cache_read_price: number | null + image_input_price: number | null image_output_price: number | null per_request_price: number | null intervals: PricingInterval[] @@ -154,6 +155,7 @@ export interface ModelDefaultPricing { output_price?: number cache_write_price?: number cache_read_price?: number + image_input_price?: number image_output_price?: number } diff --git a/frontend/src/api/channels.ts b/frontend/src/api/channels.ts index 8eaf7eed95..46e350809f 100644 --- a/frontend/src/api/channels.ts +++ b/frontend/src/api/channels.ts @@ -39,6 +39,7 @@ export interface UserSupportedModelPricing { output_price: number | null cache_write_price: number | null cache_read_price: number | null + image_input_price: number | null image_output_price: number | null per_request_price: number | null intervals: UserPricingInterval[] diff --git a/frontend/src/components/admin/channel/PricingEntryCard.vue b/frontend/src/components/admin/channel/PricingEntryCard.vue index 9fc103ef56..086a9cd307 100644 --- a/frontend/src/components/admin/channel/PricingEntryCard.vue +++ b/frontend/src/components/admin/channel/PricingEntryCard.vue @@ -101,7 +101,7 @@ {{ t('admin.channels.form.defaultPrices') }} $/MTok -
+
+
+ + +
R
+
+
+ + {{ row.image_input_tokens.toLocaleString() }} +
+
@@ -246,10 +252,18 @@
{{ t('usage.tokenDetails') }}
-
+
{{ t('admin.usage.inputTokens') }} {{ tokenTooltipData.input_tokens.toLocaleString() }}
+
+ {{ t('admin.usage.inputTokens') }} + {{ textInputTokens(tokenTooltipData).toLocaleString() }} +
+
+ {{ t('usage.imageInputTokens') }} + {{ tokenTooltipData.image_input_tokens.toLocaleString() }} +
{{ t('admin.usage.outputTokens') }} {{ tokenTooltipData.output_tokens.toLocaleString() }} @@ -327,6 +341,10 @@ {{ t('admin.usage.inputCost') }} ${{ tooltipData.input_cost.toFixed(6) }}
+
+ {{ t('usage.imageInputCost') }} + ${{ tooltipData.image_input_cost.toFixed(6) }} +
{{ t('admin.usage.outputCost') }} ${{ tooltipData.output_cost.toFixed(6) }} @@ -337,9 +355,13 @@