From 716fcc6f3d6591e602924edcc2a8cc0054b49692 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=8F?= <2829624376@qq.com> Date: Wed, 15 Jul 2026 20:59:56 +0800 Subject: [PATCH 01/26] fix(openai): reject malformed WS v2 events --- .../handler/openai_gateway_handler_test.go | 37 +++++ .../openai_sse_concatenated_json_test.go | 132 ++++++++++++++++++ .../service/openai_ws_forwarder_v2.go | 20 +++ 3 files changed, 189 insertions(+) diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index d855f364a7..66bffe5a83 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -221,6 +221,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() 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_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 9054e60502..eefadbb82d 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" @@ -454,6 +455,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) From 858f3e4a78ecc974b3fedd6af27e0afe7d5b69cc Mon Sep 17 00:00:00 2001 From: drgnchan Date: Wed, 15 Jul 2026 21:57:08 +0800 Subject: [PATCH 02/26] fix(frontend): preserve horizontal table scrolling on tablets --- .../src/components/layout/TablePageLayout.vue | 4 --- .../layout/__tests__/TablePageLayout.spec.ts | 26 +++++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) create mode 100644 frontend/src/components/layout/__tests__/TablePageLayout.spec.ts diff --git a/frontend/src/components/layout/TablePageLayout.vue b/frontend/src/components/layout/TablePageLayout.vue index e4d3d447d1..3620697add 100644 --- a/frontend/src/components/layout/TablePageLayout.vue +++ b/frontend/src/components/layout/TablePageLayout.vue @@ -100,10 +100,6 @@ onUnmounted(() => { @apply flex-none min-h-fit; } -.table-page-layout.mobile-mode .table-scroll-container :deep(.table-wrapper) { - @apply overflow-visible; -} - .table-page-layout.mobile-mode .table-scroll-container :deep(table) { @apply flex-none; display: table; diff --git a/frontend/src/components/layout/__tests__/TablePageLayout.spec.ts b/frontend/src/components/layout/__tests__/TablePageLayout.spec.ts new file mode 100644 index 0000000000..3631d09313 --- /dev/null +++ b/frontend/src/components/layout/__tests__/TablePageLayout.spec.ts @@ -0,0 +1,26 @@ +import { readFileSync } from 'node:fs' +import { dirname, resolve } from 'node:path' +import { fileURLToPath } from 'node:url' + +import { describe, expect, it } from 'vitest' + +const componentPath = resolve(dirname(fileURLToPath(import.meta.url)), '../TablePageLayout.vue') +const componentSource = readFileSync(componentPath, 'utf8') + +describe('TablePageLayout responsive table scrolling', () => { + it('does not disable the table horizontal scroll container in mobile mode', () => { + const tableWrapperBlocks = Array.from( + componentSource.matchAll(/([^{}]*:deep\(\.table-wrapper\)[^{}]*)\{([^{}]*)\}/g) + ) + + expect(tableWrapperBlocks.length).toBeGreaterThan(0) + + const baseBlock = tableWrapperBlocks.find(([selector]) => !selector.includes('.mobile-mode')) + const mobileBlocks = tableWrapperBlocks.filter(([selector]) => selector.includes('.mobile-mode')) + + expect(baseBlock?.[2]).toContain('overflow-x-auto') + expect(mobileBlocks.every(([, , declarations]) => !declarations.includes('overflow-visible'))).toBe( + true + ) + }) +}) From 134179085c1df63c8e51a32d39a273882bfec1aa Mon Sep 17 00:00:00 2001 From: haruka <1628615876@qq.com> Date: Wed, 15 Jul 2026 22:13:37 +0800 Subject: [PATCH 03/26] feat: add async image generation tasks --- backend/cmd/server/wire_gen.go | 7 +- backend/go.sum | 2 + backend/internal/handler/endpoint.go | 3 + backend/internal/handler/endpoint_test.go | 2 + backend/internal/handler/handler.go | 1 + .../internal/handler/image_task_handler.go | 270 ++++++++++++++++++ .../handler/image_task_handler_test.go | 107 +++++++ backend/internal/handler/wire.go | 3 + .../internal/repository/image_task_store.go | 48 ++++ .../repository/image_task_store_test.go | 43 +++ backend/internal/repository/wire.go | 1 + .../server/middleware/api_key_auth.go | 17 +- .../api_key_auth_image_task_test.go | 15 + backend/internal/server/routes/gateway.go | 6 + .../internal/server/routes/gateway_test.go | 20 ++ backend/internal/service/image_task.go | 210 ++++++++++++++ backend/internal/service/image_task_test.go | 91 ++++++ backend/internal/service/wire.go | 2 + 18 files changed, 843 insertions(+), 5 deletions(-) create mode 100644 backend/internal/handler/image_task_handler.go create mode 100644 backend/internal/handler/image_task_handler_test.go create mode 100644 backend/internal/repository/image_task_store.go create mode 100644 backend/internal/repository/image_task_store_test.go create mode 100644 backend/internal/server/middleware/api_key_auth_image_task_test.go create mode 100644 backend/internal/service/image_task.go create mode 100644 backend/internal/service/image_task_test.go diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 0c466a6e8a..996db5b7c1 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -70,7 +70,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { concurrencyCache := repository.ProvideConcurrencyCache(redisClient, configConfig) schedulerCache := repository.ProvideSchedulerCache(redisClient, configConfig) accountRepository := repository.NewAccountRepository(client, db, schedulerCache) - adminAccountRepository := repository.NewAdminAccountRepository(client, db, schedulerCache) concurrencyService := service.ProvideConcurrencyService(concurrencyCache, accountRepository, configConfig) apiKeyService := service.ProvideAPIKeyService(apiKeyRepository, userRepository, groupRepository, userSubscriptionRepository, userGroupRateRepository, apiKeyCache, configConfig, billingCacheService, concurrencyService) apiKeyAuthCacheInvalidator := service.ProvideAPIKeyAuthCacheInvalidator(apiKeyService) @@ -174,6 +173,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { leaderLockCache := repository.NewLeaderLockCache(redisClient) dashboardAggregationService := service.ProvideDashboardAggregationService(dashboardAggregationRepository, timingWheelService, leaderLockCache, db, configConfig) dashboardHandler := admin.NewDashboardHandler(dashboardService, dashboardAggregationService) + adminAccountRepository := repository.NewAdminAccountRepository(client, db, schedulerCache) proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig) proxyLatencyCache := repository.NewProxyLatencyCache(redisClient) adminService := service.NewAdminService(userRepository, groupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService) @@ -261,6 +261,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService) paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry) availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService) + imageTaskStore := repository.NewImageTaskStore(redisClient) + imageTaskService := service.NewImageTaskService(imageTaskStore) + asyncImageHandler := handler.NewAsyncImageHandler(imageTaskService, openAIGatewayHandler) batchImageRepository := repository.NewBatchImageRepository(db) batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig) batchImageModelPricingResolver := service.ProvideBatchImageModelPricingResolver(modelPricingResolver) @@ -271,7 +274,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { batchImageHandler := handler.NewBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService) idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig) idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig) - handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, batchImageHandler, idempotencyCoordinator, idempotencyCleanupService) + handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, asyncImageHandler, batchImageHandler, idempotencyCoordinator, idempotencyCleanupService) jwtAuthMiddleware := middleware.NewJWTAuthMiddleware(authService, userService) adminAuthMiddleware := middleware.NewAdminAuthMiddleware(authService, userService, settingService) apiKeyAuthMiddleware := middleware.NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, configConfig) diff --git a/backend/go.sum b/backend/go.sum index 3d9989bb43..104dc56a4e 100644 --- a/backend/go.sum +++ b/backend/go.sum @@ -166,6 +166,8 @@ github.com/google/go-querystring v1.1.0/go.mod h1:Kcdr2DB4koayq7X8pmAG4sNG59So17 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= +github.com/google/subcommands v1.2.0 h1:vWQspBTo2nEqTUFita5/KeEWlUL8kQObDFbub/EN9oE= +github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/wire v0.7.0 h1:JxUKI6+CVBgCO2WToKy/nQk0sS+amI9z9EjVmdaocj4= diff --git a/backend/internal/handler/endpoint.go b/backend/internal/handler/endpoint.go index 5e4d84ba72..0e0d28bb81 100644 --- a/backend/internal/handler/endpoint.go +++ b/backend/internal/handler/endpoint.go @@ -23,6 +23,7 @@ const ( EndpointResponsesCompact = "/v1/responses/compact" EndpointImagesGenerations = "/v1/images/generations" EndpointImagesEdits = "/v1/images/edits" + EndpointImageTasks = "/v1/images/tasks" EndpointVideosGenerations = "/v1/videos/generations" EndpointVideosEdits = "/v1/videos/edits" EndpointVideosExtensions = "/v1/videos/extensions" @@ -88,6 +89,8 @@ func NormalizeInboundEndpoint(path string) string { return EndpointImagesGenerations case strings.Contains(path, EndpointImagesEdits) || strings.Contains(path, "/images/edits"): return EndpointImagesEdits + case strings.Contains(path, EndpointImageTasks) || strings.Contains(path, "/images/tasks/"): + return EndpointImageTasks case strings.Contains(path, EndpointVideosGenerations) || strings.Contains(path, "/videos/generations"): return EndpointVideosGenerations case strings.Contains(path, EndpointVideosEdits) || strings.Contains(path, "/videos/edits"): diff --git a/backend/internal/handler/endpoint_test.go b/backend/internal/handler/endpoint_test.go index dc8ad728c8..c8ae566ad5 100644 --- a/backend/internal/handler/endpoint_test.go +++ b/backend/internal/handler/endpoint_test.go @@ -31,6 +31,7 @@ func TestNormalizeInboundEndpoint(t *testing.T) { {"/v1/responses/compact/detail", EndpointResponsesCompact}, {"/v1/images/generations", EndpointImagesGenerations}, {"/v1/images/edits", EndpointImagesEdits}, + {"/v1/images/tasks/imgtask_123", EndpointImageTasks}, {"/v1/videos/generations", EndpointVideosGenerations}, {"/v1/videos/req_123", EndpointVideos}, {"/v1beta/models", EndpointGeminiModels}, @@ -52,6 +53,7 @@ func TestNormalizeInboundEndpoint(t *testing.T) { {"/responses/compact", EndpointResponsesCompact}, {"/responses/compact/detail", EndpointResponsesCompact}, {"/alpha/search", EndpointAlphaSearch}, + {"/images/tasks/imgtask_123", EndpointImageTasks}, // Bare Codex direct alias route — root vs. compact. {"/backend-api/codex/responses", EndpointResponses}, diff --git a/backend/internal/handler/handler.go b/backend/internal/handler/handler.go index 58c524889c..d8ffdb2865 100644 --- a/backend/internal/handler/handler.go +++ b/backend/internal/handler/handler.go @@ -58,6 +58,7 @@ type Handlers struct { Payment *PaymentHandler PaymentWebhook *PaymentWebhookHandler AvailableChannel *AvailableChannelHandler + AsyncImage *AsyncImageHandler BatchImage *BatchImageHandler } diff --git a/backend/internal/handler/image_task_handler.go b/backend/internal/handler/image_task_handler.go new file mode 100644 index 0000000000..adb46d2531 --- /dev/null +++ b/backend/internal/handler/image_task_handler.go @@ -0,0 +1,270 @@ +package handler + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type AsyncImageHandler struct { + tasks *service.ImageTaskService + openAI *OpenAIGatewayHandler + execute func(platform string, c *gin.Context) +} + +func NewAsyncImageHandler(tasks *service.ImageTaskService, openAI *OpenAIGatewayHandler) *AsyncImageHandler { + h := &AsyncImageHandler{tasks: tasks, openAI: openAI} + h.execute = h.executeWithGateway + return h +} + +// Submit accepts the same payload as the synchronous Images endpoint and +// returns before the upstream image generation begins. +func (h *AsyncImageHandler) Submit(c *gin.Context) { + apiKey, ok := middleware2.GetAPIKeyFromContext(c) + if !ok || apiKey == nil || apiKey.UserID <= 0 || apiKey.ID <= 0 { + imageTaskError(c, service.ErrImageTaskForbidden) + return + } + platform := "" + if apiKey.Group != nil { + platform = apiKey.Group.Platform + } + if platform != service.PlatformOpenAI && platform != service.PlatformGrok { + imageTaskJSONError(c, http.StatusNotFound, "not_found_error", "Images API is not supported for this platform") + return + } + if !service.GroupAllowsImageGeneration(apiKey.Group) { + imageTaskJSONError(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage()) + return + } + if h == nil || h.tasks == nil || h.execute == nil { + imageTaskError(c, service.ErrImageTaskUnavailable) + return + } + + body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + if err != nil { + if maxErr, ok := extractMaxBytesError(err); ok { + imageTaskJSONError(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) + return + } + imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body") + return + } + if len(body) == 0 { + imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty") + return + } + if asyncImageRequestStreams(c.GetHeader("Content-Type"), body) { + imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", "streaming image requests cannot be submitted as asynchronous tasks") + return + } + if err := h.validateRequest(c, platform, body); err != nil { + imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", err.Error()) + return + } + + taskCtx, recorder, cancel := newAsyncImageContext(c, body, h.tasks.ExecutionTimeout()) + task, err := h.tasks.Create(c.Request.Context(), service.ImageTaskOwner{UserID: apiKey.UserID, APIKeyID: apiKey.ID}) + if err != nil { + cancel() + imageTaskError(c, err) + return + } + + pollURL := imageTaskPollURL(c.Request.URL.Path, task.ID) + c.Header("Cache-Control", "no-store") + c.Header("Location", pollURL) + c.Header("Retry-After", "3") + c.JSON(http.StatusAccepted, gin.H{ + "id": task.ID, + "task_id": task.TaskID, + "object": task.Object, + "status": task.Status, + "created_at": task.CreatedAt, + "expires_at": task.ExpiresAt, + "poll_url": pollURL, + }) + + go h.run(task.ID, platform, taskCtx, recorder, cancel) +} + +func (h *AsyncImageHandler) Get(c *gin.Context) { + apiKey, ok := middleware2.GetAPIKeyFromContext(c) + if !ok || apiKey == nil || apiKey.UserID <= 0 || apiKey.ID <= 0 { + imageTaskError(c, service.ErrImageTaskForbidden) + return + } + if h == nil || h.tasks == nil { + imageTaskError(c, service.ErrImageTaskUnavailable) + return + } + task, err := h.tasks.Get(c.Request.Context(), service.ImageTaskOwner{UserID: apiKey.UserID, APIKeyID: apiKey.ID}, c.Param("task_id")) + if err != nil { + imageTaskError(c, err) + return + } + c.Header("Cache-Control", "no-store") + if task.Status == service.ImageTaskStatusProcessing { + c.Header("Retry-After", "3") + } + c.JSON(http.StatusOK, task) +} + +func (h *AsyncImageHandler) validateRequest(c *gin.Context, platform string, body []byte) error { + if h.openAI == nil || h.openAI.gatewayService == nil { + return nil + } + if platform == service.PlatformGrok { + parsed := service.ParseGrokMediaRequest(c.GetHeader("Content-Type"), body) + if strings.TrimSpace(parsed.Model) == "" { + return errors.New("model is required") + } + return nil + } + parsed, err := h.openAI.gatewayService.ParseOpenAIImagesRequest(c, body) + if err != nil { + return err + } + if parsed.Stream { + return errors.New("streaming image requests cannot be submitted as asynchronous tasks") + } + return nil +} + +func (h *AsyncImageHandler) executeWithGateway(platform string, c *gin.Context) { + if h.openAI == nil { + imageTaskJSONError(c, http.StatusServiceUnavailable, "api_error", "image gateway is unavailable") + return + } + if platform == service.PlatformGrok { + h.openAI.GrokImages(c) + return + } + h.openAI.Images(c) +} + +func (h *AsyncImageHandler) run(taskID, platform string, taskCtx *gin.Context, recorder *httptest.ResponseRecorder, cancel context.CancelFunc) { + defer cancel() + defer func() { + if recovered := recover(); recovered != nil { + logger.L().Error("image_task.execution_panicked", zap.String("task_id", taskID), zap.Any("panic", recovered)) + h.failTask(taskID, http.StatusInternalServerError, imageTaskErrorPayload("api_error", "image generation task panicked")) + } + }() + + h.execute(platform, taskCtx) + body := bytes.TrimSpace(recorder.Body.Bytes()) + if err := taskCtx.Request.Context().Err(); err != nil && len(body) == 0 { + h.failTask(taskID, http.StatusGatewayTimeout, imageTaskErrorPayload("timeout_error", "image generation task timed out")) + return + } + statusCode := recorder.Code + if statusCode == 0 { + statusCode = http.StatusOK + } + if statusCode >= http.StatusOK && statusCode < http.StatusMultipleChoices { + if len(body) == 0 || !json.Valid(body) { + h.failTask(taskID, http.StatusBadGateway, imageTaskErrorPayload("api_error", "upstream returned an invalid image response")) + return + } + if err := h.tasks.Complete(context.Background(), taskID, statusCode, json.RawMessage(body)); err != nil { + logger.L().Error("image_task.complete_store_failed", zap.String("task_id", taskID), zap.Error(err)) + } + return + } + h.failTask(taskID, statusCode, extractImageTaskError(body)) +} + +func (h *AsyncImageHandler) failTask(taskID string, statusCode int, taskErr json.RawMessage) { + if err := h.tasks.Fail(context.Background(), taskID, statusCode, taskErr); err != nil { + logger.L().Error("image_task.failure_store_failed", zap.String("task_id", taskID), zap.Error(err)) + } +} + +func newAsyncImageContext(c *gin.Context, body []byte, timeoutDuration time.Duration) (*gin.Context, *httptest.ResponseRecorder, context.CancelFunc) { + base := context.WithoutCancel(c.Request.Context()) + executionCtx, cancel := context.WithTimeout(base, timeoutDuration) + request := c.Request.Clone(executionCtx) + request.Body = io.NopCloser(bytes.NewReader(body)) + request.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(bytes.NewReader(body)), nil + } + request.ContentLength = int64(len(body)) + request.URL.Path = strings.TrimSuffix(request.URL.Path, "/async") + + taskCtx := c.Copy() + recorder := httptest.NewRecorder() + recorderCtx, _ := gin.CreateTestContext(recorder) + taskCtx.Writer = recorderCtx.Writer + taskCtx.Request = request + return taskCtx, recorder, cancel +} + +func asyncImageRequestStreams(contentType string, body []byte) bool { + if isMultipartImagesContentType(contentType) { + return false + } + var envelope struct { + Stream bool `json:"stream"` + } + return json.Unmarshal(body, &envelope) == nil && envelope.Stream +} + +func imageTaskPollURL(submitPath, taskID string) string { + if strings.HasPrefix(submitPath, "/v1/") { + return "/v1/images/tasks/" + taskID + } + return "/images/tasks/" + taskID +} + +func extractImageTaskError(body []byte) json.RawMessage { + if json.Valid(body) { + var envelope struct { + Error json.RawMessage `json:"error"` + } + if json.Unmarshal(body, &envelope) == nil && len(envelope.Error) > 0 && json.Valid(envelope.Error) { + return envelope.Error + } + return json.RawMessage(body) + } + return imageTaskErrorPayload("api_error", "image generation failed") +} + +func imageTaskErrorPayload(errorType, message string) json.RawMessage { + data, _ := json.Marshal(gin.H{"type": errorType, "message": message}) + return data +} + +func imageTaskError(c *gin.Context, err error) { + status := infraerrors.Code(err) + code := infraerrors.Reason(err) + message := infraerrors.Message(err) + if status <= 0 { + status = http.StatusInternalServerError + } + if strings.TrimSpace(code) == "" { + code = "IMAGE_TASK_ERROR" + } + imageTaskJSONError(c, status, code, message) +} + +func imageTaskJSONError(c *gin.Context, status int, code, message string) { + c.Header("Cache-Control", "no-store") + c.JSON(status, gin.H{"error": gin.H{"type": code, "code": code, "message": message}}) +} diff --git a/backend/internal/handler/image_task_handler_test.go b/backend/internal/handler/image_task_handler_test.go new file mode 100644 index 0000000000..84002b7ce9 --- /dev/null +++ b/backend/internal/handler/image_task_handler_test.go @@ -0,0 +1,107 @@ +package handler + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + 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 asyncImageMemoryStore struct { + mu sync.RWMutex + tasks map[string]*service.ImageTaskRecord +} + +func (s *asyncImageMemoryStore) Save(_ context.Context, task *service.ImageTaskRecord, _ time.Duration) error { + s.mu.Lock() + defer s.mu.Unlock() + copy := *task + copy.Result = append(json.RawMessage(nil), task.Result...) + copy.Error = append(json.RawMessage(nil), task.Error...) + s.tasks[task.ID] = © + return nil +} + +func (s *asyncImageMemoryStore) Get(_ context.Context, id string) (*service.ImageTaskRecord, error) { + s.mu.RLock() + defer s.mu.RUnlock() + task := s.tasks[id] + if task == nil { + return nil, service.ErrImageTaskNotFound + } + copy := *task + copy.Result = append(json.RawMessage(nil), task.Result...) + copy.Error = append(json.RawMessage(nil), task.Error...) + return ©, nil +} + +func TestAsyncImageHandlerSubmitAndPoll(t *testing.T) { + gin.SetMode(gin.TestMode) + store := &asyncImageMemoryStore{tasks: make(map[string]*service.ImageTaskRecord)} + tasks := service.NewImageTaskServiceWithOptions(store, time.Hour, time.Minute) + release := make(chan struct{}) + h := &AsyncImageHandler{tasks: tasks} + h.execute = func(_ string, c *gin.Context) { + <-release + c.JSON(http.StatusOK, gin.H{"created": 123, "data": []gin.H{{"url": "https://example.test/image.png"}}}) + } + + router := gin.New() + router.Use(func(c *gin.Context) { + groupID := int64(3) + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ + ID: 9, + UserID: 7, + GroupID: &groupID, + Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, AllowImageGeneration: true}, + }) + c.Next() + }) + router.POST("/v1/images/generations/async", h.Submit) + router.GET("/v1/images/tasks/:task_id", h.Get) + + requestCtx, cancelRequest := context.WithCancel(context.Background()) + req := httptest.NewRequest(http.MethodPost, "/v1/images/generations/async", strings.NewReader(`{"model":"gpt-image-1","prompt":"cat"}`)).WithContext(requestCtx) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + require.Equal(t, http.StatusAccepted, w.Code) + require.Equal(t, "no-store", w.Header().Get("Cache-Control")) + require.Equal(t, "3", w.Header().Get("Retry-After")) + + var accepted struct { + TaskID string `json:"task_id"` + Status string `json:"status"` + PollURL string `json:"poll_url"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &accepted)) + require.Equal(t, service.ImageTaskStatusProcessing, accepted.Status) + require.Equal(t, "/v1/images/tasks/"+accepted.TaskID, accepted.PollURL) + require.Equal(t, accepted.PollURL, w.Header().Get("Location")) + + // The detached background request must survive completion of/cancellation + // from the short submission request. + cancelRequest() + close(release) + require.Eventually(t, func() bool { + got, err := tasks.Get(context.Background(), service.ImageTaskOwner{UserID: 7, APIKeyID: 9}, accepted.TaskID) + return err == nil && got.Status == service.ImageTaskStatusCompleted + }, time.Second, 10*time.Millisecond) + + pollReq := httptest.NewRequest(http.MethodGet, accepted.PollURL, nil) + pollWriter := httptest.NewRecorder() + router.ServeHTTP(pollWriter, pollReq) + require.Equal(t, http.StatusOK, pollWriter.Code) + require.Equal(t, "no-store", pollWriter.Header().Get("Cache-Control")) + require.Empty(t, pollWriter.Header().Get("Retry-After")) + require.Contains(t, pollWriter.Body.String(), "https://example.test/image.png") +} diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index 67380fed33..4831aaf70e 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -115,6 +115,7 @@ func ProvideHandlers( paymentHandler *PaymentHandler, paymentWebhookHandler *PaymentWebhookHandler, availableChannelHandler *AvailableChannelHandler, + asyncImageHandler *AsyncImageHandler, batchImageHandler *BatchImageHandler, _ *service.IdempotencyCoordinator, _ *service.IdempotencyCleanupService, @@ -136,6 +137,7 @@ func ProvideHandlers( Payment: paymentHandler, PaymentWebhook: paymentWebhookHandler, AvailableChannel: availableChannelHandler, + AsyncImage: asyncImageHandler, BatchImage: batchImageHandler, } } @@ -158,6 +160,7 @@ var ProviderSet = wire.NewSet( NewPaymentHandler, NewPaymentWebhookHandler, NewAvailableChannelHandler, + NewAsyncImageHandler, NewBatchImageHandler, // Admin handlers diff --git a/backend/internal/repository/image_task_store.go b/backend/internal/repository/image_task_store.go new file mode 100644 index 0000000000..000346154a --- /dev/null +++ b/backend/internal/repository/image_task_store.go @@ -0,0 +1,48 @@ +package repository + +import ( + "context" + "encoding/json" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/redis/go-redis/v9" +) + +const imageTaskKeyPrefix = "image_task:" + +type imageTaskStore struct { + rdb *redis.Client +} + +func NewImageTaskStore(rdb *redis.Client) service.ImageTaskStore { + return &imageTaskStore{rdb: rdb} +} + +func (s *imageTaskStore) Save(ctx context.Context, task *service.ImageTaskRecord, ttl time.Duration) error { + data, err := json.Marshal(task) + if err != nil { + return err + } + return s.rdb.Set(ctx, imageTaskKey(task.ID), data, ttl).Err() +} + +func (s *imageTaskStore) Get(ctx context.Context, id string) (*service.ImageTaskRecord, error) { + data, err := s.rdb.Get(ctx, imageTaskKey(id)).Bytes() + if err != nil { + if err == redis.Nil { + return nil, service.ErrImageTaskNotFound + } + return nil, err + } + var task service.ImageTaskRecord + if err := json.Unmarshal(data, &task); err != nil { + return nil, err + } + return &task, nil +} + +func imageTaskKey(id string) string { + return imageTaskKeyPrefix + strings.TrimSpace(id) +} diff --git a/backend/internal/repository/image_task_store_test.go b/backend/internal/repository/image_task_store_test.go new file mode 100644 index 0000000000..dc4ace5a63 --- /dev/null +++ b/backend/internal/repository/image_task_store_test.go @@ -0,0 +1,43 @@ +package repository + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +func TestImageTaskStoreRoundTripAndTTL(t *testing.T) { + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + store := NewImageTaskStore(rdb) + task := &service.ImageTaskRecord{ + ID: "imgtask_123", + UserID: 7, + APIKeyID: 9, + Status: service.ImageTaskStatusProcessing, + CreatedAt: 100, + ExpiresAt: 200, + } + + require.NoError(t, store.Save(context.Background(), task, 24*time.Hour)) + got, err := store.Get(context.Background(), task.ID) + require.NoError(t, err) + require.Equal(t, task, got) + require.Equal(t, 24*time.Hour, mr.TTL(imageTaskKey(task.ID))) +} + +func TestImageTaskStoreMissing(t *testing.T) { + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + store := NewImageTaskStore(rdb) + + _, err := store.Get(context.Background(), "imgtask_missing") + require.ErrorIs(t, err, service.ErrImageTaskNotFound) +} diff --git a/backend/internal/repository/wire.go b/backend/internal/repository/wire.go index 0ee91d146b..d7fb702b39 100644 --- a/backend/internal/repository/wire.go +++ b/backend/internal/repository/wire.go @@ -117,6 +117,7 @@ var ProviderSet = wire.NewSet( NewRedeemCache, NewUpdateCache, NewGeminiTokenCache, + NewImageTaskStore, NewBatchImageQueue, NewBatchImageDownloadLimiter, NewLeaderLockCache, diff --git a/backend/internal/server/middleware/api_key_auth.go b/backend/internal/server/middleware/api_key_auth.go index 082766bfd5..a512d8bcd7 100644 --- a/backend/internal/server/middleware/api_key_auth.go +++ b/backend/internal/server/middleware/api_key_auth.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "net/http" "strings" "github.com/Wei-Shaw/sub2api/internal/config" @@ -25,7 +26,7 @@ func NewAPIKeyAuthMiddleware(apiKeyService *service.APIKeyService, subscriptionS // - 鉴权（Authentication）：验证 Key 有效性、用户状态、IP 限制 —— 始终执行 // - 计费执行（Billing Enforcement）：过期/配额/订阅/余额检查 —— skipBilling 时整块跳过 // -// /v1/usage 端点只需鉴权，不需要计费执行（允许过期/配额耗尽的 Key 查询自身用量）。 +// /v1/usage 与异步生图任务查询只需鉴权，不需要计费执行（允许已耗尽额度的 Key 查询自身结果）。 func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc { return func(c *gin.Context) { // ── 1. 提取 API Key ────────────────────────────────────────── @@ -146,8 +147,11 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti // ── 5. 加载订阅（订阅模式时始终加载） ─────────────────────── - // skipBilling: /v1/usage 只需鉴权，跳过所有计费执行 - skipBilling := c.Request.URL.Path == "/v1/usage" + // skipBilling: usage and async image task polling only read data that + // already belongs to the authenticated key. In particular, polling must + // remain available after the completed generation consumes the key's + // remaining balance. + skipBilling := c.Request.URL.Path == "/v1/usage" || isAsyncImageTaskRead(c.Request.Method, c.Request.URL.Path) var subscription *service.UserSubscription isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType() @@ -243,6 +247,13 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti } } +func isAsyncImageTaskRead(method, path string) bool { + if method != http.MethodGet { + return false + } + return strings.HasPrefix(path, "/v1/images/tasks/") || strings.HasPrefix(path, "/images/tasks/") +} + // GetAPIKeyFromContext 从上下文中获取API key func GetAPIKeyFromContext(c *gin.Context) (*service.APIKey, bool) { value, exists := c.Get(string(ContextKeyAPIKey)) diff --git a/backend/internal/server/middleware/api_key_auth_image_task_test.go b/backend/internal/server/middleware/api_key_auth_image_task_test.go new file mode 100644 index 0000000000..653fd8ba3c --- /dev/null +++ b/backend/internal/server/middleware/api_key_auth_image_task_test.go @@ -0,0 +1,15 @@ +package middleware + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestIsAsyncImageTaskRead(t *testing.T) { + require.True(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/tasks/imgtask_123")) + require.True(t, isAsyncImageTaskRead(http.MethodGet, "/images/tasks/imgtask_123")) + require.False(t, isAsyncImageTaskRead(http.MethodPost, "/v1/images/tasks/imgtask_123")) + require.False(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/generations")) +} diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index f73379981e..e79cd4a56e 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -191,6 +191,9 @@ func RegisterGatewayRoutes( }) gateway.POST("/images/generations", imagesHandler) gateway.POST("/images/edits", imagesHandler) + gateway.POST("/images/generations/async", h.AsyncImage.Submit) + gateway.POST("/images/edits/async", h.AsyncImage.Submit) + gateway.GET("/images/tasks/:task_id", h.AsyncImage.Get) gateway.POST("/images/batches", h.BatchImage.Submit) gateway.GET("/images/batches", h.BatchImage.List) gateway.GET("/images/batches/models", h.BatchImage.Models) @@ -271,6 +274,9 @@ func RegisterGatewayRoutes( }) r.POST("/images/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler) r.POST("/images/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler) + r.POST("/images/generations/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.AsyncImage.Submit) + r.POST("/images/edits/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.AsyncImage.Submit) + r.GET("/images/tasks/:task_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.AsyncImage.Get) r.POST("/videos/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoGenerationHandler) r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoEditHandler) r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoExtensionHandler) diff --git a/backend/internal/server/routes/gateway_test.go b/backend/internal/server/routes/gateway_test.go index 65c6824440..f8abdb4062 100644 --- a/backend/internal/server/routes/gateway_test.go +++ b/backend/internal/server/routes/gateway_test.go @@ -28,6 +28,7 @@ func newGatewayRoutesTestRouter(platform ...string) *gin.Engine { &handler.Handlers{ Gateway: &handler.GatewayHandler{}, OpenAIGateway: &handler.OpenAIGatewayHandler{}, + AsyncImage: handler.NewAsyncImageHandler(nil, nil), }, servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) { groupID := int64(1) @@ -113,6 +114,25 @@ func TestGatewayRoutesOpenAIImagesPathsAreRegistered(t *testing.T) { } } +func TestGatewayRoutesAsyncImagesPathsAreRegistered(t *testing.T) { + router := newGatewayRoutesTestRouter() + registered := make(map[string]bool) + for _, route := range router.Routes() { + registered[route.Method+" "+route.Path] = true + } + + for _, route := range []string{ + "POST /v1/images/generations/async", + "POST /v1/images/edits/async", + "GET /v1/images/tasks/:task_id", + "POST /images/generations/async", + "POST /images/edits/async", + "GET /images/tasks/:task_id", + } { + require.True(t, registered[route], "%s should be registered", route) + } +} + func TestGatewayRoutesGrokImagesAndVideosPathsAreRegistered(t *testing.T) { router := newGatewayRoutesTestRouter(service.PlatformGrok) diff --git a/backend/internal/service/image_task.go b/backend/internal/service/image_task.go new file mode 100644 index 0000000000..e41b73dd6f --- /dev/null +++ b/backend/internal/service/image_task.go @@ -0,0 +1,210 @@ +package service + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "strings" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/google/uuid" +) + +const ( + ImageTaskStatusProcessing = "processing" + ImageTaskStatusCompleted = "completed" + ImageTaskStatusFailed = "failed" + + defaultImageTaskTTL = 24 * time.Hour + defaultImageTaskExecutionTimeout = 30 * time.Minute +) + +var ( + ErrImageTaskNotFound = infraerrors.New(http.StatusNotFound, "IMAGE_TASK_NOT_FOUND", "image task not found") + ErrImageTaskForbidden = infraerrors.New(http.StatusForbidden, "IMAGE_TASK_FORBIDDEN", "image task does not belong to this API key") + ErrImageTaskUnavailable = infraerrors.New(http.StatusServiceUnavailable, "IMAGE_TASK_UNAVAILABLE", "image task storage is unavailable") +) + +// ImageTaskRecord is the private Redis representation of an asynchronous image +// request. Ownership fields are intentionally omitted from the public view. +type ImageTaskRecord struct { + ID string `json:"id"` + UserID int64 `json:"user_id"` + APIKeyID int64 `json:"api_key_id"` + Status string `json:"status"` + HTTPStatus int `json:"http_status,omitempty"` + Result json.RawMessage `json:"result,omitempty"` + Error json.RawMessage `json:"error,omitempty"` + CreatedAt int64 `json:"created_at"` + CompletedAt *int64 `json:"completed_at,omitempty"` + ExpiresAt int64 `json:"expires_at"` +} + +// ImageTask is the API-safe task representation returned to callers. +type ImageTask struct { + ID string `json:"id"` + TaskID string `json:"task_id"` + Object string `json:"object"` + Status string `json:"status"` + HTTPStatus int `json:"http_status,omitempty"` + ImageURL string `json:"image_url,omitempty"` + Result json.RawMessage `json:"result,omitempty"` + Error json.RawMessage `json:"error,omitempty"` + CreatedAt int64 `json:"created_at"` + CompletedAt *int64 `json:"completed_at,omitempty"` + ExpiresAt int64 `json:"expires_at"` +} + +type ImageTaskOwner struct { + UserID int64 + APIKeyID int64 +} + +type ImageTaskStore interface { + Save(ctx context.Context, task *ImageTaskRecord, ttl time.Duration) error + Get(ctx context.Context, id string) (*ImageTaskRecord, error) +} + +type ImageTaskService struct { + store ImageTaskStore + ttl time.Duration + executionTimeout time.Duration +} + +func NewImageTaskService(store ImageTaskStore) *ImageTaskService { + return NewImageTaskServiceWithOptions(store, defaultImageTaskTTL, defaultImageTaskExecutionTimeout) +} + +func NewImageTaskServiceWithOptions(store ImageTaskStore, ttl, executionTimeout time.Duration) *ImageTaskService { + if ttl <= 0 { + ttl = defaultImageTaskTTL + } + if executionTimeout <= 0 { + executionTimeout = defaultImageTaskExecutionTimeout + } + return &ImageTaskService{store: store, ttl: ttl, executionTimeout: executionTimeout} +} + +func (s *ImageTaskService) ExecutionTimeout() time.Duration { + if s == nil || s.executionTimeout <= 0 { + return defaultImageTaskExecutionTimeout + } + return s.executionTimeout +} + +func (s *ImageTaskService) Create(ctx context.Context, owner ImageTaskOwner) (*ImageTask, error) { + if s == nil || s.store == nil { + return nil, ErrImageTaskUnavailable + } + now := time.Now().UTC() + task := &ImageTaskRecord{ + ID: "imgtask_" + strings.ReplaceAll(uuid.NewString(), "-", ""), + UserID: owner.UserID, + APIKeyID: owner.APIKeyID, + Status: ImageTaskStatusProcessing, + CreatedAt: now.Unix(), + ExpiresAt: now.Add(s.ttl).Unix(), + } + if err := s.store.Save(ctx, task, s.ttl); err != nil { + return nil, ErrImageTaskUnavailable.WithCause(err) + } + return imageTaskToPublic(task), nil +} + +func (s *ImageTaskService) Get(ctx context.Context, owner ImageTaskOwner, id string) (*ImageTask, error) { + if s == nil || s.store == nil { + return nil, ErrImageTaskUnavailable + } + task, err := s.store.Get(ctx, strings.TrimSpace(id)) + if err != nil { + if errors.Is(err, ErrImageTaskNotFound) { + return nil, ErrImageTaskNotFound + } + return nil, ErrImageTaskUnavailable.WithCause(err) + } + if task.UserID != owner.UserID || task.APIKeyID != owner.APIKeyID { + // Do not reveal whether a random task ID exists for another caller. + return nil, ErrImageTaskNotFound + } + return imageTaskToPublic(task), nil +} + +func (s *ImageTaskService) Complete(ctx context.Context, id string, statusCode int, result json.RawMessage) error { + if !json.Valid(result) { + return s.Fail(ctx, id, http.StatusBadGateway, imageTaskErrorJSON("api_error", "upstream returned a non-JSON image response")) + } + return s.finish(ctx, id, ImageTaskStatusCompleted, statusCode, result, nil) +} + +func (s *ImageTaskService) Fail(ctx context.Context, id string, statusCode int, taskErr json.RawMessage) error { + if !json.Valid(taskErr) { + taskErr = imageTaskErrorJSON("api_error", "image generation failed") + } + return s.finish(ctx, id, ImageTaskStatusFailed, statusCode, nil, taskErr) +} + +func (s *ImageTaskService) finish(ctx context.Context, id, status string, statusCode int, result, taskErr json.RawMessage) error { + if s == nil || s.store == nil { + return ErrImageTaskUnavailable + } + task, err := s.store.Get(ctx, id) + if err != nil { + if errors.Is(err, ErrImageTaskNotFound) { + return ErrImageTaskNotFound + } + return ErrImageTaskUnavailable.WithCause(err) + } + now := time.Now().UTC() + completedAt := now.Unix() + task.Status = status + task.HTTPStatus = statusCode + task.Result = result + task.Error = taskErr + task.CompletedAt = &completedAt + task.ExpiresAt = now.Add(s.ttl).Unix() + if err := s.store.Save(ctx, task, s.ttl); err != nil { + return ErrImageTaskUnavailable.WithCause(err) + } + return nil +} + +func imageTaskToPublic(task *ImageTaskRecord) *ImageTask { + if task == nil { + return nil + } + return &ImageTask{ + ID: task.ID, + TaskID: task.ID, + Object: "image.generation.task", + Status: task.Status, + HTTPStatus: task.HTTPStatus, + ImageURL: firstImageTaskURL(task.Result), + Result: task.Result, + Error: task.Error, + CreatedAt: task.CreatedAt, + CompletedAt: task.CompletedAt, + ExpiresAt: task.ExpiresAt, + } +} + +func firstImageTaskURL(result json.RawMessage) string { + if len(result) == 0 || !json.Valid(result) { + return "" + } + var response struct { + Data []struct { + URL string `json:"url"` + } `json:"data"` + } + if json.Unmarshal(result, &response) != nil || len(response.Data) == 0 { + return "" + } + return strings.TrimSpace(response.Data[0].URL) +} + +func imageTaskErrorJSON(errorType, message string) json.RawMessage { + data, _ := json.Marshal(map[string]string{"type": errorType, "message": message}) + return data +} diff --git a/backend/internal/service/image_task_test.go b/backend/internal/service/image_task_test.go new file mode 100644 index 0000000000..6eb9ce50c9 --- /dev/null +++ b/backend/internal/service/image_task_test.go @@ -0,0 +1,91 @@ +package service + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type imageTaskMemoryStore struct { + task *ImageTaskRecord + ttl time.Duration + saveErr error + getErr error +} + +func (s *imageTaskMemoryStore) Save(_ context.Context, task *ImageTaskRecord, ttl time.Duration) error { + if s.saveErr != nil { + return s.saveErr + } + copy := *task + s.task = © + s.ttl = ttl + return nil +} + +func (s *imageTaskMemoryStore) Get(_ context.Context, _ string) (*ImageTaskRecord, error) { + if s.getErr != nil { + return nil, s.getErr + } + if s.task == nil { + return nil, ErrImageTaskNotFound + } + copy := *s.task + return ©, nil +} + +func TestImageTaskServiceLifecycleAndOwnership(t *testing.T) { + store := &imageTaskMemoryStore{} + svc := NewImageTaskServiceWithOptions(store, time.Hour, 10*time.Minute) + owner := ImageTaskOwner{UserID: 7, APIKeyID: 9} + + created, err := svc.Create(context.Background(), owner) + require.NoError(t, err) + require.Equal(t, ImageTaskStatusProcessing, created.Status) + require.Equal(t, created.ID, created.TaskID) + require.Equal(t, "image.generation.task", created.Object) + require.Equal(t, time.Hour, store.ttl) + require.Equal(t, owner.UserID, store.task.UserID) + require.Equal(t, owner.APIKeyID, store.task.APIKeyID) + + _, err = svc.Get(context.Background(), ImageTaskOwner{UserID: 7, APIKeyID: 10}, created.ID) + require.ErrorIs(t, err, ErrImageTaskNotFound) + + result := json.RawMessage(`{"created":123,"data":[{"url":"https://example.test/image.png"}]}`) + require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, result)) + + completed, err := svc.Get(context.Background(), owner, created.ID) + require.NoError(t, err) + require.Equal(t, ImageTaskStatusCompleted, completed.Status) + require.Equal(t, http.StatusOK, completed.HTTPStatus) + require.Equal(t, "https://example.test/image.png", completed.ImageURL) + require.JSONEq(t, string(result), string(completed.Result)) + require.NotNil(t, completed.CompletedAt) +} + +func TestImageTaskServiceInvalidResultBecomesFailed(t *testing.T) { + store := &imageTaskMemoryStore{} + svc := NewImageTaskServiceWithOptions(store, time.Hour, time.Minute) + created, err := svc.Create(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2}) + require.NoError(t, err) + + require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, json.RawMessage(`not-json`))) + got, err := svc.Get(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2}, created.ID) + require.NoError(t, err) + require.Equal(t, ImageTaskStatusFailed, got.Status) + require.Equal(t, http.StatusBadGateway, got.HTTPStatus) + require.Contains(t, string(got.Error), "non-JSON") +} + +func TestImageTaskServiceMapsStoreFailures(t *testing.T) { + store := &imageTaskMemoryStore{saveErr: errors.New("redis down")} + svc := NewImageTaskService(store) + + _, err := svc.Create(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2}) + require.ErrorIs(t, err, ErrImageTaskUnavailable) +} diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index ace69e0c34..6dd2c29abc 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -637,6 +637,7 @@ var ProviderSet = wire.NewSet( NewAdminService, NewGatewayService, NewOpenAIGatewayService, + NewImageTaskService, ProvideBatchImageModelPricingResolver, NewBatchImagePublicService, NewBatchImageDownloadService, @@ -646,6 +647,7 @@ var ProviderSet = wire.NewSet( NewOAuthService, ProvideOpenAIOAuthService, NewGrokOAuthService, + wire.Bind(new(GrokOAuthTokenService), new(*GrokOAuthService)), NewGeminiOAuthService, NewGeminiQuotaService, NewCompositeTokenCacheInvalidator, From df247b4364035b6399982785b39282b17fa4029d Mon Sep 17 00:00:00 2001 From: haruka <1628615876@qq.com> Date: Wed, 15 Jul 2026 22:14:34 +0800 Subject: [PATCH 04/26] docs: document async image task API --- .gitignore | 1 + README.md | 6 +++ docs/ASYNC_IMAGE_TASKS.md | 111 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 118 insertions(+) create mode 100644 docs/ASYNC_IMAGE_TASKS.md diff --git a/.gitignore b/.gitignore index d45d27f79e..8c7e67f7e2 100644 --- a/.gitignore +++ b/.gitignore @@ -135,6 +135,7 @@ docs/* !docs/PAYMENT.md !docs/PAYMENT_CN.md !docs/ADMIN_PAYMENT_INTEGRATION_API.md +!docs/ASYNC_IMAGE_TASKS.md !docs/legal/ !docs/legal/*.md .serena/ diff --git a/README.md b/README.md index b06839745f..0d1fe23aee 100644 --- a/README.md +++ b/README.md @@ -688,6 +688,12 @@ Simple Mode is designed for individual developers or internal teams who want qui --- +## Asynchronous Image Tasks + +Long-running OpenAI/Grok image generation and editing can be submitted through `/v1/images/generations/async` or `/v1/images/edits/async`, then polled at `/v1/images/tasks/{task_id}` without holding a CDN connection open. See [Asynchronous Image Tasks](docs/ASYNC_IMAGE_TASKS.md) for request and response examples. + +--- + ## Grok / xAI Support Sub2API supports both Grok subscription accounts through xAI OAuth and standard xAI API-key accounts. Both account types forward OpenAI-compatible Responses traffic to xAI. diff --git a/docs/ASYNC_IMAGE_TASKS.md b/docs/ASYNC_IMAGE_TASKS.md new file mode 100644 index 0000000000..69db7b4932 --- /dev/null +++ b/docs/ASYNC_IMAGE_TASKS.md @@ -0,0 +1,111 @@ +# Asynchronous Image Tasks + +Asynchronous image tasks let clients submit long-running OpenAI-compatible image requests without keeping one HTTP connection open. This avoids proxy/CDN response timeouts such as Cloudflare 524 while preserving the existing image routing, billing, moderation, concurrency, and failover behavior. + +## Endpoints + +The authenticated gateway exposes both `/v1` paths and their existing no-prefix aliases: + +```text +POST /v1/images/generations/async +POST /v1/images/edits/async +GET /v1/images/tasks/{task_id} +``` + +The aliases are `/images/generations/async`, `/images/edits/async`, and `/images/tasks/{task_id}`. + +Only OpenAI and Grok groups are supported. Requests use the same JSON or multipart payload as the corresponding synchronous endpoint. Streaming image requests are rejected because a polled task returns one final JSON result. + +## Submit a task + +```bash +curl -i https://api.example.com/v1/images/generations/async \ + -H 'Authorization: Bearer sk-...' \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "gpt-image-1", + "prompt": "A lighthouse during a winter storm", + "size": "1536x1024" + }' +``` + +The server stores the initial task in Redis and responds with `202 Accepted`: + +```json +{ + "id": "imgtask_0123456789abcdef", + "task_id": "imgtask_0123456789abcdef", + "object": "image.generation.task", + "status": "processing", + "created_at": 1784092800, + "expires_at": 1784179200, + "poll_url": "/v1/images/tasks/imgtask_0123456789abcdef" +} +``` + +`Location` contains the polling path and `Retry-After: 3` provides the recommended polling interval. + +## Poll a task + +Use the same API key that submitted the task: + +```bash +curl https://api.example.com/v1/images/tasks/imgtask_0123456789abcdef \ + -H 'Authorization: Bearer sk-...' +``` + +While work is in progress: + +```json +{ + "id": "imgtask_0123456789abcdef", + "task_id": "imgtask_0123456789abcdef", + "object": "image.generation.task", + "status": "processing", + "created_at": 1784092800, + "expires_at": 1784179200 +} +``` + +On success, `result` is the unmodified JSON body from the synchronous image API, so URL and base64 response formats both remain supported: + +```json +{ + "id": "imgtask_0123456789abcdef", + "task_id": "imgtask_0123456789abcdef", + "object": "image.generation.task", + "status": "completed", + "http_status": 200, + "image_url": "https://...", + "result": { + "created": 1784092923, + "data": [{"url": "https://..."}] + }, + "created_at": 1784092800, + "completed_at": 1784092923, + "expires_at": 1784179323 +} +``` + +For URL responses, `image_url` mirrors the first `data[].url` for simple clients. On failure, the task reaches `failed` and exposes the original OpenAI-compatible error object where available: + +```json +{ + "id": "imgtask_0123456789abcdef", + "task_id": "imgtask_0123456789abcdef", + "object": "image.generation.task", + "status": "failed", + "http_status": 502, + "error": { + "type": "api_error", + "message": "Upstream request failed" + }, + "created_at": 1784092800, + "completed_at": 1784092923, + "expires_at": 1784179323 +} +``` + +All submit and poll responses include `Cache-Control: no-store`, preventing a CDN from caching the `processing` state. Tasks and results expire 24 hours after their latest state update. A task executes for at most 30 minutes. + +Task ownership is scoped to both user and API key. Unknown task IDs and IDs owned by another key both return `404`, avoiding task-existence disclosure. Polling remains available when the completed generation used the key's remaining balance; normal authentication, disabled-key, user, IP, and group checks still apply. From 4f641208a0a08e2b38df98ec226c6c1570077eee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=8F?= <2829624376@qq.com> Date: Wed, 15 Jul 2026 22:05:53 +0800 Subject: [PATCH 05/26] fix(openai-ws): close ingress reads cleanly --- .../handler/openai_gateway_handler.go | 54 ++++--- .../handler/openai_gateway_handler_test.go | 64 ++++++-- .../internal/service/openai_ws_client_read.go | 70 +++++++++ .../service/openai_ws_client_read_test.go | 143 ++++++++++++++++++ .../service/openai_ws_forwarder_ingress.go | 22 ++- ...penai_ws_forwarder_ingress_session_test.go | 8 + .../service/openai_ws_http_bridge_test.go | 8 + 7 files changed, 322 insertions(+), 47 deletions(-) create mode 100644 backend/internal/service/openai_ws_client_read.go create mode 100644 backend/internal/service/openai_ws_client_read_test.go diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 83c4b4ac23..af4746531b 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -1360,6 +1360,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, @@ -1381,32 +1403,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)) diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index d855f364a7..b1eb260fe3 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -820,17 +820,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 +921,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/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..0232e7a72d 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -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 } 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..d2fc0ebae5 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -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 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 From 2fe7df9b813a949bfbd9ef5bbb975773fd63049c Mon Sep 17 00:00:00 2001 From: wucm667 Date: Wed, 15 Jul 2026 22:22:51 +0800 Subject: [PATCH 06/26] fix(openai): reject image models on chat completions --- .../chat_completions_image_model_test.go | 96 +++++++++++++++++++ .../gateway_handler_chat_completions.go | 4 + .../handler/openai_chat_completions.go | 4 + backend/internal/service/openai_images.go | 7 +- 4 files changed, 110 insertions(+), 1 deletion(-) create mode 100644 backend/internal/handler/chat_completions_image_model_test.go 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/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/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index 62a2b70b59..37f532bb21 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)) diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index 081b5e741b..7a5d456d20 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 { From 56650d6aed98d9c307892f54548c03a3bdf0abe5 Mon Sep 17 00:00:00 2001 From: wucm667 Date: Wed, 15 Jul 2026 22:23:57 +0800 Subject: [PATCH 07/26] fix(openai): normalize Responses Lite reasoning context --- .../service/openai_responses_lite_tools.go | 39 +++++++++++++--- .../openai_responses_lite_tools_test.go | 44 +++++++++++++++++++ ...penai_ws_forwarder_ingress_session_test.go | 7 +++ 3 files changed, 83 insertions(+), 7 deletions(-) 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_ws_forwarder_ingress_session_test.go b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go index 3752c1bd0e..2578c623da 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -531,6 +531,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 +584,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 +595,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 +980,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 +1014,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) { From 5e4da92dec788b31e01db30dbfc340c6eaf18dbb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=8F?= <2829624376@qq.com> Date: Wed, 15 Jul 2026 22:12:00 +0800 Subject: [PATCH 08/26] fix(openai): retry explicitly rejected Responses fields --- .../service/openai_gateway_forward.go | 16 +- .../openai_responses_rejected_field_retry.go | 133 ++++++++++ ...nai_responses_rejected_field_retry_test.go | 230 ++++++++++++++++++ 3 files changed, 376 insertions(+), 3 deletions(-) create mode 100644 backend/internal/service/openai_responses_rejected_field_retry.go create mode 100644 backend/internal/service/openai_responses_rejected_field_retry_test.go diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 5259aa843c..c6adda595a 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 { 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)), + } +} From f59a6ed74c914b2980ec98999e346e8deea04104 Mon Sep 17 00:00:00 2001 From: Tian Lee <498756723@qq.com> Date: Wed, 15 Jul 2026 22:42:53 +0800 Subject: [PATCH 09/26] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=20API=20Key=20?= =?UTF-8?q?=E8=AE=A1=E8=B4=B9=E5=80=8D=E7=8E=87=E8=87=AA=E7=9C=81=E6=8E=A5?= =?UTF-8?q?=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/cmd/server/wire_gen.go | 2 +- backend/internal/handler/gateway_handler.go | 3 + .../internal/handler/gateway_key_billing.go | 108 ++++++ .../handler/gateway_key_billing_test.go | 311 ++++++++++++++++++ .../server/middleware/api_key_auth.go | 21 +- .../server/middleware/api_key_auth_test.go | 130 +++++++- backend/internal/server/routes/gateway.go | 1 + .../server/routes/gateway_key_billing_test.go | 172 ++++++++++ .../internal/service/gateway_usage_billing.go | 7 +- .../internal/service/openai_gateway_usage.go | 18 +- 10 files changed, 756 insertions(+), 17 deletions(-) create mode 100644 backend/internal/handler/gateway_key_billing.go create mode 100644 backend/internal/handler/gateway_key_billing_test.go create mode 100644 backend/internal/server/routes/gateway_key_billing_test.go diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 86721fbac0..3700e2257e 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -254,7 +254,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/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_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/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/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/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index 410b4d944d..677239c3c6 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 倍率缓存。 From 40b8f04a6af732e5ab0f71fe685a5c8429789060 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=8F?= <2829624376@qq.com> Date: Wed, 15 Jul 2026 21:56:13 +0800 Subject: [PATCH 10/26] fix(openai): scope transient cooldowns by model --- backend/internal/handler/grok_media.go | 6 +- .../internal/handler/openai_alpha_search.go | 6 +- .../handler/openai_chat_completions.go | 8 +- backend/internal/handler/openai_embeddings.go | 6 +- .../handler/openai_gateway_handler.go | 26 +-- .../handler/openai_gateway_handler_test.go | 13 ++ backend/internal/handler/openai_images.go | 10 +- .../service/openai_account_model_transient.go | 177 ++++++++++++++++++ .../openai_account_model_transient_test.go | 130 +++++++++++++ .../openai_account_runtime_block_fastpath.go | 97 +++++++++- .../openai_account_runtime_transient_test.go | 116 ++++++++++++ .../service/openai_account_scheduler.go | 9 +- .../service/openai_account_scheduler_test.go | 30 ++- .../service/openai_gateway_passthrough.go | 6 +- .../service/openai_gateway_scheduling.go | 10 +- .../service/openai_gateway_service.go | 53 ++++-- .../service/openai_gateway_upstream_errors.go | 7 +- .../service/openai_ws_forwarder_ingress.go | 29 +-- ...penai_ws_forwarder_ingress_session_test.go | 8 +- .../openai_ws_forwarder_success_test.go | 51 +++++ .../service/openai_ws_forwarder_support.go | 94 +++++++++- .../service/openai_ws_forwarder_test.go | 79 ++++++++ .../service/openai_ws_forwarder_v2.go | 31 +-- .../internal/service/openai_ws_http_bridge.go | 29 +-- .../openai_ws_ratelimit_signal_test.go | 46 +++++ .../openai_ws_v2_passthrough_adapter.go | 48 +++-- 26 files changed, 1003 insertions(+), 122 deletions(-) create mode 100644 backend/internal/service/openai_account_model_transient.go create mode 100644 backend/internal/service/openai_account_model_transient_test.go create mode 100644 backend/internal/service/openai_account_runtime_transient_test.go diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 01928028e6..61c32344b3 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -273,7 +273,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) @@ -321,7 +321,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") } @@ -332,7 +332,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 a1d185a1cb..e8a1429f03 100644 --- a/backend/internal/handler/openai_alpha_search.go +++ b/backend/internal/handler/openai_alpha_search.go @@ -163,7 +163,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) } @@ -172,7 +172,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") } @@ -180,7 +180,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_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index 62a2b70b59..2a7e5703b4 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -251,7 +251,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) @@ -296,7 +296,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 { @@ -312,9 +312,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 99f869f14c..5a159cf027 100644 --- a/backend/internal/handler/openai_embeddings.go +++ b/backend/internal/handler/openai_embeddings.go @@ -192,7 +192,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), @@ -216,7 +216,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") } @@ -227,7 +227,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 83c4b4ac23..83a618fc58 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 "" @@ -481,7 +485,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) @@ -541,7 +545,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 { @@ -566,9 +570,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） @@ -1009,7 +1013,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) @@ -1061,7 +1065,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), @@ -1072,9 +1076,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") @@ -1554,7 +1558,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() { @@ -1773,7 +1777,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) @@ -1853,7 +1857,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), diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index d855f364a7..21f5319092 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"}]}`) 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/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 a0a4fbff3e..d85cf3e4dd 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -1108,7 +1108,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 未设置的账号并标记异常 @@ -1366,7 +1366,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 @@ -1841,7 +1841,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 2ff7c25e5d..7252a4fcd2 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -1032,7 +1032,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() @@ -2079,6 +2079,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) @@ -2534,7 +2558,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() @@ -2871,7 +2895,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_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index bbf80c97c9..d82ee3d9fb 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -576,7 +576,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, @@ -637,7 +638,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_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index f318adc604..aff87faffe 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 } @@ -864,7 +864,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash) } else if !openAIStickyAccountMatchesGroup(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) @@ -927,7 +927,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) { @@ -1162,7 +1162,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 @@ -1207,7 +1207,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_upstream_errors.go b/backend/internal/service/openai_gateway_upstream_errors.go index d8823fe4ed..239a34b595 100644 --- a/backend/internal/service/openai_gateway_upstream_errors.go +++ b/backend/internal/service/openai_gateway_upstream_errors.go @@ -262,9 +262,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) @@ -385,6 +385,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_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 51aace6419..8f12fe5ab2 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -674,6 +674,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 +820,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 +954,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 +983,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..cc47f3cafd 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") 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..77234757b0 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -199,6 +199,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 +352,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( flushedBufferedEventCount := 0 firstEventType := "" lastEventType := "" + upstreamTerminalEvent := "" var flusher http.Flusher if reqStream { @@ -552,6 +554,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 +653,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 +725,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_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()) From a8a8a58f65ff76d004e0ee27aadfc35cfb92cd6d Mon Sep 17 00:00:00 2001 From: jjaw Date: Thu, 16 Jul 2026 00:11:17 +0800 Subject: [PATCH 11/26] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=89=B9=E9=87=8F?= =?UTF-8?q?=E8=B4=A6=E6=88=B7=E4=BA=8B=E4=BB=B6=E7=9A=84=E8=B0=83=E5=BA=A6?= =?UTF-8?q?=E9=87=8D=E5=BB=BA=E8=8C=83=E5=9B=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../scheduler_snapshot_bulk_event_test.go | 211 ++++++++++++++++++ .../service/scheduler_snapshot_service.go | 74 +++++- 2 files changed, 279 insertions(+), 6 deletions(-) create mode 100644 backend/internal/service/scheduler_snapshot_bulk_event_test.go 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..5b186aaf6d 100644 --- a/backend/internal/service/scheduler_snapshot_service.go +++ b/backend/internal/service/scheduler_snapshot_service.go @@ -529,11 +529,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 +546,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 { From 410ea8490cae48615894772915212593f30b0b0b Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Thu, 16 Jul 2026 00:16:21 +0800 Subject: [PATCH 12/26] fix(grok): allow passive Codex image tool declarations --- .../handler/gateway_handler_responses.go | 2 +- .../handler/openai_gateway_handler.go | 4 +- .../openai_grok_image_intent_gate_test.go | 101 ++++++++++++++++ .../service/image_generation_intent.go | 106 ++++++++++++++++- .../image_generation_intent_grok_test.go | 108 ++++++++++++++++++ .../service/openai_ws_forwarder_ingress.go | 2 +- 6 files changed, 316 insertions(+), 7 deletions(-) create mode 100644 backend/internal/handler/openai_grok_image_intent_gate_test.go create mode 100644 backend/internal/service/image_generation_intent_grok_test.go 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/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 83c4b4ac23..6a634bb595 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -272,7 +272,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 @@ -1461,7 +1461,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return } - if service.IsImageGenerationIntent("/v1/responses", reqModel, firstMessage) && !service.GroupAllowsImageGeneration(apiKey.Group) { + if service.IsImageGenerationIntentForPlatform("/v1/responses", reqModel, firstMessage, openAICompatibleRequestPlatform(apiKey)) && !service.GroupAllowsImageGeneration(apiKey.Group) { closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, service.ImageGenerationPermissionMessage()) return } 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/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/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 51aace6419..89846e1541 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) } From e8cc1fe26b07152ee9a5a3cb584d2499e2d21f3d Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Thu, 16 Jul 2026 00:15:38 +0800 Subject: [PATCH 13/26] feat(frontend): add Grok key setup templates --- frontend/src/components/keys/UseKeyModal.vue | 136 +++++++++++++++- .../keys/__tests__/UseKeyModal.spec.ts | 153 +++++++++++++++++- frontend/src/i18n/locales/en/dashboard.ts | 9 +- frontend/src/i18n/locales/zh/dashboard.ts | 11 +- 4 files changed, 297 insertions(+), 12 deletions(-) diff --git a/frontend/src/components/keys/UseKeyModal.vue b/frontend/src/components/keys/UseKeyModal.vue index 1dda0c51eb..4663b2b7d8 100644 --- a/frontend/src/components/keys/UseKeyModal.vue +++ b/frontend/src/components/keys/UseKeyModal.vue @@ -29,11 +29,12 @@

-
-